Retrain simple_motion_search_based_split on speed 1-2

The model is retrained using simple_motion_search results from
PARTITION_NONE and PARTITION_SPLIT without two_pass_partition_search.

The new model is only enabled on lowres and midres, where
there is significan speed up. The model's thresholds can be further
refined on lowres. On hdres, the performance is fairly neutral, but we
might get gains if we disable two pass partition search

Performance:
 TESTSET | AVG_PSNR | OVR_PSNR |   SSIM   | AVG_SPDUP | OVR_SPDUP |
  MIDRES |  -0.127% |  -0.141% | -0.213%  |  +8.202%  |  +9.012%  |
  LOWRES |  +0.075% |  +0.080% | +0.038%  |  +8.468%  | +10.568%  |

STATS_CHANGED

Change-Id: Ibc8047a9c19ea1c93d0d2e62b840c76976117926
diff --git a/av1/encoder/context_tree.h b/av1/encoder/context_tree.h
index 205ac8a..74d4b2f 100644
--- a/av1/encoder/context_tree.h
+++ b/av1/encoder/context_tree.h
@@ -99,7 +99,15 @@
   PC_TREE_STATS pc_tree_stats;
   CB_TREE_SEARCH cb_search_range;
   int index;
+
+  // Simple motion search_features
   MV mv_ref_fulls[REF_FRAMES];
+  unsigned int sms_none_feat[2];
+  unsigned int sms_split_feat[8];
+  unsigned int sms_rect_feat[8];
+  int sms_none_valid;
+  int sms_split_valid;
+  int sms_rect_valid;
 } PC_TREE;
 
 void av1_setup_pc_tree(struct AV1Common *cm, struct ThreadData *td);
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index a03d8c4..0dfd7e2 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -3000,9 +3000,6 @@
 // TODO(jinging,jimbankoski,rbultje): properly skip partition types that are
 // unlikely to be selected depending on previous rate-distortion optimization
 // results, for encoding speed-up.
-// TODO(chiyotsai@google.com): Move these ml related varables to a seprate file
-// to separate low level ml logic from partition logic
-#define NUM_SIMPLE_MOTION_FEATURES 28
 static void rd_pick_partition(AV1_COMP *const cpi, ThreadData *td,
                               TileDataEnc *tile_data, TOKENEXTRA **tp,
                               int mi_row, int mi_col, BLOCK_SIZE bsize,
@@ -3189,14 +3186,14 @@
   // Use simple_motion_search to prune partitions. This must be done prior to
   // PARTITION_SPLIT to propagate the initial mvs to a smaller blocksize.
   const int try_split_only =
-      cpi->sf.simple_motion_search_split_only && bsize >= BLOCK_8X8 &&
-      do_square_split && mi_row + mi_size_high[bsize] <= cm->mi_rows &&
+      cpi->sf.simple_motion_search_split_only && do_square_split &&
+      bsize >= BLOCK_8X8 && mi_row + mi_size_high[bsize] <= cm->mi_rows &&
       mi_col + mi_size_wide[bsize] <= cm->mi_cols && !frame_is_intra_only(cm) &&
       !av1_superres_scaled(cm);
 
   if (try_split_only) {
     av1_simple_motion_search_based_split(
-        cpi, x, mi_row, mi_col, bsize, &partition_none_allowed,
+        cpi, x, pc_tree, mi_row, mi_col, bsize, &partition_none_allowed,
         &partition_horz_allowed, &partition_vert_allowed, &do_rectangular_split,
         &do_square_split);
   }
@@ -3208,15 +3205,11 @@
        (prune_horz && prune_vert)) &&
       (partition_horz_allowed || partition_vert_allowed) && bsize >= BLOCK_8X8;
 
-  float simple_motion_features[NUM_SIMPLE_MOTION_FEATURES] = { 0.0f };
-  int simple_motion_features_are_valid = 0;
-
   if (try_prune_rect) {
     av1_simple_motion_search_prune_part(
         cpi, x, pc_tree, mi_row, mi_col, bsize, &partition_none_allowed,
         &partition_horz_allowed, &partition_vert_allowed, &do_square_split,
-        &do_rectangular_split, &prune_horz, &prune_vert, simple_motion_features,
-        &simple_motion_features_are_valid);
+        &do_rectangular_split, &prune_horz, &prune_vert);
   }
 
   // Max and min square partition levels are defined as the partition nodes that
@@ -3370,10 +3363,9 @@
             this_rdc.rdcost < INT64_MAX && this_rdc.rdcost >= 0 &&
             this_rdc.rate < INT_MAX && this_rdc.rate >= 0 &&
             (do_square_split || do_rectangular_split)) {
-          av1_simple_motion_search_early_term_none(
-              cpi, x, pc_tree, mi_row, mi_col, bsize, &this_rdc,
-              &terminate_partition_search, simple_motion_features,
-              &simple_motion_features_are_valid);
+          av1_simple_motion_search_early_term_none(cpi, x, pc_tree, mi_row,
+                                                   mi_col, bsize, &this_rdc,
+                                                   &terminate_partition_search);
         }
       }
     }
@@ -4648,7 +4640,8 @@
     PC_TREE *const pc_root = td->pc_root[mib_size_log2 - MIN_MIB_SIZE_LOG2];
     pc_root->index = 0;
 
