Add flat tensor to cmake

CMakeLists.txt

@@ -182,6 +182,10 @@ option(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER "Build the Data Loader extension"
        OFF
 )
 
+option(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR "Build the Flat Tensor extension"
+       OFF
+)
+
 option(EXECUTORCH_BUILD_EXTENSION_MODULE "Build the Module extension" OFF)
 
 option(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL "Build the Runner Util extension"
@@ -694,6 +698,11 @@ if(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/data_loader)
 endif()
 
+if(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/flat_tensor)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/flat_tensor/serialize)
+endif()
+
 if(EXECUTORCH_BUILD_EXTENSION_MODULE)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/module)
 endif()

build/cmake_deps.toml

@@ -148,6 +148,27 @@ deps = [
   "executorch",
 ]
 
+[targets.extension_flat_tensor_schema]
+buck_targets = [
+  "//extension/flat_tensor/serialize:generated_headers",
+]
+filters = [
+  ".fbs$",
+]
+
+[targets.extension_flat_tensor]
+buck_targets = [
+  "//extension/flat_tensor:flat_tensor_data_map",
+]
+filters = [
+  ".cpp$",
+]
+deps = [
+  "extension_flat_tensor_schema",
+  "executorch_core",
+  "executorch",
+]
+
 [targets.extension_module]
 buck_targets = [
   "//extension/module:module",

extension/flat_tensor/CMakeLists.txt

@@ -0,0 +1,39 @@
+# 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.
+
+# Please this file formatted by running:
+# ~~~
+# cmake-format -i CMakeLists.txt
+# ~~~
+
+cmake_minimum_required(VERSION 3.19)
+
+# Source root directory for executorch.
+if(NOT EXECUTORCH_ROOT)
+  set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..)
+endif()
+
+list(TRANSFORM _extension_flat_tensor__srcs PREPEND "${EXECUTORCH_ROOT}/")
+add_library(extension_flat_tensor ${_extension_flat_tensor__srcs})
+target_link_libraries(extension_flat_tensor executorch extension_data_loader)
+target_include_directories(extension_flat_tensor PUBLIC
+  ${EXECUTORCH_ROOT}/..
+  "${CMAKE_BINARY_DIR}/extension/flat_tensor/include"
+  "${EXECUTORCH_ROOT}/third-party/flatbuffers/include"
+  ${_common_include_directories})
+target_compile_options(extension_flat_tensor PUBLIC ${_common_compile_options})
+
+# Install libraries
+install(
+  TARGETS extension_flat_tensor
+  DESTINATION lib
+  INCLUDES
+  DESTINATION ${_common_include_directories}
+)
+
+if(BUILD_TESTING)
+  add_subdirectory(test)
+endif()

extension/flat_tensor/flat_tensor_data_map.cpp

@@ -8,9 +8,10 @@
 
 #include <executorch/extension/flat_tensor/flat_tensor_data_map.h>
 
+#include <executorch/extension/flat_tensor/serialize/flat_tensor_generated.h>
 #include <executorch/extension/flat_tensor/serialize/flat_tensor_header.h>
-#include <executorch/extension/flat_tensor/serialize/schema_generated.h>
 
+#include <executorch/runtime/core/data_loader.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
 #include <executorch/runtime/core/freeable_buffer.h>

extension/flat_tensor/serialize/CMakeLists.txt

