add flat tensor to cmake

pytorch · Feb 10, 2025 · b4cd0b9 · b4cd0b9
1 parent 974373e
commit b4cd0b9
Show file tree

Hide file tree

Showing 14 changed files with 192 additions and 9 deletions.
diff --git a/CMakeLists.txt b/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()

diff --git a/build/cmake_deps.toml b/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",

diff --git a/extension/flat_tensor/CMakeLists.txt b/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()
diff --git a/extension/flat_tensor/flat_tensor_data_map.cpp b/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>

diff --git a/extension/flat_tensor/serialize/CMakeLists.txt b/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)
diff --git a/extension/flat_tensor/serialize/serialize.cpp b/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>

diff --git a/extension/flat_tensor/serialize/targets.bzl b/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"],

diff --git a/extension/flat_tensor/test/CMakeLists.txt b/extension/flat_tensor/test/CMakeLists.txt
@@ -0,0 +1,24 @@
+# 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)
+
+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 )
diff --git a/extension/flat_tensor/test/flat_tensor_data_map_test.cpp b/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>

diff --git a/extension/flat_tensor/test/test_serialize.cpp b/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>

diff --git a/runtime/executor/test/CMakeLists.txt b/runtime/executor/test/CMakeLists.txt
@@ -25,12 +25,18 @@ add_custom_command(
          "${CMAKE_BINARY_DIR}/ModuleLinear.pte"
          "${CMAKE_BINARY_DIR}/ModuleMultipleEntry.pte"
          "${CMAKE_BINARY_DIR}/ModuleSimpleTrain.pte"
+         "${CMAKE_BINARY_DIR}/ModuleLinearProgram.pte"
+         "${CMAKE_BINARY_DIR}/_default_external_constant.ptd"
   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_program --modules
+    "ModuleLinear" --external-constants
+    --outdir "${CMAKE_BINARY_DIR}" 2> /dev/null
+  COMMAND
+    python -m test.models.export_delegated_program --modules "ModuleAddMul"
     --backend_id "StubBackend" --outdir "${CMAKE_BINARY_DIR}" || true
   WORKING_DIRECTORY ${EXECUTORCH_ROOT}
 )
@@ -44,6 +50,8 @@ add_custom_target(
           "${CMAKE_BINARY_DIR}/ModuleLinear.pte"
           "${CMAKE_BINARY_DIR}/ModuleMultipleEntry.pte"
           "${CMAKE_BINARY_DIR}/ModuleSimpleTrain.pte"
+          "${CMAKE_BINARY_DIR}/ModuleLinearProgram.pte"
+          "${CMAKE_BINARY_DIR}/_default_external_constant.ptd"
 )
 
 set(test_env
@@ -53,6 +61,8 @@ set(test_env
     "ET_MODULE_DYNAMIC_CAT_UNALLOCATED_IO_PATH=${CMAKE_BINARY_DIR}/ModuleDynamicCatUnallocatedIO.pte"
     "ET_MODULE_INDEX_PATH=${CMAKE_BINARY_DIR}/ModuleIndex.pte"
     "ET_MODULE_LINEAR_PATH=${CMAKE_BINARY_DIR}/ModuleLinear.pte"
+    "ET_MODULE_LINEAR_PROGRAM_PATH=${CMAKE_BINARY_DIR}/ModuleLinearProgram.pte"
+    "ET_MODULE_LINEAR_DATA_PATH=${CMAKE_BINARY_DIR}/_default_external_constant.ptd"
     "ET_MODULE_MULTI_ENTRY_PATH=${CMAKE_BINARY_DIR}/ModuleMultipleEntry.pte"
     "ET_MODULE_SIMPLE_TRAIN_PATH=${CMAKE_BINARY_DIR}/ModuleSimpleTrain.pte"
 )
@@ -87,6 +97,7 @@ et_cxx_test(
   portable_ops_lib
   portable_kernels
   extension_data_loader
+  extension_flat_tensor
   extension_runner_util
 )
 add_dependencies(method_test generated_pte_files)

diff --git a/test/models/export_program.py b/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,

diff --git a/test/run_oss_cpp_tests.sh b/test/run_oss_cpp_tests.sh
@@ -33,11 +33,11 @@ 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 \
     -DEXECUTORCH_BUILD_DEVTOOLS=ON \
-    -DEXECUTORCH_BUILD_VULKAN=$BUILD_VULKAN \
     -DEXECUTORCH_BUILD_XNNPACK=ON \
     -DEXECUTORCH_BUILD_TESTS=ON \
     -Bcmake-out

diff --git a/test/utils/OSSTestConfig.json b/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": [