Backend data separation test

exir/backend/test/demos/rpc/ExecutorBackend.cpp

@@ -18,6 +18,7 @@
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/evalue.h>
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
+#include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/executor/method.h>
 #include <executorch/runtime/executor/program.h>
 
@@ -37,6 +38,7 @@ using ::executorch::runtime::MemoryAllocator;
 using ::executorch::runtime::MemoryManager;
 using ::executorch::runtime::Method;
 using ::executorch::runtime::MethodMeta;
+using ::executorch::runtime::NamedDataMap;
 using ::executorch::runtime::Program;
 using ::executorch::runtime::Result;
 using ::executorch::runtime::Span;
@@ -156,9 +158,13 @@ class ExecutorBackend final : public ::executorch::runtime::BackendInterface {
     new (client_memory_manager)
         MemoryManager(client_method_allocator, client_planned_memory);
 
+    const NamedDataMap* named_data_map = context.get_named_data_map();
     // Construct the client Method
-    Result<Method> method_res =
-        client_program->load_method("forward", client_memory_manager);
+    Result<Method> method_res = client_program->load_method(
+        "forward",
+        client_memory_manager,
+        /*event_tracer=*/nullptr,
+        named_data_map);
     if (!method_res.ok()) {
       ET_LOG(
           Error,

exir/backend/test/demos/rpc/TARGETS

@@ -11,6 +11,7 @@ runtime.python_library(
     ],
     visibility = [
         "//executorch/exir/backend/test/...",
+        "//executorch/test/...",
     ],
     deps = [
         "//caffe2:torch",

exir/backend/test/demos/rpc/executor_backend_preprocess.py

@@ -8,6 +8,8 @@
 
 from typing import final, List
 
+from executorch.exir import ExecutorchBackendConfig
+
 from executorch.exir.backend.backend_details import (
     BackendDetails,
     ExportedProgram,
@@ -24,10 +26,14 @@ def preprocess(
         edge_program: ExportedProgram,
         compile_specs: List[CompileSpec],
     ) -> PreprocessResult:
+        config = ExecutorchBackendConfig()
+        for spec in compile_specs:
+            if spec.key == "external_constants":
+                config.external_constants = True
         return PreprocessResult(
             processed_bytes=EdgeProgramManager(
                 edge_programs=edge_program,
             )
-            .to_executorch()
+            .to_executorch(config)
             .buffer,
         )

exir/backend/test/demos/rpc/targets.bzl

@@ -40,6 +40,7 @@ def define_common_targets():
         ],
         visibility = [
             "//executorch/exir/backend/test/...",
+            "//executorch/runtime/executor/test/...",
         ],
         deps = [
             ":executor_backend",

runtime/executor/method.cpp

@@ -329,6 +329,8 @@ Result<size_t> Method::get_num_external_constants() {
 }
 
 Error Method::parse_external_constants(const NamedDataMap* named_data_map) {
+  ET_CHECK_OR_RETURN_ERROR(
+      named_data_map != nullptr, InvalidState, "named_data_map is null");
   auto flatbuffer_values = serialization_plan_->values();
   size_t n_value = flatbuffer_values->size();
 
@@ -372,6 +374,7 @@ Error Method::parse_external_constants(const NamedDataMap* named_data_map) {
     Result<const TensorLayout> tensor_layout =
         named_data_map->get_metadata(key);
     if (!tensor_layout.ok()) {
+      ET_LOG(Info, "Failed to get metadata for key %s", key);
       return tensor_layout.error();
     }
     // Check external tensor compatibility.

runtime/executor/test/backend_data_separation_test.cpp

@@ -0,0 +1,101 @@
+
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/exir/backend/test/demos/rpc/ExecutorBackend.h>
+#include <executorch/extension/data_loader/file_data_loader.h>
+#include <executorch/extension/flat_tensor/flat_tensor_data_map.h>
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/result.h>
+#include <executorch/runtime/executor/method.h>
+#include <executorch/runtime/executor/program.h>
+#include <executorch/runtime/executor/test/managed_memory_manager.h>
+#include <executorch/runtime/platform/runtime.h>
+
+#include <gtest/gtest.h>
+
+using namespace ::testing;
+using executorch::extension::FlatTensorDataMap;
+using executorch::runtime::DataLoader;
+using executorch::runtime::Error;
+using executorch::runtime::Method;
+using executorch::runtime::Program;
+using executorch::runtime::Result;
+using executorch::runtime::testing::ManagedMemoryManager;
+using torch::executor::util::FileDataLoader;
+
+constexpr size_t kDefaultNonConstMemBytes = 32 * 1024U;
+constexpr size_t kDefaultRuntimeMemBytes = 32 * 1024U;
+
+class BackendDataSeparationTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    // Since these tests cause ET_LOG to be called, the PAL must be initialized
+    // first.
+    executorch::runtime::runtime_init();
+
+    // Make sure that the backend has been registered. Safe to call multiple
+    // times. Doing this at runtime ensures that it's only registered if these
+    // tests are run.
+    ASSERT_EQ(example::register_executor_backend(), Error::Ok);
+
+    // Create data loaders.
+    Result<FileDataLoader> linear_program_loader = FileDataLoader::from(
+        std::getenv("ET_MODULE_LINEAR_DELEGATE_PROGRAM_PATH"));
+    ASSERT_EQ(linear_program_loader.error(), Error::Ok);
+    linear_program_loader_ = std::make_unique<FileDataLoader>(
+        std::move(linear_program_loader.get()));
+
+    Result<FileDataLoader> linear_data_loader =
+        FileDataLoader::from(std::getenv("ET_MODULE_LINEAR_DATA_PATH"));
+    ASSERT_EQ(linear_data_loader.error(), Error::Ok);
+    linear_data_loader_ =
+        std::make_unique<FileDataLoader>(std::move(linear_data_loader.get()));
+
+    // Create programs.
+    Result<Program> linear_program = Program::load(
+        linear_program_loader_.get(),
+        Program::Verification::InternalConsistency);
+    ASSERT_EQ(linear_program.error(), Error::Ok);
+    linear_program_ =
+        std::make_unique<Program>(std::move(linear_program.get()));
+
+    Result<FlatTensorDataMap> linear_data_map =
+        FlatTensorDataMap::load(linear_data_loader_.get());
+    EXPECT_EQ(linear_data_map.error(), Error::Ok);
+    linear_data_map_ =
+        std::make_unique<FlatTensorDataMap>(std::move(linear_data_map.get()));
+
+    ET_LOG(
+        Info,
+        "setup done, named_data_map_ = %lu",
+        linear_data_map_->get_num_keys().get());
+  }
+
+ private:
+  std::unique_ptr<FileDataLoader> linear_program_loader_;
+  std::unique_ptr<FileDataLoader> linear_data_loader_;
+
+ protected:
+  std::unique_ptr<Program> linear_program_;
+  std::unique_ptr<FlatTensorDataMap> linear_data_map_;
+};
+
+TEST_F(BackendDataSeparationTest, TestSeparation) {
+  ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
+  Result<Method> method = linear_program_->load_method(
+      "forward",
+      &mmm.get(),
+      /*event_tracer=*/nullptr,
+      /*named_data_map=*/linear_data_map_.get());
+  ASSERT_EQ(method.error(), Error::Ok);
+
+  // Can execute the method.
+  Error err = method->execute();
+  ASSERT_EQ(err, Error::Ok);
+}

runtime/executor/test/targets.bzl

@@ -240,6 +240,28 @@ def define_common_targets(is_fbcode = False):
             },
         )
 
+        runtime.cxx_test(
+            name = "backend_data_separation_test",
+            srcs = [
+                "backend_data_separation_test.cpp",
+            ],
+            deps = [
+                ":managed_memory_manager",
+                "//executorch/runtime/executor:program",
+                "//executorch/extension/data_loader:file_data_loader",
+                "//executorch/exir/backend/test/demos/rpc:executor_backend",
+                "//executorch/exir/backend/test/demos/rpc:executor_backend_register",
+                "//executorch/extension/flat_tensor:flat_tensor_data_map",
+            ],
+            env = {
+                # The tests use these vars to find the program files to load.
+                # Uses an fbcode target path because the authoring/export tools
+                # intentionally don't work in xplat (since they're host-only
+                # tools).
+                "ET_MODULE_LINEAR_DELEGATE_PROGRAM_PATH": "$(location fbcode//executorch/test/models:exported_executor_backend_program_and_data[ModuleLinear-e.pte])",
+                "ET_MODULE_LINEAR_DATA_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleLinear.ptd])",
+            },
+        )
         runtime.cxx_test(
             name = "memory_manager_test",
             srcs = [

test/models/export_delegated_program.py

@@ -20,9 +20,13 @@
 from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower
 from executorch.exir.backend.backend_api import to_backend
 from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult
+from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.backend.test.backend_with_compiler_demo import (
     BackendWithCompilerDemo,
 )
+from executorch.exir.backend.test.demos.rpc.executor_backend_preprocess import (  # noqa: F401
+    ExecutorBackend,
+)
 from executorch.exir.passes.external_constants_pass import (
     delegate_external_constants_pass,
 )
@@ -150,13 +154,18 @@ def __init__(self, fn, method_name=method_name):
         def forward(self, *args, **kwargs):
             return getattr(self.fn, self.method_name)(*args, **kwargs)
 
-    exported_program = export(WrapperModule(eager_module), args=inputs, strict=True)
+    if method_name != "forward":
+        # Only require wrapper module if we're exporting a specific method other than forward.
+        exported_program = export(WrapperModule(eager_module), args=inputs, strict=True)
+    else:
+        exported_program = export(eager_module, args=inputs, strict=True)
 
     edge_config = EdgeCompileConfig(_check_ir_validity=False)
     et_config = exir.ExecutorchBackendConfig(
         extract_delegate_segments=extract_delegate_segments,
         constant_tensor_alignment=constant_tensor_alignment,
         delegate_alignment=delegate_alignment,
+        external_constants=external_constants,
     )
 
     if backend_id == "XnnpackBackend":
@@ -181,7 +190,10 @@ def forward(self, *args, **kwargs):
     else:
         edge: exir.EdgeProgramManager = to_edge(exported_program)
         lowered_module = to_backend(  # type: ignore[call-arg]
-            backend_id, edge.exported_program(), compile_specs=[]
+            backend_id,
+            edge.exported_program(),
+            # Just for the demo executor_backend.
+            compile_specs=[CompileSpec(key="external_constants", value=b"")],
         )
 
         class CompositeModule(nn.Module):

test/models/export_program.py

@@ -146,6 +146,19 @@ def get_random_inputs(self):
         return (torch.ones(2, 2, dtype=torch.float),)
 
 
+# Used for program-data-separation.
+class ModuleLinear(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = torch.nn.Linear(3, 3)
+
+    def forward(self, x: torch.Tensor):
+        return self.linear(x)
+
+    def get_random_inputs(self):
+        return (torch.randn(3),)
+
+
 class ModuleMultipleEntry(torch.nn.Module):
     def __init__(self):
         super().__init__()

test/models/targets.bzl

@@ -95,6 +95,7 @@ def define_common_targets():
     # Class names of nn.Modules for :exported_programs to export.
     MODULES_AND_DATA_TO_EXPORT = [
         "ModuleAddMul",
+        "ModuleLinear",
         "ModuleSimpleTrain",
     ]
 
@@ -104,6 +105,8 @@ def define_common_targets():
         outs = {
             "ModuleAddMul.pte": ["ModuleAddMulProgram.pte"],
             "ModuleAddMul.ptd": ["ModuleAddMulProgram.ptd"],
+            "ModuleLinear.pte": ["ModuleLinearProgram.pte"],
+            "ModuleLinear.ptd": ["ModuleLinearProgram.ptd"],
             "ModuleSimpleTrainProgram.pte": ["ModuleSimpleTrainProgram.pte"],
             "ModuleSimpleTrain.ptd": ["ModuleSimpleTrainProgram.ptd"],
         },
@@ -146,7 +149,7 @@ def define_common_targets():
         deps = [
             ":export_delegated_program_lib",
             "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
-
+            "//executorch/exir/backend/test/demos/rpc:executor_backend_preprocess",
         ],
         visibility = [],  # Private
     )
@@ -225,3 +228,23 @@ def define_common_targets():
             "//executorch/test/...",
         ],
     )
+
+    # Export with demo ExecutorBackend for program-data separation test.
+    runtime.genrule(
+        name = "exported_executor_backend_program_and_data",
+        cmd = "$(exe :export_delegated_program)" +
+            " --modules ModuleLinear" + 
+            " --backend_id ExecutorBackend" +
+            " --external_constants" +
+            " --outdir $OUT",
+
+        outs = {
+            "ModuleLinear-e.pte": ["ModuleLinear-e.pte"],
+        },
+        default_outs = ["."],
+        visibility = [
+            "//executorch/runtime/executor/test/...",
+            "//executorch/extension/flat_tensor/test/...",
+            "//executorch/test/...",
+        ],
+    )