@@ -0,0 +1,69 @@
+# 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.
+
+# Flatbuffer schema header lib. Please this file formatted by running:
+# ~~~
+# cmake-format -i CMakeLists.txt
+# ~~~
+
+if(NOT FLATC_EXECUTABLE)
+  set(FLATC_EXECUTABLE flatc)
+endif()
+
+# The include directory that will contain the generated schema headers.
+set(_flat_tensor_schema__include_dir "${CMAKE_BINARY_DIR}/extension/flat_tensor/include")
+set(_flat_tensor_schema__output_dir "${_flat_tensor_schema__include_dir}/executorch/extension/flat_tensor/serialize")
+# Source root directory for executorch.
+if(NOT EXECUTORCH_ROOT)
+  set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/..)
+endif()
+
+function(generate_flat_tensor_schema _schema_srcs _schema_name)
+  set(_schema_outputs)
+  foreach(fbs_file ${_schema_srcs})
+    string(REGEX REPLACE "[.]fbs$" "_generated.h" generated "${fbs_file}")
+    list(APPEND _schema_outputs
+         "${_flat_tensor_schema__output_dir}/${generated}"
+    )
+  endforeach()
+
+  # Generate the headers from the .fbs files.
+  add_custom_command(
+    OUTPUT ${_schema_outputs}
+    COMMAND
+      ${FLATC_EXECUTABLE} --cpp --cpp-std c++11 --gen-mutable --scoped-enums -o
+      "${_flat_tensor_schema__output_dir}" ${_schema_srcs}
+    WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+    DEPENDS ${FLATC_EXECUTABLE} ${_schema_srcs}
+    COMMENT "Generating ${_schema_name} headers"
+    VERBATIM
+  )
+
+  add_library(${_schema_name} INTERFACE ${_schema_outputs})
+  set_target_properties(${_schema_name} PROPERTIES LINKER_LANGUAGE CXX)
+
+  # exir lets users set the alignment of tensor data embedded in the flatbuffer,
+  # and some users need an alignment larger than the default, which is typically
+  # 32.
+  target_compile_definitions(
+    ${_schema_name} INTERFACE FLATBUFFERS_MAX_ALIGNMENT=1024
+  )
+
+  target_include_directories(
+    ${_schema_name}
+    INTERFACE ${_flat_tensor_schema__include_dir}
+              ${EXECUTORCH_ROOT}/third-party/flatbuffers/include
+  )
+endfunction()
+
+# Generate common schema
+set(scalar_type_schema_srcs scalar_type.fbs)
+generate_flat_tensor_schema("${scalar_type_schema_srcs}" "scalar_type_schema")
+
+# For the other schemas
+set(flat_tensor_schema_srcs flat_tensor.fbs)
+generate_flat_tensor_schema("${flat_tensor_schema_srcs}" "flat_tensor_schema")
+add_dependencies(flat_tensor_schema scalar_type_schema)

extension/flat_tensor/serialize/serialize.cpp

@@ -8,9 +8,9 @@
 
 #include <executorch/extension/flat_tensor/serialize/serialize.h>
 
+#include <executorch/extension/flat_tensor/serialize/flat_tensor_generated.h>
 #include <executorch/extension/flat_tensor/serialize/flat_tensor_header.h>
 #include <executorch/extension/flat_tensor/serialize/scalar_type_generated.h>
-#include <executorch/extension/flat_tensor/serialize/schema_generated.h>
 
 #include <fstream>
 #include <string>

extension/flat_tensor/serialize/targets.bzl

