Add cmake support for flat tensor

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"
@@ -240,6 +244,9 @@ cmake_dependent_option(
   "NOT EXECUTORCH_BUILD_ARM_BAREMETAL" OFF
 )
 
+if(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR)
+  set(EXECUTORCH_BUILF_EXTENSION_DATA_LOADER ON)
+endif()
 
 if(EXECUTORCH_BUILD_EXTENSION_TRAINING)
   set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON)
@@ -694,6 +701,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/Utils.cmake

@@ -67,6 +67,9 @@ function(executorch_print_configuration_summary)
   message(STATUS "  EXECUTORCH_BUILD_EXTENSION_DATA_LOADER : "
                  "${EXECUTORCH_BUILD_EXTENSION_DATA_LOADER}"
   )
+  message(STATUS "  EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR : "
+                 "${EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR}"
+  )
   message(STATUS "  EXECUTORCH_BUILD_EXTENSION_MODULE      : "
                  "${EXECUTORCH_BUILD_EXTENSION_MODULE}"
   )

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,48 @@
+# 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(
+  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(_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,7 @@ 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");
     Result<FileDataLoader> loader = FileDataLoader::from(path);
     ASSERT_EQ(loader.error(), Error::Ok);
 

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/test/CMakeLists.txt

@@ -26,11 +26,11 @@ add_custom_command(
          "${CMAKE_BINARY_DIR}/ModuleMultipleEntry.pte"
          "${CMAKE_BINARY_DIR}/ModuleSimpleTrain.pte"
   COMMAND
-    python3 -m test.models.export_program --modules
+    python -m test.models.export_program --modules
     "ModuleAdd,ModuleAddHalf,ModuleDynamicCatUnallocatedIO,ModuleIndex,ModuleLinear,ModuleMultipleEntry,ModuleSimpleTrain"
     --outdir "${CMAKE_BINARY_DIR}" 2> /dev/null
   COMMAND
-    python3 -m test.models.export_delegated_program --modules "ModuleAddMul"
+    python -m test.models.export_delegated_program --modules "ModuleAddMul"
     --backend_id "StubBackend" --outdir "${CMAKE_BINARY_DIR}" || true
   WORKING_DIRECTORY ${EXECUTORCH_ROOT}
 )

test/models/export_program.py

@@ -264,6 +264,8 @@ def main() -> None:
             # Skip type promotion to keep the model in fp16.
             # Type promotion will convert to fp32.
             skip_type_promotion = True
+        if args.external_constants:
+            module_name = f"{module_name}Program"
         outfile = os.path.join(args.outdir, f"{module_name}.pte")
         prog = export_module_to_program(
             module_class,

test/run_oss_cpp_tests.sh

@@ -33,6 +33,7 @@ build_executorch() {
     -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \
     -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \
     -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
+    -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \
     -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
     -DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=ON \
     -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \

test/utils/OSSTestConfig.json

@@ -18,6 +18,13 @@
             "../print_evalue.cpp"
         ]
     },
+    {
+        "directory": "extension/flat_tensor/test",
+        "sources": [
+            "flat_tensor_data_map_test.cpp",
+            "flat_tensor_header_test.cpp"
+        ]
+    },
     {
         "directory": "extension/kernel_util/test",
         "sources": [