Update on "Introduce public MergedDataMap"

Add public merged data map. Module can use this to resolve multiple named data maps.

Differential Revision: [D83527299](https://our.internmc.facebook.com/intern/diff/D83527299/)

[ghstack-poisoned]

Author: lucylq

extension/named_data_map/test/CMakeLists.txt

@@ -20,35 +20,23 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake)
 
 add_custom_command(
   OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.ptd"
-         "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.pte"
          "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.ptd"
-         "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.pte"
-         "${CMAKE_CURRENT_BINARY_DIR}/ModuleSimpleTrainProgram.ptd"
-         "${CMAKE_CURRENT_BINARY_DIR}/ModuleSimpleTrainProgram.pte"
   COMMAND
     ${PYTHON_EXECUTABLE} -m test.models.export_program --modules
-    "ModuleAddMul,ModuleLinear,ModuleSimpleTrain" --external-constants --outdir
+    "ModuleAddMul,ModuleLinear" --external-constants --outdir
     "${CMAKE_CURRENT_BINARY_DIR}"
   WORKING_DIRECTORY ${EXECUTORCH_ROOT}
 )
 
 add_custom_target(
   extension_named_data_map_test_resources
   DEPENDS "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.ptd"
-          "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.pte"
           "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.ptd"
-          "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.pte"
-          "${CMAKE_CURRENT_BINARY_DIR}/ModuleSimpleTrainProgram.ptd"
-          "${CMAKE_CURRENT_BINARY_DIR}/ModuleSimpleTrainProgram.pte"
 )
 
 set(test_env
     "ET_MODULE_ADD_MUL_DATA_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.ptd"
-    "ET_MODULE_ADD_MUL_PROGRAM_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.pte"
     "ET_MODULE_LINEAR_DATA_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.ptd"
-    "ET_MODULE_LINEAR_PROGRAM_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.pte"
-    "ET_MODULE_SIMPLE_TRAIN_DATA_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleSimpleTrainProgram.ptd"
-    "ET_MODULE_SIMPLE_TRAIN_PROGRAM_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleSimpleTrainProgram.pte"
 )
 
 set(_test_srcs merged_data_map_test.cpp)

extension/named_data_map/test/merged_data_map_test.cpp

@@ -60,8 +60,6 @@ class MergedDataMapTest : public ::testing::Test {
         std::getenv("ET_MODULE_ADD_MUL_DATA_PATH"), "addmul");
     load_flat_tensor_data_map(
         std::getenv("ET_MODULE_LINEAR_DATA_PATH"), "linear");
-    load_flat_tensor_data_map(
-        std::getenv("ET_MODULE_SIMPLE_TRAIN_DATA_PATH"), "simple_train");
   }
 
  private:
@@ -155,14 +153,6 @@ TEST_F(MergedDataMapTest, LoadDuplicateDataMapsFail) {
   Result<MergedDataMap> merged_map =
       MergedDataMap::load(Span<const NamedDataMap*>(ndms.data(), ndms.size()));
   EXPECT_EQ(merged_map.error(), Error::InvalidArgument);
-
-  std::vector<const NamedDataMap*> ndms2 = {
-      data_maps_["addmul"].get(),
-      data_maps_["linear"].get(),
-      data_maps_["simple_train"].get()};
-  Result<MergedDataMap> merged_map2 =
-      MergedDataMap::load(Span<const NamedDataMap*>(ndms.data(), ndms.size()));
-  EXPECT_EQ(merged_map2.error(), Error::InvalidArgument);
 }
 
 TEST_F(MergedDataMapTest, CheckDataMapContents) {
@@ -181,23 +171,4 @@ TEST_F(MergedDataMapTest, CheckDataMapContents) {
   // API calls produce equivalent results.
   compare_ndm_api_calls(data_maps_["addmul"].get(), &merged_map.get());
   compare_ndm_api_calls(data_maps_["linear"].get(), &merged_map.get());
-
-  // Check with addmul and simple_train as well.
-  std::vector<const NamedDataMap*> ndms2 = {
-      data_maps_["addmul"].get(), data_maps_["simple_train"].get()};
-  Result<MergedDataMap> merged_map2 = MergedDataMap::load(
-      Span<const NamedDataMap*>(ndms2.data(), ndms2.size()));
-  EXPECT_EQ(merged_map2.error(), Error::Ok);
-
-  // Num keys.
-  size_t addmul_num_keys2 = data_maps_["addmul"]->get_num_keys().get();
-  size_t simple_train_num_keys =
-      data_maps_["simple_train"]->get_num_keys().get();
-  EXPECT_EQ(
-      merged_map2->get_num_keys().get(),
-      addmul_num_keys2 + simple_train_num_keys);
-
-  // API calls produce equivalent results.
-  compare_ndm_api_calls(data_maps_["addmul"].get(), &merged_map2.get());
-  compare_ndm_api_calls(data_maps_["simple_train"].get(), &merged_map2.get());
 }

extension/named_data_map/test/targets.bzl

@@ -8,7 +8,6 @@ def define_common_targets(is_fbcode=False):
             # intentionally don't work in xplat (since they're host-only tools).
             "ET_MODULE_ADD_MUL_DATA_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleAddMul.ptd])",
             "ET_MODULE_LINEAR_DATA_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleLinear.ptd])",
-            "ET_MODULE_SIMPLE_TRAIN_DATA_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleSimpleTrain.ptd])",
         }
 
         runtime.cxx_test(