@@ -8,7 +8,7 @@ def define_common_targets():
             "scalar_type.fbs",
         ],
         outs = {
-            "schema_generated.h": ["flat_tensor_generated.h"],
+            "flat_tensor_generated.h": ["flat_tensor_generated.h"],
             "scalar_type_generated.h": ["scalar_type_generated.h"]
         },
         cmd = " ".join([
@@ -29,7 +29,7 @@ def define_common_targets():
             "//executorch/...",
         ],
         exported_headers = {
-            "schema_generated.h": ":gen_schema[schema_generated.h]",
+            "flat_tensor_generated.h": ":gen_schema[flat_tensor_generated.h]",
             "scalar_type_generated.h": ":gen_schema[scalar_type_generated.h]",
         },
         exported_external_deps = ["flatbuffers-api"],

extension/flat_tensor/test/CMakeLists.txt

@@ -0,0 +1,51 @@
+# 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.
+
+# @generated by test/utils/generate_gtest_cmakelists.py
+#
+# This file should be formatted with
+# ~~~
+# cmake-format -i CMakeLists.txt
+# ~~~
+# It should also be cmake-lint clean.
+#
+
+cmake_minimum_required(VERSION 3.19)
+
+set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..)
+
+include(${EXECUTORCH_ROOT}/build/Test.cmake)
+
+add_custom_command(
+  COMMENT "Lucy: Generating from extension/flat_tensor ${CMAKE_CURRENT_BINARY_DIR}"
+  OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.pte"
+         "${CMAKE_CURRENT_BINARY_DIR}/_default_external_constant.ptd"
+
+  COMMAND
+    python -m test.models.export_program --modules "ModuleLinear" --external-constants --outdir "${CMAKE_CURRENT_BINARY_DIR}" 2> /dev/null
+
+  WORKING_DIRECTORY ${EXECUTORCH_ROOT}
+)
+
+add_custom_target(
+  program_data_files
+  DEPENDS "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.pte"
+          "${CMAKE_CURRENT_BINARY_DIR}/_default_external_constant.ptd"
+)
+
+set(test_env_
+    "ET_MODULE_LINEAR_PROGRAM_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.pte"
+    "ET_MODULE_LINEAR_DATA_PATH=${CMAKE_CURRENT_BINARY_DIR}/_default_external_constant.ptd"
+)
+
+# set(ENV{RESOURCES_PATH} "${CMAKE_BINARY_DIR}")
+
+set(_test_srcs flat_tensor_data_map_test.cpp flat_tensor_header_test.cpp)
+
+et_cxx_test(extension_flat_tensor_test SOURCES ${_test_srcs} EXTRA_LIBS extension_flat_tensor extension_data_loader)
+
+add_dependencies(extension_flat_tensor_test program_data_files)
+set_property(TEST extension_flat_tensor_test PROPERTY ENVIRONMENT ${test_env_})

extension/flat_tensor/test/flat_tensor_data_map_test.cpp

@@ -8,8 +8,8 @@
 
 #include <executorch/extension/data_loader/file_data_loader.h>
 #include <executorch/extension/flat_tensor/flat_tensor_data_map.h>
+#include <executorch/extension/flat_tensor/serialize/flat_tensor_generated.h>
 #include <executorch/extension/flat_tensor/serialize/flat_tensor_header.h>
-#include <executorch/extension/flat_tensor/serialize/schema_generated.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/result.h>
 #include <executorch/runtime/platform/runtime.h>