-    if ((sf->simple_motion_search_prune_rect ||
+    if ((sf->simple_motion_search_split_only ||
+         sf->simple_motion_search_prune_rect ||
          sf->simple_motion_search_early_term_none ||
          sf->firstpass_simple_motion_search_early_term) &&
         !frame_is_intra_only(cm)) {
diff --git a/av1/encoder/partition_model_weights.h b/av1/encoder/partition_model_weights.h
index b754c88..0a30622 100644
--- a/av1/encoder/partition_model_weights.h
+++ b/av1/encoder/partition_model_weights.h
@@ -2442,6 +2442,7 @@
 #undef LABEL_SIZE
 
 // Below are the models used for simple_motion_search_based_split
+// The model for the faster model
 static const float av1_simple_motion_search_based_split_thresh_128 = 2.0f;
 static const float av1_simple_motion_search_based_split_thresh_64 = 2.0f;
 static const float av1_simple_motion_search_based_split_thresh_32 = 2.0f;
@@ -2700,49 +2701,639 @@
 #undef NUM_LAYER_0_UNITS_16
 #undef NUM_LOGITS_16
 
-#if !CONFIG_DISABLE_FULL_PIXEL_SPLIT_8X8
-// BLOCK_8X8
+// Slower higher quality split model
+static const float av1_simple_motion_search_split_thresh_128 =
+    5.839480f;  // p = 0.997098
+static const float av1_simple_motion_search_split_thresh_64 =
+    1.877167f;  // p = 0.867285
+static const float av1_simple_motion_search_split_thresh_32 =
+    3.073499f;  // p = 0.955783
+static const float av1_simple_motion_search_split_thresh_16 =
+    1.405601f;  // p = 0.803071
+static const float av1_simple_motion_search_split_thresh_8 =
+    2.555636f;  // p = 0.927951
+
+static const float av1_simple_motion_search_split_mean_128[17] = {
+  14.119120f, 14.087010f, 12.016185f, 11.966075f, 12.042454f, 11.994805f,
+  12.152105f, 12.100394f, 12.178377f, 12.128937f, 4.779944f,  0.714786f,
+  3.535450f,  3.566207f,  0.835913f,  3.315452f,  3.302908f,
+};
+
+static const float av1_simple_motion_search_split_std_128[17] = {
+  1.832420f, 1.835338f, 2.019207f, 2.020793f, 2.008731f, 2.008403f,
+  1.900999f, 1.907081f, 1.908915f, 1.913122f, 2.109345f, 0.451517f,
+  1.407097f, 1.372501f, 0.370355f, 1.321495f, 1.319665f,
+};
+
+static const float av1_simple_motion_search_split_mean_64[17] = {
+  12.363721f, 12.314348f, 10.404341f, 10.333541f, 10.405775f, 10.336996f,
+  10.402246f, 10.330084f, 10.405584f, 10.334330f, 4.554232f,  0.896393f,
+  2.819613f,  2.855845f,  0.926296f,  2.808782f,  2.798229f,
+};
+
+static const float av1_simple_motion_search_split_std_64[17] = {
+  1.878920f, 1.882255f, 1.950167f, 1.953289f, 1.913869f, 1.914781f,
+  1.920096f, 1.924454f, 1.880200f, 1.882499f, 2.050922f, 0.304750f,
+  1.144391f, 1.125088f, 0.261289f, 1.145059f, 1.131215f,
+};
+
+static const float av1_simple_motion_search_split_mean_32[17] = {
+  10.750278f, 10.679627f, 8.745625f, 8.644149f, 8.757436f, 8.656657f,
+  8.759780f,  8.656299f,  8.772563f, 8.669839f, 4.208026f, 0.958573f,
+  2.308769f,  2.347375f,  0.961685f, 2.323464f, 2.296322f,
+};
+
+static const float av1_simple_motion_search_split_std_32[17] = {
+  1.879269f, 1.883531f, 1.935828f, 1.935677f, 1.915823f, 1.914773f,
+  1.909733f, 1.910315f, 1.890451f, 1.890032f, 1.913318f, 0.199276f,
+  0.988825f, 0.972115f, 0.191956f, 0.977131f, 0.951418f,
+};
+
+static const float av1_simple_motion_search_split_mean_16[17] = {
+  9.076768f, 8.974986f, 7.078364f, 6.926072f, 7.088739f, 6.936111f,
+  7.096697f, 6.942841f, 7.114978f, 6.961046f, 3.865480f, 0.982632f,
+  1.886023f, 1.912892f, 0.981492f, 1.926059f, 1.891233f,
+};
+
+static const float av1_simple_motion_search_split_std_16[17] = {
+  1.922965f, 1.925609f, 1.851980f, 1.847558f, 1.848410f, 1.843990f,
+  1.843931f, 1.839582f, 1.840304f, 1.836144f, 1.760042f, 0.130639f,
+  0.841086f, 0.833523f, 0.134780f, 0.840790f, 0.831309f,
+};
+
+static const float av1_simple_motion_search_split_mean_8[17] = {
+  7.120238f, 6.957731f, 5.176309f, 4.889594f, 5.178396f, 4.886607f,
+  5.195322f, 4.905566f, 5.198845f, 4.904745f, 3.648933f, 0.993198f,
+  1.496831f, 1.520804f, 0.991864f, 1.489763f, 1.460761f,
+};
+
+static const float av1_simple_motion_search_split_std_8[17] = {
+  1.698498f, 1.696000f, 1.629605f, 1.614641f, 1.632476f, 1.618831f,
+  1.618352f, 1.603742f, 1.623089f, 1.609674f, 1.668587f, 0.082193f,
+  0.759407f, 0.759684f, 0.089830f, 0.742797f, 0.730632f,
+};
+
+#define NUM_HIDDEN_LAYERS_128 1
+#define NUM_FEATURES_128 17
+#define NUM_LAYER_0_UNITS_128 20
+#define NUM_LOGITS_128 1
+
+static const float av1_simple_motion_search_split_hiddenlayer_0_kernel_128[] = {
+  0.24095f,    -0.397761f,  -0.388619f,  -0.0629548f, -0.44577f,   0.688212f,
+  -0.20889f,   -1.08227f,   -0.0313894f, -0.615505f,  -0.401839f,  0.40233f,
+  -0.171305f,  0.439803f,   1.58527f,    -0.968535f,  -1.29255f,   1.14846f,
+  0.885777f,   0.116412f,   -0.225704f,  0.316506f,   0.793951f,   -0.63591f,
+  0.097789f,   -0.327027f,  -0.778396f,  -0.231667f,  -0.9622f,    1.0044f,
+  0.32594f,    0.179768f,   -0.115529f,  -0.499395f,  -1.14727f,   -1.26111f,
+  0.269818f,   -0.0882028f, -0.349107f,  0.100901f,   0.0249506f,  0.528929f,
+  0.113961f,   0.929794f,   0.242494f,   -0.122828f,  -0.0477379f, 0.170659f,
+  0.0500187f,  0.28859f,    0.78783f,    0.482412f,   0.795298f,   0.179517f,
+  0.453911f,   -0.298029f,  -0.903332f,  0.510615f,   0.691994f,   0.433383f,
+  -0.140802f,  -1.11635f,   -0.547326f,  1.11318f,    0.71905f,    0.978538f,
+  0.097444f,   -0.0386012f, 0.713599f,   0.465164f,   0.391278f,   -0.472864f,
+  0.230224f,   -0.279508f,  0.558192f,   -0.468625f,  0.55995f,    -0.57507f,
+  -1.39947f,   -0.755819f,  -1.04512f,   -0.411552f,  -0.830444f,  -0.106571f,
+  -0.0972184f, 0.251842f,   0.269955f,   0.230492f,   -0.290581f,  -0.484799f,
+  0.0151041f,  0.171047f,   0.829999f,   -0.384581f,  0.220301f,   -0.121687f,
+  1.88848f,    -0.482809f,  -0.48185f,   1.34482f,    -0.716438f,  -0.284482f,
+  -1.78592f,   -1.29333f,   0.886867f,   0.80106f,    0.456415f,   0.649095f,
+  0.231093f,   0.361562f,   0.290018f,   0.128009f,   -0.196343f,  0.0607802f,
+  0.576761f,   -0.0413836f, 0.0300984f,  -0.318998f,  0.204434f,   -0.712524f,
+  0.833394f,   -0.81168f,   0.765488f,   -0.720973f,  1.12866f,    -0.838694f,
+  1.295f,      -0.159127f,  1.05404f,    0.736519f,   0.248662f,   0.229233f,
+  0.0434302f,  0.0551856f,  0.197862f,   0.354823f,   -0.32429f,   -0.227353f,
+  -0.132198f,  -0.438118f,  -0.210401f,  -0.81046f,   0.653555f,   0.826737f,
+  0.154235f,   0.228945f,   0.123089f,   0.614964f,   -0.0940471f, -0.00676807f,
+  0.24996f,    0.949233f,   0.746526f,   -0.044474f,  0.386414f,   0.503221f,
+  0.155133f,   -0.698848f,  -0.735356f,  -0.255091f,  0.413235f,   -0.335295f,
+  -0.145757f,  0.326299f,   -0.602629f,  -0.844474f,  -0.346722f,  -0.42598f,
+  -0.491016f,  -0.447732f,  -0.965366f,  -0.0242841f, 0.836606f,   -0.104877f,
+  1.23236f,    0.683986f,   0.787005f,   -0.0253437f, 1.2145f,     1.29554f,
+  -1.24302f,   -0.229495f,  0.439415f,   0.885087f,   -0.408704f,  -0.119299f,
+  -0.0960972f, 0.60148f,    0.683271f,   -0.057129f,  -0.180295f,  -0.264815f,
+  -0.363184f,  0.638271f,   0.631083f,   -0.252899f,  -0.164364f,  -1.31274f,
+  0.354408f,   0.0429172f,  0.371154f,   -1.0978f,    0.0433642f,  -0.467394f,
+  -0.706572f,  1.57198f,    -0.0701271f, 1.93149f,    -0.446267f,  1.4519f,
+  -1.29567f,   0.309978f,   -0.878062f,  0.891494f,   0.364005f,   -0.209611f,
+  -0.125927f,  0.184097f,   0.0629695f,  -0.43375f,   -0.0980562f, 1.08547f,
+  0.578312f,   0.16566f,    -0.198852f,  -0.241854f,  -0.523934f,  -0.206037f,
+  -0.867721f,  1.00041f,    1.09848f,    -2.12562f,   -0.19992f,   -0.186128f,
+  -0.03507f,   0.0484884f,  0.160856f,   0.10802f,    -0.805141f,  -1.06902f,
+  0.290363f,   0.0222096f,  -0.849266f,  0.112932f,   0.148682f,   -0.0457585f,
+  1.139f,      1.79141f,    0.194122f,   -0.342508f,  -0.403572f,  0.133678f,
+  0.217553f,   -0.263759f,  0.18441f,    0.254529f,   0.0471115f,  0.733178f,
+  -0.416205f,  0.441447f,   -0.443335f,  0.725005f,   -0.78946f,   0.71301f,
+  -0.644969f,  1.5445f,     0.365277f,   -0.455775f,  -0.365066f,  0.4742f,
+  -0.381714f,  -0.545794f,  -0.0464861f, -0.222768f,  -0.0106466f, -0.069743f,
+  0.0335566f,  0.378348f,   -0.249663f,  0.922286f,   0.125711f,   -0.894619f,
+  0.444682f,   0.447893f,   -1.98936f,   -1.41978f,   0.0406667f,  -0.199928f,
+  -0.199786f,  0.463481f,   0.334931f,   -0.396222f,  -0.0732259f, 0.796684f,
+  -0.140817f,  -0.26878f,   0.194642f,   0.895784f,   -0.369976f,  -2.26981f,
+  -0.0791776f, -0.0492268f, 0.6715f,     0.281805f,   0.0156664f,  -0.779785f,
+  0.17743f,    0.188786f,   -0.588077f,  -0.359153f,  0.258319f,   0.881688f,
+  0.846894f,   1.00292f,    0.838134f,   0.680632f,   0.273098f,   -0.329261f,
+  0.217757f,   -0.506726f,  -0.336523f,  -0.695875f,  -0.252006f,  0.751216f,
+  0.334409f,   -0.0151467f, 0.0885474f,  0.0973114f,  -0.248754f,  -0.263716f,
+  0.369906f,   -0.213749f,  -0.0355395f, -0.137799f,  2.43233f,    -0.944233f,
+  -0.745167f,  0.318558f,   0.316608f,   0.568678f
+};
+
+static const float av1_simple_motion_search_split_hiddenlayer_0_bias_128[] = {
+  0.821344f,  1.11542f,   -1.24172f,  1.03642f,  1.13511f,
+  1.16414f,   -0.278655f, -1.35558f,  -1.26788f, -1.63189f,
+  -0.323271f, 1.21319f,   -0.888415f, 0.987145f, -1.16767f,
+  0.255833f,  -0.1392f,   1.43265f,   -1.54952f, 1.65159f
+};
+
+static const float av1_simple_motion_search_split_logits_kernel_128[] = {
+  0.3565753f, 0.5490161f, -1.015597f, 0.565366f,   0.751604f,
+  0.922747f,  -1.931846f, 1.759353f,  -0.7362949f, 0.5707034f,
+  -1.092127f, 0.936767f,  2.034499f,  2.08148f,    0.9509507f,
+  -1.342504f, -0.834566f, 0.618184f,  0.844113f,   1.182693f
+};
+
+static const float av1_simple_motion_search_split_logits_bias_128[] = {
+  1.819351f
+};
+
+static const NN_CONFIG av1_simple_motion_search_split_nn_config_128 = {
+  NUM_FEATURES_128,
+  NUM_LOGITS_128,
+  NUM_HIDDEN_LAYERS_128,
+  {
+      NUM_LAYER_0_UNITS_128,
+  },
+  {
+      av1_simple_motion_search_split_hiddenlayer_0_kernel_128,
+      av1_simple_motion_search_split_logits_kernel_128,
+  },
+  {
+      av1_simple_motion_search_split_hiddenlayer_0_bias_128,
+      av1_simple_motion_search_split_logits_bias_128,
+  },
+};
+
+#undef NUM_HIDDEN_LAYERS_128
+#undef NUM_FEATURES_128
+#undef NUM_LAYER_0_UNITS_128
+#undef NUM_LOGITS_128
+
+#define NUM_HIDDEN_LAYERS_64 1
+#define NUM_FEATURES_64 17
+#define NUM_LAYER_0_UNITS_64 24
+#define NUM_LOGITS_64 1
+
+static const float av1_simple_motion_search_split_hiddenlayer_0_kernel_64[] = {
+  -1.40663f,    -0.851503f,   -0.0613111f,  0.741591f,    0.302754f,
+  0.184001f,    0.0474853f,   0.371096f,    0.0541624f,   0.381508f,
+  0.355427f,    0.0428822f,   0.154916f,    -0.00490099f, 0.025484f,
+  0.0208921f,   0.140596f,    -0.292525f,   -0.459067f,   -0.081393f,
+  0.109824f,    -0.290183f,   0.720236f,    0.385835f,    -0.150643f,
+  -0.078518f,   0.0979819f,   -0.102135f,   0.137152f,    -0.0786457f,
+  0.0171441f,   0.991338f,    -0.546583f,   -1.0714f,     -0.0842851f,
+  0.244072f,    0.427379f,    0.146775f,    -0.921613f,   -0.912093f,
+  0.393566f,    -0.232375f,   0.19963f,     0.312355f,    0.55659f,
+  -0.104714f,   -0.137563f,   0.0985237f,   0.0788307f,   -0.225514f,
+  0.0228832f,   -0.288733f,   -0.00737685f, -0.711657f,   -0.256796f,
+  0.0869605f,   0.583977f,    0.384306f,    1.46692f,     -0.741126f,
+  -0.21105f,    -0.276604f,   -0.0151463f,  -0.0227997f,  -0.0403232f,
+  0.044122f,    0.0185784f,   -0.0451951f,  0.00489513f,  -0.387131f,
+  0.0966724f,   -0.599174f,   -0.00243351f, -0.21439f,    0.302043f,
+  0.130334f,    -0.191251f,   0.863261f,    -1.50112f,    0.00901057f,
+  0.000324294f, -0.0572545f,  0.0117685f,   -0.0734682f,  -0.0570435f,
+  -0.126253f,   1.2313f,      -0.328267f,   0.211788f,    -0.175438f,
+  -0.0419298f,  0.166447f,    -0.178739f,   -0.326221f,   -0.0439188f,
+  1.01182f,     -0.390678f,   -0.426343f,   0.0944665f,   -0.225042f,
+  -0.183344f,   0.0500763f,   -0.377393f,   -0.673401f,   -0.436907f,
+  -0.00366876f, -0.363412f,   0.195194f,    0.250248f,    -0.397193f,
+  -0.0917222f,  -0.0221579f,  1.7693f,      -0.0694484f,  -0.0410764f,
+  -0.134571f,   -0.159992f,   -0.170359f,   -0.249333f,   -0.128056f,
+  -0.617054f,   -0.808701f,   -0.540642f,   0.396391f,    0.147787f,
+  0.346916f,    0.709852f,    0.116064f,    0.0509731f,   0.073713f,
+  -0.365082f,   -1.09287f,    -0.618214f,   0.20545f,     0.126161f,
+  -0.140012f,   0.62592f,     0.316326f,    -0.392765f,   -0.15934f,
+  0.337617f,    -0.41669f,    -0.295225f,   0.0602025f,   -0.0150657f,
+  -0.319629f,   0.783729f,    -0.0661199f,  -0.362657f,   0.390042f,
+  -0.043614f,   -0.0414596f,  0.121155f,    -0.309775f,   -0.284761f,
+  -0.243932f,   0.279855f,    -0.266823f,   0.734824f,    -0.164028f,
+  0.261776f,    -0.105585f,   0.10733f,     -0.180469f,   1.18875f,
+  -1.12836f,    -0.173008f,   0.150221f,    0.111598f,    0.148306f,
+  -1.2833f,     -1.06346f,    0.233546f,    0.16432f,     0.00142378f,
+  0.340574f,    -0.0140885f,  0.634761f,    -0.122096f,   0.821487f,
+  0.421424f,    -0.0256687f,  -0.035503f,   -0.0453547f,  -0.0215179f,
+  -0.0671277f,  -0.0486862f,  -0.962761f,   -0.208383f,   0.109573f,
+  -0.210668f,   -0.176485f,   0.421279f,    0.41605f,     0.342084f,
+  0.619364f,    0.103718f,    -0.00341643f, 0.00266677f,  0.249089f,
+  -0.22848f,    -0.0368968f,  1.12092f,     -0.64912f,    -0.456579f,
+  0.477823f,    0.418345f,    1.41515f,     0.0936279f,   0.886155f,
+  -0.785656f,   -0.217109f,   -0.561829f,   -0.286435f,   -0.884068f,
+  -0.148839f,   -0.282848f,   0.0683745f,   0.0962815f,   -0.111975f,
+  0.0509158f,   -0.211274f,   0.744909f,    -0.8982f,     0.315232f,
+  -0.78624f,    0.598387f,    -0.530952f,   0.677357f,    0.0371339f,
+  0.99209f,     -0.681899f,   -0.291416f,   -0.224822f,   -0.26049f,
+  -0.0436525f,  -0.380004f,   -0.27187f,    0.534779f,    0.717939f,
+  0.418197f,    -0.152539f,   -0.0684039f,  -0.186308f,   -0.0653121f,
+  0.194145f,    -0.196367f,   0.256997f,    -0.726269f,   -0.307672f,
+  -0.153362f,   0.450827f,    0.708842f,    -0.0667079f,  0.555564f,
+  0.0486892f,   0.0715072f,   -0.7211f,     -0.849797f,   0.0650271f,
+  1.2747f,      -0.646738f,   -0.53042f,    0.182197f,    0.928203f,
+  0.180621f,    -0.00640791f, -0.171416f,   0.092688f,    -0.391275f,
+  -0.0650657f,  0.0843773f,   0.170824f,    0.378085f,    0.0596657f,
+  0.844398f,    -1.3083f,     -1.27828f,    -0.199179f,   0.557855f,
+  0.241479f,    0.385804f,    0.169533f,    -0.0028072f,  0.0538041f,
+  0.00136234f,  0.0130481f,   0.0349449f,   -0.0366494f,  -0.000474055f,
+  0.437956f,    0.286724f,    -0.298187f,   0.461967f,    0.43065f,
+  -0.0877194f,  -0.19133f,    0.379121f,    -0.687751f,   -1.64077f,
+  -0.375191f,   -0.336836f,   -0.323904f,   -0.101859f,   0.0126672f,
+  -0.346332f,   0.112303f,    -0.863336f,   0.155538f,    0.366509f,
+  -0.0976829f,  0.635278f,    -0.681967f,   -0.527729f,   0.591839f,
+  0.366678f,    0.189981f,    0.0208007f,   -0.565809f,   0.70183f,
+  -0.282844f,   -0.327485f,   0.347243f,    -1.13014f,    -0.373378f,
+  -0.514978f,   0.662994f,    -0.144931f,   0.1402f,      -0.820049f,
+  0.711498f,    0.681156f,    1.06515f,     -0.423409f,   -0.0392664f,
+  0.0675396f,   -0.0508602f,  0.0431443f,   0.0212639f,   -0.0279887f,
+  -0.62611f,    -0.202064f,   0.701934f,    1.28452f,     -0.00858481f,
+  -0.517249f,   0.0615832f,   -0.260215f,   0.0949119f,   -0.28423f,
+  -0.39573f,    -0.0574246f,  -0.318658f,   0.0601775f,   -0.0629386f,
+  -0.134208f,   0.111686f,    -0.23355f,    0.078667f,    0.741023f,
+  0.828523f,    -0.345067f,   -0.315135f,   -0.0957154f,  0.522825f,
+  -0.190057f,   -0.473789f,   -0.390489f,   0.200677f,    -0.0271802f,
+  0.110336f,    0.493302f,    0.663126f,    0.570148f,    -0.380042f,
+  -0.437349f,   -0.660884f,   0.301908f,    0.0644179f,   0.172494f,
+  0.461917f,    0.330938f,    -0.140041f,   -0.0430205f,  -1.51003f,
+  -0.410984f,   -0.182161f,   0.0235313f,   -0.364849f,   0.154183f,
+  -0.592465f,   0.272701f,    0.192389f,    -0.0497777f,  -0.924467f,
+  -0.179513f,   -0.592217f,   0.436363f,    -0.0716164f,  0.189094f,
+  -0.574697f,   -0.304303f,   0.326441f,    -0.0865553f,  0.735948f,
+  0.266912f,    0.435824f,    -0.123322f
+};
+
+static const float av1_simple_motion_search_split_hiddenlayer_0_bias_64[] = {
+  -1.19333f,  1.01834f,   -1.10844f,  0.0454873f, -1.45506f,   0.580864f,
+  -0.040979f, -0.505681f, -1.15072f,  0.692697f,  -0.520812f,  -0.479384f,
+  0.529652f,  0.507252f,  -1.08619f,  0.0586375f, 0.0929614f,  -0.46753f,
+  -0.701857f, -0.362933f, -0.291983f, -0.133933f, -0.0131351f, -0.267582f
+};
+
+static const float av1_simple_motion_search_split_logits_kernel_64[] = {
+  -3.32501f,  0.43082f,   -1.060692f, 1.328908f,  0.8892894f,  0.6488833f,
+  -1.096516f, -0.664786f, -1.301339f, 0.508805f,  -2.128406f,  -0.757304f,
+  0.383839f,  0.694763f,  -0.591725f, 0.770385f,  1.021594f,   0.589181f,
+  -0.76238f,  1.488826f,  0.709135f,  -0.575738f, 0.26421759f, -0.2484219f
+};
+
+static const float av1_simple_motion_search_split_logits_bias_64[] = {
+  0.699037f
+};
+
+static const NN_CONFIG av1_simple_motion_search_split_nn_config_64 = {
+  NUM_FEATURES_64,
+  NUM_LOGITS_64,
+  NUM_HIDDEN_LAYERS_64,
+  {
+      NUM_LAYER_0_UNITS_64,
+  },
+  {
+      av1_simple_motion_search_split_hiddenlayer_0_kernel_64,
+      av1_simple_motion_search_split_logits_kernel_64,
+  },
+  {
+      av1_simple_motion_search_split_hiddenlayer_0_bias_64,
+      av1_simple_motion_search_split_logits_bias_64,
+  },
+};
+
+#undef NUM_HIDDEN_LAYERS_64
+#undef NUM_FEATURES_64
+#undef NUM_LAYER_0_UNITS_64
+#undef NUM_LOGITS_64
+
+#define NUM_HIDDEN_LAYERS_32 1
+#define NUM_FEATURES_32 17
+#define NUM_LAYER_0_UNITS_32 20
+#define NUM_LOGITS_32 1
+
+static const float av1_simple_motion_search_split_hiddenlayer_0_kernel_32[] = {
+  -0.980626f,   -0.946611f,    0.103761f,    0.408899f,    0.498149f,
+  0.0490161f,   0.253279f,     0.332029f,    0.00367441f,  0.364401f,
+  -0.236433f,   0.0592119f,    -0.0978848f,  0.159733f,    -0.018052f,
+  -1.10726f,    1.16167f,      -0.244982f,   -0.147819f,   -0.147095f,
+  0.111404f,    -0.349502f,    0.441178f,    0.0984191f,   -0.135537f,
+  -0.0423312f,  0.0123079f,    0.358012f,    -0.266796f,   0.0125811f,
+  0.196563f,    0.337093f,     -1.07266f,    -1.25134f,    0.57337f,
+  -0.521717f,   0.259824f,     0.537383f,    -0.463688f,   -0.336128f,
+  0.373385f,    0.483443f,     -0.229293f,   -0.33373f,    -0.656021f,
+  0.768647f,    0.179279f,     0.315415f,    0.187749f,    1.07839f,
+  0.0626629f,   -0.230299f,    0.662606f,    -0.414154f,   0.459334f,
+  -0.6312f,     0.427704f,     -0.249849f,   0.701056f,    -0.707969f,
+  0.057401f,    0.620434f,     0.665748f,    -0.501356f,   -0.230685f,
+  0.0722371f,   -0.0988625f,   -0.114035f,   -0.653799f,   0.571353f,
+  0.268276f,    1.13251f,      -1.0695f,     -0.225607f,   -0.984355f,
+  -0.42213f,    0.300422f,     1.21492f,     -0.139931f,   -0.000726004f,
+  0.045964f,    -0.0817352f,   -0.0278813f,  -0.0102341f,  -0.0144087f,
+  -0.475882f,   1.20682f,      -0.359919f,   0.277189f,    -0.166401f,
+  0.599211f,    -0.129872f,    0.574211f,    -0.247573f,   0.824405f,
+  -1.53329f,    -0.202151f,    -0.328698f,   -0.516322f,   -0.281416f,
+  -0.383651f,   -0.252862f,    -0.43185f,    0.456802f,    -0.430055f,
+  -0.55245f,    -0.6884f,      -0.541456f,   -0.281376f,   1.10425f,
+  -0.140706f,   1.59816f,      -0.0343895f,  -0.00920039f, -0.0307667f,
+  0.0560132f,   -0.0340302f,   -0.10848f,    0.0593314f,   -0.951795f,
+  0.876831f,    -1.00548f,     -0.566244f,   0.430061f,    1.10109f,
+  -0.634212f,   -0.0755369f,   -0.108953f,   1.03191f,     0.109036f,
+  -0.0415309f,  0.0681162f,    -0.0611775f,  -0.0231938f,  0.0973158f,
+  -0.0558169f,  -0.823484f,    -0.918509f,   0.16756f,     0.27087f,
+  0.286074f,    0.174069f,     0.1304f,      0.386074f,    0.433953f,
+  0.0291467f,   -1.74087f,     0.0296094f,   -0.00793714f, -0.13041f,
+  0.00990992f,  -0.0137848f,   -0.0742606f,  -0.251029f,   -0.645316f,
+  0.640029f,    0.550607f,     0.470097f,    0.549451f,    -0.285723f,
+  -0.164759f,   -0.128166f,    -0.391496f,   -0.80287f,    0.0769472f,
+  1.34391f,     0.0215005f,    0.0669497f,   0.131919f,    0.291674f,
+  0.0952889f,   -0.677953f,    -0.364054f,   0.144823f,    0.246198f,
+  -0.12393f,    0.363661f,     0.215091f,    -0.239658f,   0.18491f,
+  0.118703f,    0.0064156f,    1.38619f,     -1.3845f,     0.0567323f,
+  1.20812f,     -0.720374f,    -1.92158f,    -1.48657f,    0.335601f,
+  0.409379f,    0.373618f,     0.231274f,    0.292194f,    0.368619f,
+  0.2398f,      0.473579f,     0.83402f,     -0.0133751f,  -0.00344358f,
+  2.20688e-05f, 0.00836757f,   0.00405377f,  0.0110539f,   -0.260154f,
+  0.192112f,    -0.666986f,    0.302875f,    -0.113302f,   0.17882f,
+  -0.221493f,   0.146161f,     -0.448697f,   0.584187f,    0.122109f,
+  0.989981f,    -1.14706f,     -0.734042f,   0.0638213f,   0.213357f,
+  0.068543f,    -0.808558f,    0.404741f,    0.808313f,    1.57523f,
+  -0.113448f,   0.254102f,     -0.350065f,   -0.615f,      0.0753549f,
+  -0.540936f,   -0.0250732f,   -0.225681f,   -0.161384f,   0.0128342f,
+  -0.0933368f,  -0.286904f,    0.130133f,    -0.874747f,   0.392585f,
+  -0.493135f,   0.169708f,     0.0909804f,   1.89921f,     -0.469954f,
+  0.65165f,     -0.953401f,    -0.21595f,    -0.37479f,    0.0451146f,
+  0.0234621f,   -0.0596903f,   -0.0682308f,  -0.0830426f,  0.130011f,
+  -0.409141f,   0.0627038f,    -0.581148f,   -0.513922f,   0.631676f,
+  0.0637034f,   0.0539081f,    0.0638872f,   0.515863f,    -0.0123463f,
+  0.177238f,    0.279506f,     -0.930345f,   1.23726f,     0.202851f,
+  0.708792f,    -0.445086f,    -0.0267075f,  -0.913822f,   -0.0714978f,
+  -0.281107f,   -0.0770565f,   -0.23086f,    -0.165893f,   -0.319683f,
+  0.216235f,    -0.490999f,    2.04841f,     -0.0524071f,  -0.239043f,
+  -0.0526375f,  0.023002f,     -0.132685f,   -0.155354f,   -0.186503f,
+  -0.904296f,   0.166478f,     0.063268f,    -0.302842f,   -0.27179f,
+  -0.428299f,   0.50193f,      0.480717f,    -0.864275f,   0.317096f,
+  0.40698f,     0.0286107f,    0.189432f,    -0.0374374f,  0.0671728f,
+  0.203681f,    -0.457959f,    -0.155776f,   0.340948f,    0.542841f,
+  0.342675f,    -0.000952399f, 0.470957f,    0.744418f,    -1.11763f,
+  -0.658812f,   -0.044832f,    0.0688237f,   -0.357766f,   0.428662f,
+  -0.087152f,   -0.291903f,    0.373244f,    -0.587853f,   0.415895f,
+  -0.535694f,   0.621785f,     -0.143648f,   0.0451373f,   0.00068827f,
+  1.84432f,     -1.26239f,     -0.432087f,   -0.152307f,   0.0293551f,
+  0.184744f,    -0.0173156f,   -0.00572154f, -0.0305062f,  -0.0900071f
+};
+
+static const float av1_simple_motion_search_split_hiddenlayer_0_bias_32[] = {
+  0.160011f,  0.903856f,   -0.13738f,  0.358221f, -0.0906044f,
+  -0.606558f, -0.0215651f, -0.03377f,  -1.67017f, -0.144554f,
+  -0.201482f, -0.87719f,   0.639815f,  -0.51976f, -0.309922f,
+  -1.33421f,  0.721328f,   -0.889354f, -1.7158f,  -0.285963f
+};
+
+static const float av1_simple_motion_search_split_logits_kernel_32[] = {
+  -0.2745374f,  0.333548f,  -0.2437388f, 0.288009f,   0.55635f,
+  0.4560176f,   0.2970518f, 0.391192f,   1.311854f,   -0.231219f,
+  -0.2968651f,  -1.819984f, 0.2775824f,  0.28929857f, 0.419126f,
+  -0.32868411f, -0.916399f, -0.1921077f, -0.617489f,  0.637953f
+};
+
+static const float av1_simple_motion_search_split_logits_bias_32[] = {
+  0.208473f
+};
+
+static const NN_CONFIG av1_simple_motion_search_split_nn_config_32 = {
+  NUM_FEATURES_32,
+  NUM_LOGITS_32,
+  NUM_HIDDEN_LAYERS_32,
+  {
+      NUM_LAYER_0_UNITS_32,
+  },
+  {
+      av1_simple_motion_search_split_hiddenlayer_0_kernel_32,
+      av1_simple_motion_search_split_logits_kernel_32,
+  },
+  {
+      av1_simple_motion_search_split_hiddenlayer_0_bias_32,
+      av1_simple_motion_search_split_logits_bias_32,
+  },
+};
+
+#undef NUM_HIDDEN_LAYERS_32
+#undef NUM_FEATURES_32
+#undef NUM_LAYER_0_UNITS_32
+#undef NUM_LOGITS_32
+
+#define NUM_HIDDEN_LAYERS_16 1
+#define NUM_FEATURES_16 17
+#define NUM_LAYER_0_UNITS_16 20
+#define NUM_LOGITS_16 1
+
+static const float av1_simple_motion_search_split_hiddenlayer_0_kernel_16[] = {
+  0.0136957f,   0.182135f,    -0.583394f,    0.0556956f,   0.211152f,
+  0.168234f,    -0.694203f,   -0.678216f,    0.289943f,    1.00014f,
+  -0.0427784f,  -0.0427538f,  -0.0276009f,   -0.00133608f, 0.0901944f,
+  0.0674892f,   0.104068f,    -0.308582f,    -0.43596f,    0.855997f,
+  -0.223414f,   0.0390026f,   0.366492f,     0.216065f,    -0.386863f,
+  -0.148823f,   -0.297022f,   0.0529546f,    -0.202885f,   1.26471f,
+  -0.861163f,   -0.0949431f,  0.573627f,     -0.00277083f, -0.616063f,
+  -0.626927f,   0.371583f,    -0.411743f,    0.173387f,    -0.209734f,
+  0.293697f,    -0.260714f,   0.442728f,     -0.594486f,   1.38987f,
+  0.208025f,    -0.0433776f,  0.01173f,      0.921766f,    -0.168379f,
+  0.000697326f, 0.209967f,    -0.304577f,    0.149551f,    -0.196658f,
+  0.389251f,    -0.449106f,   -0.456329f,    0.669073f,    -0.163806f,
+  0.083348f,    -0.0783998f,  0.0678355f,    0.0510435f,   0.103964f,
+  0.104537f,    -0.778093f,   -1.0641f,      -0.626102f,   -2.02131f,
+  0.159591f,    0.254161f,    -0.000362642f, 0.289859f,    0.192713f,
+  0.139801f,    -0.0251327f,  0.164002f,     1.22892f,     -0.0852193f,
+  0.0769487f,   0.0296408f,   -0.0418688f,   0.0936023f,   0.0448523f,
+  0.674015f,    -0.0732944f,  0.313575f,     -0.593432f,   0.642067f,
+  -1.06063f,    0.468223f,    -0.769085f,    -0.173798f,   -0.175663f,
+  0.692808f,    0.00753295f,  -0.123327f,    -0.0234937f,  -0.0923153f,
+  0.0216917f,   -0.0690157f,  -0.397488f,    0.426628f,    0.264475f,
+  0.342074f,    -0.139817f,   0.215915f,     0.422544f,    -0.321102f,
+  0.0355587f,   0.460193f,    0.0315326f,    0.080556f,    -0.0256533f,
+  -0.0857874f,  -0.488283f,   -0.299653f,    -0.245987f,   0.104383f,
+  0.203731f,    0.328734f,    0.668104f,     -0.586909f,   -0.501335f,
+  -0.661292f,   -0.359811f,   0.00951363f,   0.816315f,    -0.0124104f,
+  0.0545827f,   0.089863f,    0.0125486f,    0.043609f,    -0.0259544f,
+  0.0123911f,   0.12557f,     -0.539875f,    -0.0556721f,  0.16532f,
+  0.265834f,    -0.384171f,   0.646496f,     0.366147f,    -0.111272f,
+  0.262096f,    -0.0845724f,  0.382724f,     0.165783f,    0.1025f,
+  0.392988f,    0.290525f,    0.038659f,     0.540269f,    -0.485586f,
+  -0.273065f,   -0.154052f,   -0.0896895f,   -0.35394f,    0.193214f,
+  -0.423728f,   0.654576f,    -0.373321f,    0.814914f,    0.026278f,
+  -0.0328304f,  -0.220913f,   -0.0442121f,   0.487545f,    -0.509537f,
+  -0.777581f,   -1.23886f,    0.223482f,     0.206009f,    0.20391f,
+  0.194628f,    0.226762f,    0.171609f,     -0.219037f,   0.557892f,
+  -0.312011f,   1.27709f,     0.064013f,     0.105384f,    0.0493933f,
+  0.074059f,    -0.0100078f,  -0.0176888f,   -0.440005f,   0.302922f,
+  -0.197456f,   0.296128f,    -0.326647f,    0.305323f,    -0.30696f,
+  0.201951f,    -0.15874f,    -0.793042f,    0.0197254f,   0.0569867f,
+  -0.0295468f,  -0.0215012f,  0.025855f,     -0.0196102f,  0.215558f,
+  -0.253069f,   0.298469f,    0.261269f,     0.435305f,    0.0120354f,
+  -0.384789f,   -0.2772f,     0.0366613f,    -0.494994f,   0.149072f,
+  1.32981f,     -0.427717f,   0.43938f,      -0.16375f,    -0.444342f,
+  0.548214f,    0.127955f,    -1.24387f,     0.0863676f,   0.175071f,
+  0.172673f,    -0.0906204f,  0.444454f,     -0.546669f,   0.215857f,
+  -0.100621f,   0.200699f,    -0.0985915f,   0.134706f,    -0.256396f,
+  0.393427f,    0.119606f,    -0.214278f,    -0.0183637f,  0.194266f,
+  -0.238025f,   0.182203f,    0.599718f,     0.846933f,    0.0607852f,
+  -0.183434f,   -0.723743f,   -0.72414f,     -0.124701f,   0.0227527f,
+  -0.0664636f,  -0.0385867f,  -0.0257377f,   -0.149054f,   0.12077f,
+  0.678029f,    -0.624456f,   0.189644f,     -0.518604f,   0.134397f,
+  -0.189777f,   -0.309376f,   -0.00377086f,  0.701132f,    -0.170915f,
+  0.00736111f,  -0.121906f,   0.329136f,     0.165514f,    0.0328356f,
+  0.171275f,    0.248619f,    0.247704f,     -0.449933f,   0.0841684f,
+  0.136982f,    0.122703f,    -0.0169439f,   -0.0726496f,  0.302648f,
+  -0.128556f,   0.0667425f,   -0.289717f,    -0.207532f,   -1.20269f,
+  -0.68892f,    0.045259f,    0.0973945f,    0.0988314f,   -0.944748f,
+  -0.180401f,   0.134331f,    0.033834f,     0.109023f,    0.265723f,
+  0.38063f,     -0.106518f,   -0.0686953f,   0.3744f,      -1.0957f,
+  0.0302782f,   0.0515164f,   0.00188222f,   0.0014413f,   -0.0404425f,
+  0.0124618f,   -0.0828645f,  0.506166f,     -0.776352f,   -0.405138f,
+  -0.123887f,   0.0732116f,   0.379928f,     0.604524f,    -0.492317f,
+  0.439191f,    0.0744193f,   0.389101f,     0.0604518f,   0.0943165f,
+  0.0339942f,   0.0917975f,   0.0161988f,    0.512227f,    0.538021f,
+  -0.411495f,   0.307281f,    0.33746f,      -0.218639f,   0.265742f,
+  0.39738f,     -0.12442f,    0.125236f,     -0.0845223f,  -0.150396f,
+  0.0334878f,   -0.00391915f, 0.0406864f,    -0.0487059f,  0.0377073f
+};
+
+static const float av1_simple_motion_search_split_hiddenlayer_0_bias_16[] = {
+  0.0535976f, -0.0130279f, 0.150146f,   -0.511132f, -0.357698f,
+  0.6719f,    -1.27877f,   -0.0208048f, 0.0961914f, 0.263603f,
+  0.704574f,  -1.48998f,   0.728063f,   0.941829f,  -0.199981f,
+  0.797802f,  -0.29816f,   -0.60894f,   -0.116624f, -1.16723f
+};
+
+static const float av1_simple_motion_search_split_logits_kernel_16[] = {
+  0.343153f,   -0.2110482f, -0.487199f,   0.3274144f, -2.1975f,
+  -0.6051438f, 0.1901127f,  0.4741924f,   -0.24029f,  -0.185018f,
+  -0.652635f,  2.57714f,    -0.31033031f, -0.307222f, 0.329035f,
+  -0.430181f,  0.3429f,     0.742292f,    0.3269808f, 0.4142165f
+};
+
+static const float av1_simple_motion_search_split_logits_bias_16[] = {
+  -0.783658f
+};
+
+static const NN_CONFIG av1_simple_motion_search_split_nn_config_16 = {
+  NUM_FEATURES_16,
+  NUM_LOGITS_16,
+  NUM_HIDDEN_LAYERS_16,
+  {
+      NUM_LAYER_0_UNITS_16,
+  },
+  {
+      av1_simple_motion_search_split_hiddenlayer_0_kernel_16,
+      av1_simple_motion_search_split_logits_kernel_16,
+  },
+  {
+      av1_simple_motion_search_split_hiddenlayer_0_bias_16,
+      av1_simple_motion_search_split_logits_bias_16,
+  },
+};
+
+#undef NUM_HIDDEN_LAYERS_16
+#undef NUM_FEATURES_16
+#undef NUM_LAYER_0_UNITS_16
+#undef NUM_LOGITS_16
+
 #define NUM_HIDDEN_LAYERS_8 1
-#define NUM_FEATURES_8 6
-#define NUM_LAYER_0_UNITS_8 16
+#define NUM_FEATURES_8 17
+#define NUM_LAYER_0_UNITS_8 20
 #define NUM_LOGITS_8 1
 
-static const float av1_simple_motion_search_based_split_layer_0_kernel_8[] = {
-  0.0370236f,   -0.580211f,  2.0134f,    1.69637f,    2.43181f,   -0.521648f,
-  -0.00375187f, 0.122712f,   -4.74411f,  7.36187f,    5.42574f,   -5.53557f,
-  0.0993344f,   -0.358843f,  0.0765453f, -0.615987f,  -0.754633f, -0.175846f,
-  0.714976f,    0.492862f,   0.346604f,  -1.23922f,   -2.67031f,  2.12749f,
-  1.71511f,     -1.4239f,    2.09396f,   2.42478f,    2.40151f,   2.90487f,
-  0.540813f,    -0.0954257f, -4.57571f,  -4.88078f,   -4.62386f,  -5.75167f,
-  1.35351f,     -1.08114f,   1.43744f,   1.44333f,    0.608153f,  0.193742f,
-  -0.405512f,   -0.155164f,  0.0771456f, -0.473182f,  -0.057984f, 0.140435f,
-  0.743021f,    -0.418589f,  -0.377622f, -0.531411f,  -0.668025f, -0.826607f,
-  1.37834f,     -1.07753f,   0.870466f,  0.516756f,   0.708689f,  0.286795f,
-  -3.97895f,    -0.338629f,  2.79427f,   1.80561f,    1.46275f,   1.50438f,
-  0.0232533f,   -0.43174f,   -0.348251f, 0.0863006f,  0.0321103f, 0.129674f,
-  -1.12024f,    -0.0990596f, -0.283472f, -0.238713f,  -0.239175f, -0.40816f,
-  -0.00106566f, 0.0972736f,  5.19284f,   -3.70862f,   6.39657f,   -5.27588f,
-  -2.08003f,    0.38825f,    2.38771f,   -1.27501f,   -2.45619f,  3.07324f,
-  0.616966f,    -0.451472f,  -0.319365f, 0.00807278f, -0.303261f, -0.351679f
+static const float av1_simple_motion_search_split_hiddenlayer_0_kernel_8[] = {
+  0.079443f,   -1.04068f,   0.336819f,    -0.20901f,   0.796251f,
+  0.181066f,   0.0118876f,  -0.207145f,   0.250671f,   -0.402119f,
+  -0.0847227f, 1.88683f,    0.303469f,    0.0718458f,  0.0338589f,
+  0.158896f,   0.0540238f,  -0.385426f,   0.955925f,   0.424506f,
+  0.492584f,   -0.795058f,  -0.248667f,   -0.905349f,  -0.316989f,
+  0.545471f,   0.63762f,    -0.232613f,   -0.238947f,  -0.395338f,
+  -0.322673f,  -0.0761563f, -0.125357f,   0.0694415f,  -0.371599f,
+  0.358387f,   -0.486841f,  0.403863f,    -0.0295666f, 0.283074f,
+  -0.424396f,  0.156318f,   -0.685355f,   0.6663f,     0.337949f,
+  0.273198f,   0.517448f,   0.458911f,    0.157252f,   0.692096f,
+  0.64965f,    -0.23987f,   -1.08431f,    -0.252475f,  -0.332614f,
+  -0.712291f,  -0.380973f,  0.460545f,    0.48936f,    0.337601f,
+  0.489223f,   1.65336f,    -0.223585f,   0.17367f,    -0.235057f,
+  -0.456773f,  0.327877f,   -0.221192f,   -0.940151f,  -1.06616f,
+  0.687084f,   -0.109973f,  0.106636f,    0.445895f,   0.163432f,
+  0.378306f,   0.201902f,   0.176811f,    0.693082f,   1.62156f,
+  -0.178346f,  0.455175f,   1.61943f,     0.231376f,   0.0890932f,
+  -0.889693f,  -1.03298f,   0.778196f,    -0.0289539f, 0.137848f,
+  0.18707f,    0.171889f,   0.119157f,    0.24893f,    -0.313628f,
+  0.00250735f, -0.0758209f, 0.272974f,    -0.229825f,  2.47926f,
+  -0.0354665f, 0.175366f,   0.0411555f,   -1.52149f,   -0.0258663f,
+  0.253027f,   -0.0520839f, -0.0189782f,  0.362387f,   -0.371154f,
+  0.622929f,   0.0447056f,  0.242529f,    -0.168391f,  0.308935f,
+  -0.117294f,  2.16307f,    0.0673638f,   0.080771f,   -0.460779f,
+  -0.940176f,  0.473266f,   -0.0125302f,  0.475145f,   -0.218187f,
+  0.43258f,    -0.0380196f, 0.413607f,    -0.110856f,  -1.52076f,
+  0.0896812f,  0.246636f,   -0.0612008f,  0.189583f,   0.0106902f,
+  -0.158403f,  -0.629377f,  -0.0634279f,  -0.0864584f, -0.226568f,
+  -0.286234f,  -0.0721132f, -0.43702f,    0.113702f,   0.433372f,
+  0.743396f,   0.14312f,    0.29914f,     0.801188f,   0.7609f,
+  0.385046f,   0.480314f,   0.171119f,    -1.59058f,   -1.18853f,
+  0.150676f,   0.408123f,   -0.00677924f, 0.398145f,   0.0914611f,
+  0.176945f,   0.0677457f,  0.316478f,    0.998219f,   -0.22618f,
+  0.0756793f,  -0.0156674f, 0.105716f,    0.0496245f,  -0.0827133f,
+  -0.423119f,  -0.161033f,  0.212962f,    -0.234453f,  0.743366f,
+  1.04108f,    0.0597604f,  -0.285993f,   -0.114829f,  -0.557364f,
+  -0.840051f,  0.326509f,   -0.192508f,   -0.141769f,  0.370626f,
+  -0.126353f,  0.00672923f, 0.493623f,    -0.852076f,  0.466798f,
+  -0.226436f,  0.259268f,   -0.452662f,   0.0721126f,  0.0198245f,
+  0.2048f,     0.02506f,    0.316194f,    0.814651f,   1.01288f,
+  -0.569607f,  -0.0838994f, 1.37146f,     -0.613135f,  0.441761f,
+  -0.643901f,  0.364269f,   -0.147177f,   0.338001f,   -0.332376f,
+  0.518875f,   -0.628964f,  -0.291889f,   -0.050736f,  0.108047f,
+  1.05673f,    0.0479492f,  0.466756f,    -0.0867334f, -0.0355575f,
+  0.57626f,    -0.227583f,  -0.146421f,   0.0990489f,  0.117351f,
+  -0.103858f,  -0.0336936f, 0.0201903f,   -0.0766383f, -0.010211f,
+  0.0400779f,  0.0725462f,  0.137142f,    0.478261f,   0.287869f,
+  0.0882359f,  -0.739754f,  -0.853521f,   -0.43703f,   0.316856f,
+  0.27593f,    0.312149f,   0.175575f,    0.441839f,   0.264325f,
+  0.0148051f,  -0.005559f,  0.373176f,    0.933701f,   -0.0197615f,
+  0.0219723f,  -0.0559883f, -0.103456f,   -0.0323009f, 0.0773202f,
+  -0.390838f,  0.855488f,   -0.596525f,   -0.249093f,  0.124262f,
+  0.220172f,   0.0552478f,  1.04041f,     -0.960992f,  -0.495255f,
+  -0.211612f,  0.350007f,   -0.238998f,   -0.0265068f, 0.384686f,
+  -0.0815808f, -0.0570019f, 0.123903f,    -0.485114f,  -0.00282573f,
+  -0.0649603f, 0.163719f,   -0.469479f,   -0.439713f,  0.0602562f,
+  -0.527993f,  -0.111458f,  2.48686f,     -0.180723f,  0.0553895f,
+  0.0560679f,  -0.0978928f, -0.216063f,   0.089457f,   -1.5602f,
+  -1.62332f,   -0.147388f,  0.736155f,    0.440409f,   0.243519f,
+  0.0622638f,  0.522932f,   0.109686f,    0.422849f,   0.510589f,
+  1.01116f,    0.174019f,   0.0191171f,   -0.0717751f, -0.0068308f,
+  0.172932f,   -0.834888f,  -0.635788f,   0.32012f,    0.298656f,
+  0.274309f,   -0.155456f,  0.1755f,      -0.175171f,  0.343498f,
+  -0.122832f,  -0.107696f,  0.279924f,    -0.797633f,  -0.344658f,
+  0.162669f,   0.389092f,   0.644479f,    -0.635216f,  -0.181868f,
+  0.0579244f,  -0.0568976f, 0.433003f,    -0.591067f,  0.71013f,
+  -0.165515f,  0.225725f,   -0.358156f,   0.0541944f,  1.95485f,
+  -0.315223f,  0.61537f,    -0.0401568f,  0.22811f,    0.271147f
 };
 
-static const float av1_simple_motion_search_based_split_logits_kernel_8[] = {
-  -0.625847f, 0.381323f, 0.342475f, 0.526161f,  -0.665965f, -0.515317f,
-  -0.406218f, 0.568007f, 0.479397f, -0.426116f, 0.615638f,  0.338572f,
-  0.185583f,  0.308031f, 0.260748f, 0.531619f
+static const float av1_simple_motion_search_split_hiddenlayer_0_bias_8[] = {
+  1.63441f,  -0.616459f, -0.437775f, -0.71669f,  1.56616f,  2.28109f, 1.64054f,
+  -1.51476f, 0.0274108f, 0.935156f,  -0.966329f, 0.906069f, 1.19954f, -1.25867f,
+  -1.7376f,  -0.594211f, 0.322242f,  0.438631f,  -1.01682f, 1.30032f
 };
 
-static const float av1_simple_motion_search_based_split_layer_0_bias_8[] = {
-  4.73775f,  -1.12658f, -0.258038f, -6.06696f, 1.79131f, 2.49609f,
-  4.28388f,  0.0f,      -4.63598f,  3.06034f,  5.31994f, -0.152142f,
-  0.514738f, -1.30098f, 3.00296f,   -3.83481f
+static const float av1_simple_motion_search_split_logits_kernel_8[] = {
+  -0.463187f, 0.2936127f, 0.16762f,    -0.1663271f, -0.292418f,
+  -0.421457f, -0.378265f, 1.053049f,   0.32432879f, -0.49775575f,
+  0.427357f,  -0.239251f, -0.1631546f, 0.335468f,   0.255371f,
+  0.276901f,  -0.665683f, -0.7021493f, 0.381513f,   -0.1339761f
 };
 
-static const float av1_simple_motion_search_based_split_logits_bias_8[] = {
-  -3.44508f
+static const float av1_simple_motion_search_split_logits_bias_8[] = {
+  -1.739754f
 };
 
-static const NN_CONFIG av1_simple_motion_search_based_split_nn_config_8 = {
+static const NN_CONFIG av1_simple_motion_search_split_nn_config_8 = {
   NUM_FEATURES_8,
   NUM_LOGITS_8,
   NUM_HIDDEN_LAYERS_8,
@@ -2750,16 +3341,19 @@
       NUM_LAYER_0_UNITS_8,
   },
   {
-      av1_simple_motion_search_based_split_layer_0_kernel_8,
-      av1_simple_motion_search_based_split_logits_kernel_8,
+      av1_simple_motion_search_split_hiddenlayer_0_kernel_8,
+      av1_simple_motion_search_split_logits_kernel_8,
   },
   {
-      av1_simple_motion_search_based_split_layer_0_bias_8,
-      av1_simple_motion_search_based_split_logits_bias_8,
+      av1_simple_motion_search_split_hiddenlayer_0_bias_8,
+      av1_simple_motion_search_split_logits_bias_8,
   },
 };
 
-#endif
+#undef NUM_HIDDEN_LAYERS_8
+#undef NUM_FEATURES_8
+#undef NUM_LAYER_0_UNITS_8
+#undef NUM_LOGITS_8
 
 // Model based on simple_motion_search
 
diff --git a/av1/encoder/partition_strategy.c b/av1/encoder/partition_strategy.c
index 23c2ef4..736ddf0 100644
--- a/av1/encoder/partition_strategy.c
+++ b/av1/encoder/partition_strategy.c
@@ -21,6 +21,10 @@
 #include "av1/encoder/partition_strategy.h"
 #include "av1/encoder/rdopt.h"
 
+static void simple_motion_search_prune_part_features(
+    AV1_COMP *const cpi, MACROBLOCK *x, PC_TREE *pc_tree, int mi_row,
+    int mi_col, BLOCK_SIZE bsize, float *features, int features_to_get);
+
 // Performs a simple_motion_search with a single reference frame and extract
 // the variance of residues. Here features is assumed to be a length 6 array.
 // After this function is called, we will store the following in to features:
@@ -80,13 +84,13 @@
   }
 }
 
-void av1_simple_motion_search_based_split(
+static void simple_motion_search_based_split_fast(
     AV1_COMP *const cpi, MACROBLOCK *x, int mi_row, int mi_col,
     BLOCK_SIZE bsize, int *partition_none_allowed, int *partition_horz_allowed,
     int *partition_vert_allowed, int *do_rectangular_split,
     int *do_square_split) {
   const NN_CONFIG *nn_config = NULL;
-  float split_only_thresh = 0.0f;
+  float split_only_thresh = 1.0f;
   if (bsize == BLOCK_128X128) {
     nn_config = &av1_simple_motion_search_based_split_nn_config_128;
     split_only_thresh = av1_simple_motion_search_based_split_thresh_128;
@@ -100,35 +104,99 @@
     nn_config = &av1_simple_motion_search_based_split_nn_config_16;
     split_only_thresh = av1_simple_motion_search_based_split_thresh_16;
   } else if (bsize == BLOCK_8X8) {
-    // Disable BLOCK_8X8 for now
-#if !CONFIG_DISABLE_FULL_PIXEL_SPLIT_8X8
-    nn_config = &av1_simple_motion_search_based_split_nn_config_8;
-    split_only_thresh = av1_simple_motion_search_based_split_thresh_8;
-#endif
+    return;
   } else {
     assert(0 && "Unexpected block size in simple_motion_based_split");
+    return;
   }
-  if (nn_config) {
-    float features[6] = { 0 };
-    float score = 0;
-    get_res_var_features(cpi, x, mi_row, mi_col, bsize, features);
-    av1_nn_predict(features, nn_config, &score);
 
-    if (score > split_only_thresh) {
-      *partition_none_allowed = 0;
-      *partition_horz_allowed = 0;
-      *partition_vert_allowed = 0;
-      *do_rectangular_split = 0;
-    }
-    if (cpi->sf.simple_motion_search_split_only >= 2) {
-      if (score < -split_only_thresh) *do_square_split = 0;
-      // For larger scores (>split_only_thresh), none and rectangular partitions
-      // are skipped. As score reduces, possibility of split decreases. Hence
-      // for near larger scores (.875 * split_only_thresh to split_only_thresh)
-      // none partition is disabled, but rectangular partitions are evaluated
-      // additionally.
-      if (score > (split_only_thresh * 0.875)) *partition_none_allowed = 0;
-    }
+  float features[FEATURE_SIZE_SMS_SPLIT_FAST] = { 0.0f };
+  float score = 0.0f;
+  get_res_var_features(cpi, x, mi_row, mi_col, bsize, features);
+  av1_nn_predict(features, nn_config, &score);
+
+  if (score > split_only_thresh) {
+    *partition_none_allowed = 0;
+    *partition_horz_allowed = 0;
+    *partition_vert_allowed = 0;
+    *do_rectangular_split = 0;
+  }
+  if (cpi->sf.simple_motion_search_split_only >= 2) {
+    if (score < -split_only_thresh) *do_square_split = 0;
+    // For larger scores (>split_only_thresh), none and rectangular partitions
+    // are skipped. As score reduces, possibility of split decreases. Hence
+    // for near larger scores (.875 * split_only_thresh to split_only_thresh)
+    // none partition is disabled, but rectangular partitions are evaluated
+    // additionally.
+    if (score > (split_only_thresh * 0.875)) *partition_none_allowed = 0;
+  }
+}
+
+void av1_simple_motion_search_based_split(
+    AV1_COMP *const cpi, MACROBLOCK *x, PC_TREE *pc_tree, int mi_row,
+    int mi_col, BLOCK_SIZE bsize, int *partition_none_allowed,
+    int *partition_horz_allowed, int *partition_vert_allowed,
+    int *do_rectangular_split, int *do_square_split) {
+  if (cpi->sf.simple_motion_search_split_speed >= 2) {
+    simple_motion_search_based_split_fast(
+        cpi, x, mi_row, mi_col, bsize, partition_none_allowed,
+        partition_horz_allowed, partition_vert_allowed, do_rectangular_split,
+        do_square_split);
+
+    return;
+  }
+
+  const NN_CONFIG *nn_config = NULL;
+  const float *ml_mean = NULL, *ml_std = NULL;
+  float split_only_thresh = 1.0f;
+  if (bsize == BLOCK_128X128) {
+    ml_mean = av1_simple_motion_search_split_mean_128;
+    ml_std = av1_simple_motion_search_split_std_128;
+    nn_config = &av1_simple_motion_search_split_nn_config_128;
+    split_only_thresh = av1_simple_motion_search_split_thresh_128;
+  } else if (bsize == BLOCK_64X64) {
+    ml_mean = av1_simple_motion_search_split_mean_64;
+    ml_std = av1_simple_motion_search_split_std_64;
+    nn_config = &av1_simple_motion_search_split_nn_config_64;
+    split_only_thresh = av1_simple_motion_search_split_thresh_64;
+  } else if (bsize == BLOCK_32X32) {
+    ml_mean = av1_simple_motion_search_split_mean_32;
+    ml_std = av1_simple_motion_search_split_std_32;
+    nn_config = &av1_simple_motion_search_split_nn_config_32;
+    split_only_thresh = av1_simple_motion_search_split_thresh_32;
+  } else if (bsize == BLOCK_16X16) {
+    ml_mean = av1_simple_motion_search_split_mean_16;
+    ml_std = av1_simple_motion_search_split_std_16;
+    nn_config = &av1_simple_motion_search_split_nn_config_16;
+    split_only_thresh = av1_simple_motion_search_split_thresh_16;
+  } else if (bsize == BLOCK_8X8) {
+    ml_mean = av1_simple_motion_search_split_mean_8;
+    ml_std = av1_simple_motion_search_split_std_8;
+    nn_config = &av1_simple_motion_search_split_nn_config_8;
+    split_only_thresh = av1_simple_motion_search_split_thresh_8;
+  } else {
+    assert(0 && "Unexpected block size in simple_motion_based_split");
+    return;
+  }
+
+  float features[FEATURE_SIZE_SMS_SPLIT] = { 0.0f };
+  simple_motion_search_prune_part_features(cpi, x, pc_tree, mi_row, mi_col,
+                                           bsize, features,
+                                           FEATURE_SMS_SPLIT_MODEL_FLAG);
+  for (int idx = 0; idx < FEATURE_SIZE_SMS_SPLIT; idx++) {
+    features[idx] = (features[idx] - ml_mean[idx]) / ml_std[idx];
+  }
+
+  float score = 0.0f;
+
+  av1_nn_predict(features, nn_config, &score);
+  aom_clear_system_state();
+
+  if (score > split_only_thresh) {
+    *partition_none_allowed = 0;
+    *partition_horz_allowed = 0;
+    *partition_vert_allowed = 0;
+    *do_rectangular_split = 0;
   }
 }
 
@@ -211,12 +279,11 @@
 // features[18] = DC q_index
 static void simple_motion_search_prune_part_features(
     AV1_COMP *const cpi, MACROBLOCK *x, PC_TREE *pc_tree, int mi_row,
-    int mi_col, BLOCK_SIZE bsize, float *features) {
+    int mi_col, BLOCK_SIZE bsize, float *features, int features_to_get) {
   // TODO(chiyotsai@google.com): Cache the result of the motion search from the
   // larger bsize.
   const int w_mi = mi_size_wide[bsize];
   const int h_mi = mi_size_high[bsize];
-  int f_idx = 0;
   assert(mi_size_wide[bsize] == mi_size_high[bsize]);
   assert(cpi->ref_frame_flags & av1_ref_frame_flag_list[LAST_FRAME] ||
          cpi->ref_frame_flags & av1_ref_frame_flag_list[ALTREF_FRAME]);
@@ -227,56 +294,77 @@
   const int num_refs = 1;
   const int use_subpixel = 1;
 
-  unsigned int int_features[FEATURE_SIZE_SMS_PRUNE_PART - 1];
-
   // Doing whole block first to update the mv
-  simple_motion_search_get_best_ref(
-      cpi, x, pc_tree, mi_row, mi_col, bsize, ref_list, num_refs, use_subpixel,
-      4, &int_features[f_idx], &int_features[f_idx + 1]);
-  f_idx += 2;
+  if (!pc_tree->sms_none_valid && features_to_get & FEATURE_SMS_NONE_FLAG) {
+    simple_motion_search_get_best_ref(cpi, x, pc_tree, mi_row, mi_col, bsize,
+                                      ref_list, num_refs, use_subpixel, 4,
+                                      &pc_tree->sms_none_feat[0],
+                                      &pc_tree->sms_none_feat[1]);
+    pc_tree->sms_none_valid = 1;
+  }
 
   // Split subblocks
-  BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_SPLIT);
-  int r_idx = 0;
-  for (r_idx = 0; r_idx < 4; r_idx++) {
-    const int sub_mi_col = mi_col + (r_idx & 1) * w_mi / 2;
-    const int sub_mi_row = mi_row + (r_idx >> 1) * h_mi / 2;
+  if (!pc_tree->sms_split_valid && features_to_get & FEATURE_SMS_SPLIT_FLAG) {
+    const BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_SPLIT);
+    for (int r_idx = 0; r_idx < 4; r_idx++) {
+      const int sub_mi_col = mi_col + (r_idx & 1) * w_mi / 2;
+      const int sub_mi_row = mi_row + (r_idx >> 1) * h_mi / 2;
 
-    simple_motion_search_get_best_ref(
-        cpi, x, pc_tree, sub_mi_row, sub_mi_col, subsize, ref_list, num_refs,
-        use_subpixel, r_idx, &int_features[f_idx], &int_features[f_idx + 1]);
-    f_idx += 2;
+      simple_motion_search_get_best_ref(
+          cpi, x, pc_tree, sub_mi_row, sub_mi_col, subsize, ref_list, num_refs,
+          use_subpixel, r_idx, &pc_tree->sms_split_feat[2 * r_idx],
+          &pc_tree->sms_split_feat[2 * r_idx + 1]);
+    }
+
+    pc_tree->sms_split_valid = 1;
   }
 
-  // Horz subblocks
-  subsize = get_partition_subsize(bsize, PARTITION_HORZ);
-  for (r_idx = 0; r_idx < 2; r_idx++) {
-    const int sub_mi_col = mi_col + 0;
-    const int sub_mi_row = mi_row + r_idx * h_mi / 2;
+  // Rectangular subblocks
+  if (!pc_tree->sms_rect_valid && features_to_get & FEATURE_SMS_RECT_FLAG) {
+    // Horz subblock
+    BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_HORZ);
+    for (int r_idx = 0; r_idx < 2; r_idx++) {
+      const int sub_mi_col = mi_col + 0;
+      const int sub_mi_row = mi_row + r_idx * h_mi / 2;
 
-    simple_motion_search_get_best_ref(
-        cpi, x, pc_tree, sub_mi_row, sub_mi_col, subsize, ref_list, num_refs,
-        use_subpixel, -1, &int_features[f_idx], &int_features[f_idx + 1]);
+      simple_motion_search_get_best_ref(
+          cpi, x, pc_tree, sub_mi_row, sub_mi_col, subsize, ref_list, num_refs,
+          use_subpixel, -1, &pc_tree->sms_rect_feat[2 * r_idx],
+          &pc_tree->sms_rect_feat[2 * r_idx + 1]);
+    }
 
-    f_idx += 2;
-  }
+    // Vert subblock
+    subsize = get_partition_subsize(bsize, PARTITION_VERT);
+    for (int r_idx = 0; r_idx < 2; r_idx++) {
+      const int sub_mi_col = mi_col + r_idx * w_mi / 2;
+      const int sub_mi_row = mi_row + 0;
 
-  // Vert subblock
-  subsize = get_partition_subsize(bsize, PARTITION_VERT);
-  for (r_idx = 0; r_idx < 2; r_idx++) {
-    const int sub_mi_col = mi_col + r_idx * w_mi / 2;
-    const int sub_mi_row = mi_row + 0;
-
-    simple_motion_search_get_best_ref(
-        cpi, x, pc_tree, sub_mi_row, sub_mi_col, subsize, ref_list, num_refs,
-        use_subpixel, -1, &int_features[f_idx], &int_features[f_idx + 1]);
-
-    f_idx += 2;
+      simple_motion_search_get_best_ref(
+          cpi, x, pc_tree, sub_mi_row, sub_mi_col, subsize, ref_list, num_refs,
+          use_subpixel, -1, &pc_tree->sms_rect_feat[4 + 2 * r_idx],
+          &pc_tree->sms_rect_feat[4 + 2 * r_idx + 1]);
+    }
+    pc_tree->sms_rect_valid = 1;
   }
 
   aom_clear_system_state();
-  for (int idx = 0; idx < f_idx; idx++) {
-    features[idx] = logf(1.0f + (float)int_features[idx]);
+  int f_idx = 0;
+  if (features_to_get & FEATURE_SMS_NONE_FLAG) {
+    for (int sub_idx = 0; sub_idx < 2; sub_idx++) {
+      features[f_idx++] = logf(1.0f + pc_tree->sms_none_feat[sub_idx]);
+    }
+  }
+
+  if (features_to_get & FEATURE_SMS_SPLIT_FLAG) {
+    for (int sub_idx = 0; sub_idx < 8; sub_idx++) {
+      features[f_idx++] = logf(1.0f + pc_tree->sms_split_feat[sub_idx]);
+    }
+  }
+
+  if (features_to_get & FEATURE_SMS_RECT_FLAG) {
+    for (int sub_idx = 0; sub_idx < 8; sub_idx++) {
+      features[f_idx++] = logf(1.0f + pc_tree->sms_rect_feat[sub_idx]);
+    }
   }
 
   const MACROBLOCKD *xd = &x->e_mbd;
@@ -297,8 +385,6 @@
   features[f_idx++] = (float)has_left;
   features[f_idx++] = (float)mi_size_wide_log2[left_bsize];
   features[f_idx++] = (float)mi_size_high_log2[left_bsize];
-
-  assert(f_idx == FEATURE_SIZE_SMS_PRUNE_PART);
 }
 
 void av1_simple_motion_search_prune_part(
@@ -306,13 +392,13 @@
     int mi_col, BLOCK_SIZE bsize, int *partition_none_allowed,
     int *partition_horz_allowed, int *partition_vert_allowed,
     int *do_square_split, int *do_rectangular_split, int *prune_horz,
-    int *prune_vert, float *features, int *valid) {
+    int *prune_vert) {
   const AV1_COMMON *const cm = &cpi->common;
   // Get model parameters
   const NN_CONFIG *nn_config = NULL;
   const float *prune_thresh = NULL, *only_thresh = NULL;
   const float *ml_mean = NULL, *ml_std = NULL;
-  float normalized_features[FEATURE_SIZE_SMS_PRUNE_PART] = { 0.0f };
+  float features[FEATURE_SIZE_SMS_PRUNE_PART] = { 0.0f };
 
   if (bsize == BLOCK_128X128) {
     nn_config = &av1_simple_motion_search_prune_part_nn_config_128;
@@ -359,11 +445,10 @@
 
   // Get features
   simple_motion_search_prune_part_features(cpi, x, pc_tree, mi_row, mi_col,
-                                           bsize, features);
-  *valid = 1;
+                                           bsize, features,
+                                           FEATURE_SMS_PRUNE_PART_FLAG);
   for (int f_idx = 0; f_idx < FEATURE_SIZE_SMS_PRUNE_PART; f_idx++) {
-    normalized_features[f_idx] =
-        (features[f_idx] - ml_mean[f_idx]) / ml_std[f_idx];
+    features[f_idx] = (features[f_idx] - ml_mean[f_idx]) / ml_std[f_idx];
   }
 
   // Get probabilities
@@ -373,7 +458,7 @@
                               ? PARTITION_TYPES
                               : EXT_PARTITION_TYPES;
 
-  av1_nn_predict(normalized_features, nn_config, scores);
+  av1_nn_predict(features, nn_config, scores);
   aom_clear_system_state();
 
   av1_nn_softmax(scores, probs, num_classes);
@@ -399,24 +484,23 @@
 //  - The frame is not intra only
 //  - The current bsize is > BLOCK_8X8
 //  - blk_row + blk_height/2 < total_rows and blk_col + blk_width/2 < total_cols
-void av1_simple_motion_search_early_term_none(
-    AV1_COMP *const cpi, MACROBLOCK *x, PC_TREE *pc_tree, int mi_row,
-    int mi_col, BLOCK_SIZE bsize, const RD_STATS *none_rdc,
-    int *early_terminate, float *simple_motion_features,
-    int *simple_motion_features_are_valid) {
+void av1_simple_motion_search_early_term_none(AV1_COMP *const cpi,
+                                              MACROBLOCK *x, PC_TREE *pc_tree,
+                                              int mi_row, int mi_col,
+                                              BLOCK_SIZE bsize,
+                                              const RD_STATS *none_rdc,
+                                              int *early_terminate) {
   // TODO(chiyotsai@google.com): There are other features we can extract from
   // PARTITION_NONE. Play with this later.
-  int f_idx = 0;
-  if (!*simple_motion_features_are_valid) {
-    simple_motion_search_prune_part_features(cpi, x, pc_tree, mi_row, mi_col,
-                                             bsize, simple_motion_features);
-    *simple_motion_features_are_valid = 1;
-  }
-  f_idx = 25;
+  float features[FEATURE_SIZE_SMS_TERM_NONE] = { 0.0f };
+  simple_motion_search_prune_part_features(cpi, x, pc_tree, mi_row, mi_col,
+                                           bsize, features,
+                                           FEATURE_SMS_PRUNE_PART_FLAG);
+  int f_idx = FEATURE_SIZE_SMS_PRUNE_PART;
 
-  simple_motion_features[f_idx++] = logf(1.0f + (float)none_rdc->rate);
-  simple_motion_features[f_idx++] = logf(1.0f + (float)none_rdc->dist);
-  simple_motion_features[f_idx++] = logf(1.0f + (float)none_rdc->rdcost);
+  features[f_idx++] = logf(1.0f + (float)none_rdc->rate);
+  features[f_idx++] = logf(1.0f + (float)none_rdc->dist);
+  features[f_idx++] = logf(1.0f + (float)none_rdc->rdcost);
 
   assert(f_idx == FEATURE_SIZE_SMS_TERM_NONE);
 
@@ -447,8 +531,8 @@
   if (ml_model) {
     float score = 0.0f;
     for (f_idx = 0; f_idx < FEATURE_SIZE_SMS_TERM_NONE; f_idx++) {
-      score += ml_model[f_idx] *
-               (simple_motion_features[f_idx] - ml_mean[f_idx]) / ml_std[f_idx];
+      score +=
+          ml_model[f_idx] * (features[f_idx] - ml_mean[f_idx]) / ml_std[f_idx];
     }
     score += ml_model[FEATURE_SIZE_SMS_TERM_NONE];
 
diff --git a/av1/encoder/partition_strategy.h b/av1/encoder/partition_strategy.h
index 36b1e95..7504d67 100644
--- a/av1/encoder/partition_strategy.h
+++ b/av1/encoder/partition_strategy.h
@@ -16,20 +16,31 @@
 #include "av1/encoder/encodemb.h"
 #include "av1/encoder/encoder.h"
 
+#define FEATURE_SIZE_SMS_SPLIT_FAST 6
+#define FEATURE_SIZE_SMS_SPLIT 17
 #define FEATURE_SIZE_SMS_PRUNE_PART 25
 #define FEATURE_SIZE_SMS_TERM_NONE 28
 #define FEATURE_SIZE_FP_SMS_TERM_NONE 20
 #define FEATURE_SIZE_MAX_MIN_PART_PRED 13
 #define MAX_NUM_CLASSES_MAX_MIN_PART_PRED 4
 
+#define FEATURE_SMS_NONE_FLAG 1
+#define FEATURE_SMS_SPLIT_FLAG (1 << 1)
+#define FEATURE_SMS_RECT_FLAG (1 << 2)
+
+#define FEATURE_SMS_PRUNE_PART_FLAG \
+  (FEATURE_SMS_NONE_FLAG | FEATURE_SMS_SPLIT_FLAG | FEATURE_SMS_RECT_FLAG)
+#define FEATURE_SMS_SPLIT_MODEL_FLAG \
+  (FEATURE_SMS_NONE_FLAG | FEATURE_SMS_SPLIT_FLAG)
+
 // Performs a simple_motion_search with a single reference frame and extract
 // the variance of residues. Then use the features to determine whether we want
 // to go straight to splitting without trying PARTITION_NONE
 void av1_simple_motion_search_based_split(
-    AV1_COMP *const cpi, MACROBLOCK *x, int mi_row, int mi_col,
-    BLOCK_SIZE bsize, int *partition_none_allowed, int *partition_horz_allowed,
-    int *partition_vert_allowed, int *do_rectangular_split,
-    int *do_square_split);
+    AV1_COMP *const cpi, MACROBLOCK *x, PC_TREE *pc_tree, int mi_row,
+    int mi_col, BLOCK_SIZE bsize, int *partition_none_allowed,
+    int *partition_horz_allowed, int *partition_vert_allowed,
+    int *do_rectangular_split, int *do_square_split);
 
 // Performs a simple_motion_search with two reference frames and extract
 // the variance of residues. Then use the features to determine whether we want
@@ -39,7 +50,7 @@
     int mi_col, BLOCK_SIZE bsize, int *partition_none_allowed,
     int *partition_horz_allowed, int *partition_vert_allowed,
     int *do_square_split, int *do_rectangular_split, int *prune_horz,
-    int *prune_vert, float *features, int *valid);
+    int *prune_vert);
 
 // Early terminates PARTITION_NONE using simple_motion_search features and the
 // rate, distortion, and rdcost of PARTITION_NONE. This is only called when:
@@ -47,11 +58,12 @@
 //  - The frame is not intra only
 //  - The current bsize is > BLOCK_8X8
 //  - blk_row + blk_height/2 < total_rows and blk_col + blk_width/2 < total_cols
-void av1_simple_motion_search_early_term_none(
-    AV1_COMP *const cpi, MACROBLOCK *x, PC_TREE *pc_tree, int mi_row,
-    int mi_col, BLOCK_SIZE bsize, const RD_STATS *none_rdc,
-    int *early_terminate, float *simple_motion_features,
-    int *simple_motion_features_are_valid);
+void av1_simple_motion_search_early_term_none(AV1_COMP *const cpi,
+                                              MACROBLOCK *x, PC_TREE *pc_tree,
+                                              int mi_row, int mi_col,
+                                              BLOCK_SIZE bsize,
+                                              const RD_STATS *none_rdc,
+                                              int *early_terminate);
 
 // Early terminates after PARTITION_NONE in firstpass of two pass partition
 // search.
@@ -116,10 +128,15 @@
 }
 
 static INLINE void init_simple_motion_search_mvs(PC_TREE *pc_tree) {
-  for (int idx = 0; idx < REF_FRAMES; idx++) {
-    pc_tree->mv_ref_fulls[idx].row = 0;
-    pc_tree->mv_ref_fulls[idx].col = 0;
-  }
+  av1_zero(pc_tree->mv_ref_fulls);
+
+  av1_zero(pc_tree->sms_none_feat);
+  av1_zero(pc_tree->sms_split_feat);
+  av1_zero(pc_tree->sms_rect_feat);
+  av1_zero(pc_tree->sms_none_valid);
+  av1_zero(pc_tree->sms_split_valid);
+  av1_zero(pc_tree->sms_rect_valid);
+
   if (pc_tree->block_size >= BLOCK_8X8) {
     init_simple_motion_search_mvs(pc_tree->split[0]);
     init_simple_motion_search_mvs(pc_tree->split[1]);
diff --git a/av1/encoder/speed_features.c b/av1/encoder/speed_features.c
index 5dfc585..624d713 100644
--- a/av1/encoder/speed_features.c
+++ b/av1/encoder/speed_features.c
@@ -138,6 +138,9 @@
       sf->ml_partition_search_breakout_thresh[4] = -1;   // BLOCK_128X128
 
       sf->firstpass_simple_motion_search_early_term = 1;
+      // TODO(chiyotsai@google.com): Try to disable two pass partition search
+      // and turn on hdres
+      sf->simple_motion_search_split_speed = 1;
     }
   }
 
@@ -165,6 +168,7 @@
   }
 
   if (speed >= 3) {
+    sf->simple_motion_search_split_speed = 2;
     if (is_720p_or_larger) {
       sf->partition_search_breakout_dist_thr = (1 << 25);
       sf->partition_search_breakout_rate_thr = 200;
@@ -731,6 +735,7 @@
   sf->skip_obmc_in_uniform_mv_field = 0;
   sf->skip_wm_in_uniform_mv_field = 0;
   sf->adaptive_interp_filter_search = 0;
+  sf->simple_motion_search_split_speed = 2;
 
   for (i = 0; i < TX_SIZES; i++) {
     sf->intra_y_mode_mask[i] = INTRA_ALL;
diff --git a/av1/encoder/speed_features.h b/av1/encoder/speed_features.h
index a321192..a5a2a01 100644
--- a/av1/encoder/speed_features.h
+++ b/av1/encoder/speed_features.h
@@ -651,6 +651,12 @@
   // want to split directly without trying other partition types.
   int simple_motion_search_split_only;
 
+  // Determines the type of model used by simple_motion_search_split_only. Only
+  // valids when simple_motion_search_split_only is >= 1. Set to 1 for the
+  // slower model that uses 5 subpixel searches, and 2 for the faster model that
+  // uses 1 fullpixel search.
+  int simple_motion_search_split_speed;
+
   // Use features from simple_motion_search to terminate prediction block
   // partition after PARTITION_NONE
   int simple_motion_search_early_term_none;