Add a namespace for ATen mode

Differential Revision: D72440313

Pull Request resolved: #9894

Author: larryliu0820

codegen/templates/RegisterCodegenUnboxedKernels.cpp

@@ -22,8 +22,8 @@
 // JIT op registry instead of c10 dispatcher. JIT op registry only takes boxed
 // kernels, so we are calling unboxing functions in UnboxingFunctions.h to cast
 // arguments into C++ types (instead of IValue) and delegate to unboxed kernels.
-using KernelSpan =
-    ::executorch::runtime::Span<const ::executorch::runtime::Kernel>;
+using KernelSpan = ::executorch::runtime::Span<
+    const ::executorch::ET_RUNTIME_NAMESPACE::Kernel>;
 namespace torch {
 namespace executor {
 namespace function {

devtools/bundled_program/bundled_program.cpp

@@ -27,9 +27,9 @@ using executorch::aten::ArrayRef;
 using executorch::aten::Half;
 using executorch::aten::ScalarType;
 using executorch::aten::Tensor;
+using ::executorch::ET_RUNTIME_NAMESPACE::Method;
 using ::executorch::runtime::Error;
 using ::executorch::runtime::EValue;
-using ::executorch::runtime::Method;
 using ::executorch::runtime::Result;
 
 namespace executorch {
@@ -332,8 +332,9 @@ ET_NODISCARD Error load_bundled_input(
         static_cast<uint32_t>(status));
   }
 
-  ::executorch::runtime::internal::event_tracer_set_bundled_input_index(
-      method.get_event_tracer(), testset_idx);
+  ::executorch::ET_RUNTIME_NAMESPACE::internal::
+      event_tracer_set_bundled_input_index(
+          method.get_event_tracer(), testset_idx);
 
   return Error::Ok;
 }

devtools/bundled_program/bundled_program.h

@@ -18,7 +18,7 @@ namespace bundled_program {
  * An opaque pointer to a serialized bundled program.
  */
 using SerializedBundledProgram = const void;
-
+using ::executorch::ET_RUNTIME_NAMESPACE::Method;
 /**
  * Load testset_idx-th bundled input of method_idx-th Method test in
  * bundled_program_ptr to given Method.
@@ -31,7 +31,7 @@ using SerializedBundledProgram = const void;
  * execution.
  */
 ET_NODISCARD ::executorch::runtime::Error load_bundled_input(
-    ::executorch::runtime::Method& method,
+    Method& method,
     SerializedBundledProgram* bundled_program_ptr,
     size_t testset_idx);
 
@@ -49,7 +49,7 @@ ET_NODISCARD ::executorch::runtime::Error load_bundled_input(
  * execution.
  */
 ET_NODISCARD ::executorch::runtime::Error verify_method_outputs(
-    ::executorch::runtime::Method& method,
+    Method& method,
     SerializedBundledProgram* bundled_program_ptr,
     size_t testset_idx,
     double rtol = 1e-5,
@@ -106,7 +106,7 @@ using serialized_bundled_program =
     ::executorch::bundled_program::SerializedBundledProgram;
 
 ET_NODISCARD inline ::executorch::runtime::Error LoadBundledInput(
-    ::executorch::runtime::Method& method,
+    Method& method,
     serialized_bundled_program* bundled_program_ptr,
     size_t testset_idx) {
   return ::executorch::bundled_program::load_bundled_input(
@@ -115,7 +115,7 @@ ET_NODISCARD inline ::executorch::runtime::Error LoadBundledInput(
 
 ET_NODISCARD inline ::executorch::runtime::Error
 VerifyResultWithBundledExpectedOutput(
-    ::executorch::runtime::Method& method,
+    Method& method,
     serialized_bundled_program* bundled_program_ptr,
     size_t testset_idx,
     double rtol = 1e-5,

examples/portable/custom_ops/custom_ops_2_out.cpp

@@ -13,7 +13,7 @@ namespace native {
 
 using executorch::aten::ScalarType;
 using executorch::aten::Tensor;
-using executorch::runtime::KernelRuntimeContext;
+using executorch::ET_RUNTIME_NAMESPACE::KernelRuntimeContext;
 
 namespace {
 void check_preconditions(const Tensor& in, Tensor& out) {

extension/flat_tensor/flat_tensor_data_map.cpp

@@ -25,8 +25,8 @@ using executorch::runtime::Result;
 using executorch::runtime::Span;
 
 using executorch::aten::ScalarType;
+using executorch::ET_RUNTIME_NAMESPACE::TensorLayout;
 using executorch::runtime::DataLoader;
-using executorch::runtime::TensorLayout;
 
 namespace executorch {
 namespace extension {

extension/flat_tensor/flat_tensor_data_map.h

@@ -32,7 +32,8 @@ namespace extension {
 /**
  * A NamedDataMap implementation for FlatTensor-serialized data.
  */
-class FlatTensorDataMap final : public executorch::runtime::NamedDataMap {
+class FlatTensorDataMap final
+    : public executorch::ET_RUNTIME_NAMESPACE::NamedDataMap {
  public:
   /**
    * Creates a new DataMap that wraps FlatTensor data.
@@ -51,7 +52,8 @@ class FlatTensorDataMap final : public executorch::runtime::NamedDataMap {
    * @return Error::NotFound if the key is not present.
    */
   ET_NODISCARD
-  executorch::runtime::Result<const executorch::runtime::TensorLayout>
+  executorch::runtime::Result<
+      const executorch::ET_RUNTIME_NAMESPACE::TensorLayout>
   get_metadata(const char* key) const override;
 
   /**

extension/flat_tensor/targets.bzl

@@ -1,24 +1,26 @@
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 
 def define_common_targets():
-    runtime.cxx_library(
-        name = "flat_tensor_data_map",
-        srcs = [
-            "flat_tensor_data_map.cpp",
-        ],
-        exported_headers = ["flat_tensor_data_map.h"],
-        deps = [
-            "//executorch/runtime/core:core",
-            "//executorch/runtime/core:evalue",
-            "//executorch/runtime/core:named_data_map",
-            "//executorch/runtime/core/exec_aten:lib",
-            "//executorch/runtime/core/exec_aten/util:tensor_util",
-        ],
-        exported_deps = [
-            "//executorch/extension/flat_tensor/serialize:flat_tensor_header",
-            "//executorch/extension/flat_tensor/serialize:generated_headers",
-        ],
-        visibility = [
-            "//executorch/...",
-        ],
-    )
+    for aten_mode in [True, False]:
+        aten_suffix = "_aten" if aten_mode else ""
+        runtime.cxx_library(
+            name = "flat_tensor_data_map" + aten_suffix,
+            srcs = [
+                "flat_tensor_data_map.cpp",
+            ],
+            exported_headers = ["flat_tensor_data_map.h"],
+            deps = [
+                "//executorch/runtime/core:core",
+                "//executorch/runtime/core:evalue",
+                "//executorch/runtime/core:named_data_map" + aten_suffix,
+                "//executorch/runtime/core/exec_aten:lib" + aten_suffix,
+                "//executorch/runtime/core/exec_aten/util:tensor_util",
+            ],
+            exported_deps = [
+                "//executorch/extension/flat_tensor/serialize:flat_tensor_header",
+                "//executorch/extension/flat_tensor/serialize:generated_headers",
+            ],
+            visibility = [
+                "//executorch/...",
+            ],
+        )

extension/module/module.cpp

@@ -37,6 +37,9 @@
 namespace executorch {
 namespace extension {
 
+using ET_RUNTIME_NAMESPACE::MethodMeta;
+using ET_RUNTIME_NAMESPACE::Program;
+
 namespace {
 runtime::Result<std::unique_ptr<runtime::DataLoader>> load_file(
     const std::string& file_path,
@@ -113,7 +116,7 @@ Module::Module(
 }
 
 Module::Module(
-    std::shared_ptr<runtime::Program> program,
+    std::shared_ptr<Program> program,
     std::unique_ptr<runtime::MemoryAllocator> memory_allocator,
     std::unique_ptr<runtime::MemoryAllocator> temp_allocator,
     std::unique_ptr<runtime::EventTracer> event_tracer,
@@ -131,7 +134,7 @@ Module::Module(
   runtime::runtime_init();
 }
 
-runtime::Error Module::load(const runtime::Program::Verification verification) {
+runtime::Error Module::load(const Program::Verification verification) {
   if (!is_loaded()) {
     // Load the program
     if (!data_loader_) {
@@ -156,10 +159,10 @@ runtime::Error Module::load(const runtime::Program::Verification verification) {
     }
     // else: either the map itself was provided or we have no data map, either
     // way no work to do.
-    auto program = ET_UNWRAP_UNIQUE(
-        runtime::Program::load(data_loader_.get(), verification));
-    program_ = std::shared_ptr<runtime::Program>(
-        program.release(), [](runtime::Program* pointer) { delete pointer; });
+    auto program =
+        ET_UNWRAP_UNIQUE(Program::load(data_loader_.get(), verification));
+    program_ = std::shared_ptr<Program>(
+        program.release(), [](Program* pointer) { delete pointer; });
   }
   return runtime::Error::Ok;
 }
@@ -224,7 +227,7 @@ runtime::Error Module::load_method(
   return runtime::Error::Ok;
 }
 
-runtime::Result<runtime::MethodMeta> Module::method_meta(
+runtime::Result<MethodMeta> Module::method_meta(
     const std::string& method_name) {
   ET_CHECK_OK_OR_RETURN_ERROR(load_method(method_name));
   return methods_.at(method_name).method->method_meta();

extension/module/module.h

@@ -19,6 +19,11 @@
 namespace executorch {
 namespace extension {
 
+using ET_RUNTIME_NAMESPACE::Method;
+using ET_RUNTIME_NAMESPACE::MethodMeta;
+using ET_RUNTIME_NAMESPACE::NamedDataMap;
+using ET_RUNTIME_NAMESPACE::Program;
+
 /**
  * A facade class for loading programs and executing methods within them.
  */
@@ -95,7 +100,7 @@ class Module {
    * @param[in] data_map_loader A DataLoader used for loading external weights.
    */
   explicit Module(
-      std::shared_ptr<runtime::Program> program,
+      std::shared_ptr<Program> program,
       std::unique_ptr<runtime::MemoryAllocator> memory_allocator = nullptr,
       std::unique_ptr<runtime::MemoryAllocator> temp_allocator = nullptr,
       std::unique_ptr<runtime::EventTracer> event_tracer = nullptr,
@@ -116,8 +121,8 @@ class Module {
    */
   ET_NODISCARD
   runtime::Error load(
-      const runtime::Program::Verification verification =
-          runtime::Program::Verification::Minimal);
+      const Program::Verification verification =
+          Program::Verification::Minimal);
 
   /**
    * Checks if the program is loaded.
@@ -134,7 +139,7 @@ class Module {
    *
    * @returns Shared pointer to the program or nullptr if it's not yet loaded.
    */
-  inline std::shared_ptr<runtime::Program> program() const {
+  inline std::shared_ptr<Program> program() const {
     return program_;
   }
 
@@ -224,8 +229,7 @@ class Module {
    * @returns A method metadata, or an error if the program or method failed to
    * load.
    */
-  runtime::Result<runtime::MethodMeta> method_meta(
-      const std::string& method_name);
+  runtime::Result<MethodMeta> method_meta(const std::string& method_name);
 
   /**
    * Execute a specific method with the given input values and retrieve the
@@ -473,20 +477,20 @@ class Module {
     std::vector<runtime::Span<uint8_t>> planned_spans;
     std::unique_ptr<runtime::HierarchicalAllocator> planned_memory;
     std::unique_ptr<runtime::MemoryManager> memory_manager;
-    std::unique_ptr<runtime::Method> method;
+    std::unique_ptr<Method> method;
     std::vector<runtime::EValue> inputs;
   };
 
   std::string file_path_;
   std::string data_map_path_;
   LoadMode load_mode_{LoadMode::MmapUseMlock};
-  std::shared_ptr<runtime::Program> program_;
+  std::shared_ptr<Program> program_;
   std::unique_ptr<runtime::DataLoader> data_loader_;
   std::unique_ptr<runtime::MemoryAllocator> memory_allocator_;
   std::unique_ptr<runtime::MemoryAllocator> temp_allocator_;
   std::unique_ptr<runtime::EventTracer> event_tracer_;
   std::unique_ptr<runtime::DataLoader> data_map_loader_;
-  std::unique_ptr<runtime::NamedDataMap> data_map_;
+  std::unique_ptr<NamedDataMap> data_map_;
 
  protected:
   std::unordered_map<std::string, MethodHolder> methods_;

extension/module/targets.bzl

@@ -25,7 +25,7 @@ def define_common_targets():
                 "//executorch/extension/memory_allocator:malloc_memory_allocator",
                 "//executorch/extension/data_loader:file_data_loader",
                 "//executorch/extension/data_loader:mmap_data_loader",
-                "//executorch/extension/flat_tensor:flat_tensor_data_map",
+                "//executorch/extension/flat_tensor:flat_tensor_data_map" + aten_suffix,
             ],
             exported_deps = [
                 "//executorch/runtime/executor:program" + aten_suffix,