@@ -35,7 +35,8 @@ class FlatTensorDataMapTest : public ::testing::Test {
 
     // Load data map. The eager linear model is defined at:
     // //executorch/test/models/linear_model.py
-    const char* path = std::getenv("ET_MODULE_LINEAR_DATA");
+    const char* path = std::getenv("ET_MODULE_LINEAR_DATA_PATH");
+    ET_LOG(Info, "Lucy: %s", path);
     Result<FileDataLoader> loader = FileDataLoader::from(path);
     ASSERT_EQ(loader.error(), Error::Ok);
 
@@ -46,6 +47,8 @@ class FlatTensorDataMapTest : public ::testing::Test {
 };
 
 TEST_F(FlatTensorDataMapTest, LoadFlatTensorDataMap) {
+  const char* path = std::getenv("ET_MODULE_LINEAR_DATA_PATH");
+  ET_LOG(Info, "Lucy: %s", path);
   Result<FlatTensorDataMap> data_map =
       FlatTensorDataMap::load(data_map_loader_.get());
   EXPECT_EQ(data_map.error(), Error::Ok);

extension/flat_tensor/test/targets.bzl

@@ -35,8 +35,8 @@ def define_common_targets(is_fbcode=False):
             # The tests use this var to find the program file to load. This 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_PROGRAM": "$(location fbcode//executorch/test/models:exported_programs_with_data_separated[ModuleLinear.pte])",
-            "ET_MODULE_LINEAR_DATA": "$(location fbcode//executorch/test/models:exported_programs_with_data_separated[ModuleLinear.ptd])",
+            "ET_MODULE_LINEAR_PROGRAM_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleLinear.pte])",
+            "ET_MODULE_LINEAR_DATA_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleLinear.ptd])",
         }
 
         runtime.cxx_test(

extension/flat_tensor/test/test_serialize.cpp

@@ -8,9 +8,9 @@
 
 #include <executorch/extension/flat_tensor/serialize/serialize.h>
 
+#include <executorch/extension/flat_tensor/serialize/flat_tensor_generated.h>
 #include <executorch/extension/flat_tensor/serialize/flat_tensor_header.h>
 #include <executorch/extension/flat_tensor/serialize/scalar_type_generated.h>
-#include <executorch/extension/flat_tensor/serialize/schema_generated.h>
 
 #include <executorch/extension/tensor/tensor_ptr.h>
 #include <executorch/runtime/core/result.h>

runtime/executor/method.cpp

@@ -16,6 +16,7 @@
 #include <executorch/runtime/backend/interface.h>
 #include <executorch/runtime/core/event_tracer_hooks.h>
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
+#include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/core/span.h>
 #include <executorch/runtime/executor/memory_manager.h>
 #include <executorch/runtime/executor/platform_memory_allocator.h>
@@ -288,7 +289,7 @@ Result<bool> parse_cond_value(const EValue& cond_value) {
 
 } // namespace
 
-Error Method::parse_values() {
+Error Method::parse_values(const NamedDataMap* named_data_map) {
   auto flatbuffer_values = serialization_plan_->values();
   ET_CHECK_OR_RETURN_ERROR(
       flatbuffer_values != nullptr, InvalidProgram, "Missing values");
@@ -414,7 +415,8 @@ Error Method::parse_values() {
         auto t = deserialization::parseTensor(
             program_,
             memory_manager_,
-            static_cast<const executorch_flatbuffer::Tensor*>(val));
+            static_cast<const executorch_flatbuffer::Tensor*>(val),
+            named_data_map);
         if (!t.ok()) {
           ET_LOG(
               Error,
@@ -607,7 +609,8 @@ Result<Method> Method::load(
     executorch_flatbuffer::ExecutionPlan* s_plan,
     const Program* program,
     MemoryManager* memory_manager,
-    EventTracer* event_tracer) {
+    EventTracer* event_tracer,
+    const NamedDataMap* named_data_map) {
   MemoryAllocator* temp_allocator = memory_manager->temp_allocator();
   if (temp_allocator == nullptr) {
     PlatformMemoryAllocator* platform_allocator =
@@ -621,7 +624,7 @@ Result<Method> Method::load(
   }
   Method method(program, memory_manager, event_tracer, temp_allocator);
 
-  Error err = method.init(s_plan);
+  Error err = method.init(s_plan, named_data_map);
   if (err != Error::Ok) {
     return err;
   } else {
@@ -630,7 +633,9 @@ Result<Method> Method::load(
   }
 }
 
-Error Method::init(executorch_flatbuffer::ExecutionPlan* s_plan) {
+Error Method::init(
+    executorch_flatbuffer::ExecutionPlan* s_plan,
+    const NamedDataMap* named_data_map) {
   EXECUTORCH_SCOPE_PROF("Method::init");
   internal::EventTracerProfileMethodScope event_tracer_profile_scope =
       internal::EventTracerProfileMethodScope(event_tracer_, "Method::init");
@@ -647,7 +652,7 @@ Error Method::init(executorch_flatbuffer::ExecutionPlan* s_plan) {
 
   {
     // Parse the elements of the values_ array.
-    Error err = parse_values();
+    Error err = parse_values(named_data_map);
     if (err != Error::Ok) {
       return err;
     }