ExtPart: Add a variable to distinguish deicion mode

Change-Id: Ia84f517ee2ba85db79e2689c2745cc26eedc81c4

aom/aom_external_partition.h

diff --git a/aom/aom_external_partition.h b/aom/aom_external_partition.h
index b6ea96e..daf8fe4 100644
--- a/aom/aom_external_partition.h
+++ b/aom/aom_external_partition.h
@@ -81,6 +81,19 @@
  */
 #define SIZE_PRUNE_4_WAY 18
 
+/*!\brief Decision mode of the external partition model.
+ * WHOLE_TREE_DECISION: the external partition model should provide the
+ * whole partition tree for the superblock.
+ *
+ * RECURSIVE_DECISION: the external partition model provides the partition
+ * decision of the current block only. The decision process starts from
+ * the superblock size, down to the smallest block size (4x4) recursively.
+ */
+typedef enum aom_ext_part_decision_mode {
+  WHOLE_TREE_DECISION = 0,
+  RECURSIVE_DECISION = 1,
+} aom_ext_part_decision_mode_t;
+
 /*!\brief Config information sent to the external partition model.
  *
  * For example, the maximum superblock size determined by the sequence header.
@@ -368,6 +381,11 @@
   aom_ext_part_delete_model_fn_t delete_model;
 
   /*!
+   * The decision mode of the model.
+   */
+  aom_ext_part_decision_mode_t decision_mode;
+
+  /*!
    * Private data for the external partition model.
    */
   void *priv;

av1/encoder/external_partition.c

diff --git a/av1/encoder/external_partition.c b/av1/encoder/external_partition.c
index 542b2bb..79f8b4c 100644
--- a/av1/encoder/external_partition.c
+++ b/av1/encoder/external_partition.c
@@ -91,3 +91,8 @@
   if (status != AOM_EXT_PART_OK) return false;
   return true;
 }
+
+aom_ext_part_decision_mode_t av1_get_ext_part_decision_mode(
+    const ExtPartController *ext_part_controller) {
+  return ext_part_controller->funcs.decision_mode;
+}

av1/encoder/external_partition.h

diff --git a/av1/encoder/external_partition.h b/av1/encoder/external_partition.h
index 20f03ed..f74973e 100644
--- a/av1/encoder/external_partition.h
+++ b/av1/encoder/external_partition.h
@@ -47,6 +47,9 @@
 bool av1_ext_part_send_features(ExtPartController *ext_part_controller,
                                 const aom_partition_features_t *features);
 
+aom_ext_part_decision_mode_t av1_get_ext_part_decision_mode(
+    const ExtPartController *ext_part_controller);
+
 /*!\endcond */
 #ifdef __cplusplus
 }  // extern "C"

av1/encoder/partition_search.c

diff --git a/av1/encoder/partition_search.c b/av1/encoder/partition_search.c
index ccdd8e5..8a1e2aa 100644
--- a/av1/encoder/partition_search.c
+++ b/av1/encoder/partition_search.c
@@ -4320,12 +4320,17 @@
                              RD_STATS *best_rd_cost) {
   if (cpi->ext_part_controller.ready) {
     bool valid_search = true;
-    if (ML_PARTITION_WHOLE_TREE_DECISION) {
+    const aom_ext_part_decision_mode_t decision_mode =
+        av1_get_ext_part_decision_mode(&cpi->ext_part_controller);
+    if (decision_mode == WHOLE_TREE_DECISION) {
       valid_search = ml_partition_search_whole_tree(
           cpi, td, tile_data, tp, sms_root, mi_row, mi_col, bsize);
-    } else {
+    } else if (decision_mode == RECURSIVE_DECISION) {
       valid_search = ml_partition_search_partial(
           cpi, td, tile_data, tp, sms_root, mi_row, mi_col, bsize);
+    } else {
+      assert(0 && "Unknown decision mode.");
+      return false;
     }
     if (!valid_search) {
       assert(0 && "Invalid search from ML model, partition search failed.");

test/av1_external_partition_test.cc

diff --git a/test/av1_external_partition_test.cc b/test/av1_external_partition_test.cc
index 45ca847a..20329de 100644
--- a/test/av1_external_partition_test.cc
+++ b/test/av1_external_partition_test.cc
@@ -245,6 +245,7 @@
     if (video->frame() == 0) {
       aom_ext_part_funcs_t ext_part_funcs;
       ext_part_funcs.priv = reinterpret_cast<void *>(&test_data_);
+      ext_part_funcs.decision_mode = WHOLE_TREE_DECISION;
       ext_part_funcs.create_model = ext_part_create_model;
       ext_part_funcs.send_features = ext_part_send_features;
       ext_part_funcs.get_partition_decision =