Enable selective build in OSS

CMakeLists.txt

@@ -554,6 +554,10 @@ if(EXECUTORCH_BUILD_PTHREADPOOL AND EXECUTORCH_BUILD_CPUINFO)
 endif()
 
 if(EXECUTORCH_BUILD_PYBIND)
+
+  # Add codegen tools subdirectory for selective_build pybind module
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/codegen/tools)
+
   if(NOT EXECUTORCH_BUILD_EXTENSION_DATA_LOADER)
     add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/data_loader)
   endif()

codegen/tools/CMakeLists.txt

@@ -0,0 +1,45 @@
+# 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.
+
+# Check if pybind11 is available
+
+# Create the selective_build pybind11 module
+pybind11_add_module(selective_build SHARED selective_build.cpp)
+
+# Set the output name to match the module name
+set_target_properties(selective_build PROPERTIES OUTPUT_NAME "selective_build")
+
+# Set the module name for the pybind11 module
+target_compile_definitions(
+  selective_build PUBLIC EXECUTORCH_PYTHON_MODULE_NAME=selective_build
+)
+
+# Include directories
+target_include_directories(
+  selective_build PRIVATE
+  ${CMAKE_CURRENT_SOURCE_DIR}/../../..
+)
+
+# Compile options
+target_compile_options(
+  selective_build PUBLIC
+  -Wno-deprecated-declarations
+  -fPIC
+  -frtti
+  -fexceptions
+)
+
+# Link against required libraries
+target_link_libraries(
+  selective_build PRIVATE
+  executorch_core
+  program_schema
+)
+
+# Install the module
+install(TARGETS selective_build
+        LIBRARY DESTINATION executorch/codegen/tools
+)

codegen/tools/gen_oplist.py

@@ -20,7 +20,6 @@
     # We can use relative import instead.
     from ..parse import strip_et_fields
 
-
 from torchgen.gen import LineLoader, parse_native_yaml_struct
 from torchgen.selective_build.operator import SelectiveBuildOperator
 from torchgen.selective_build.selector import merge_et_kernel_metadata
@@ -102,7 +101,6 @@ def _get_operators(model_file: str) -> List[str]:
 
 
 def _get_kernel_metadata_for_model(model_file: str) -> Dict[str, List[str]]:
-
     from executorch.codegen.tools.selective_build import (  # type: ignore[import-not-found]
         _get_io_metadata_for_program_operators,
         _get_program_from_buffer,

codegen/tools/selective_build.cpp

@@ -0,0 +1,269 @@
+/*
+ * 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 <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include <executorch/runtime/platform/assert.h>
+#include <executorch/schema/program_generated.h>
+
+namespace py = pybind11;
+
+namespace torch {
+namespace executor {
+
+namespace {
+
+// Metadata for kernel call io variables.
+// dtype and dim_order will exist only if corresponding variable is Tensor.
+struct IOMetaData {
+  int kernel_type;
+  int dtype;
+  std::vector<unsigned int> dim_order;
+
+  // Create tensor metadata. It records tensor's dtype and dim order.
+  explicit IOMetaData(const executorch_flatbuffer::Tensor* t)
+      : kernel_type(
+            static_cast<int>(executorch_flatbuffer::KernelTypes::Tensor)),
+        dtype(static_cast<int>(t->scalar_type())) {
+    for (size_t i = 0; i < t->dim_order()->size(); i++) {
+      dim_order.push_back(static_cast<unsigned int>(t->dim_order()->Get(i)));
+    }
+  }
+
+  // Create metadata for non-tensor variable.
+  explicit IOMetaData(executorch_flatbuffer::KernelTypes type)
+      : kernel_type(static_cast<int>(type)) {
+    ET_CHECK(
+        type != executorch_flatbuffer::KernelTypes::Tensor &&
+        type != executorch_flatbuffer::KernelTypes::TensorList &&
+        type != executorch_flatbuffer::KernelTypes::OptionalTensorList);
+  }
+};
+
+struct KernelIOMetaDataComparsion {
+  bool operator()(
+      const std::vector<IOMetaData>& lhs,
+      const std::vector<IOMetaData>& rhs) const {
+    if (lhs.size() != rhs.size()) {
+      return lhs.size() < rhs.size();
+    }
+    for (size_t i = 0; i < lhs.size(); i++) {
+      if (lhs[i].kernel_type != rhs[i].kernel_type) {
+        return lhs[i].kernel_type < rhs[i].kernel_type;
+      }
+      if (lhs[i].kernel_type !=
+          static_cast<int>(executorch_flatbuffer::KernelTypes::Tensor)) {
+        continue;
+      }
+      if (lhs[i].dtype != rhs[i].dtype) {
+        return lhs[i].dtype < rhs[i].dtype;
+      }
+      if (lhs[i].dim_order != rhs[i].dim_order) {
+        return lhs[i].dim_order < rhs[i].dim_order;
+      }
+    }
+    return false;
+  }
+};
+
+using KernelIOMetadata = std::vector<IOMetaData>;
+
+using OpIOMetaData = std::set<KernelIOMetadata, KernelIOMetaDataComparsion>;
+
+std::vector<std::string> get_operators_from_execution_plan(
+    const executorch_flatbuffer::ExecutionPlan& plan) {
+  std::vector<std::string> op_names;
+  for (const executorch_flatbuffer::Operator* op : *plan.operators()) {
+    if (op->overload()->str().empty()) {
+      op_names.push_back(op->name()->str());
+    } else {
+      op_names.push_back(op->name()->str() + "." + op->overload()->str());
+    }
+  }
+  return op_names;
+}
+
+std::map<std::string, OpIOMetaData>
+get_kernel_tensor_metadatas_from_execution_plan(
+    const executorch_flatbuffer::ExecutionPlan* plan) {
+  std::map<std::string, OpIOMetaData> op_io_metadata;
+  for (const executorch_flatbuffer::Chain* chain : *plan->chains()) {
+    for (const executorch_flatbuffer::Instruction* inst :
+         *chain->instructions()) {
+      if (inst->instr_args_type() ==
+          executorch_flatbuffer::InstructionArguments::KernelCall) {
+        const executorch_flatbuffer::KernelCall* kernel_call =
+            inst->instr_args_as_KernelCall();
+        const executorch_flatbuffer::Operator* op =
+            plan->operators()->Get(kernel_call->op_index());
+        std::string op_overload_name = op->name()->str();
+        if (op->overload()->size()) {
+          op_overload_name += "." + op->overload()->str();
+        }
+
+        // create an empty entry if current kernel is not in the map.
+        if (op_io_metadata.count(op_overload_name) == 0) {
+          op_io_metadata.insert(
+              std::make_pair(op_overload_name, OpIOMetaData()));
+        }
+
+        // go through IOs of this operator and collect tensor metadatas.
+        KernelIOMetadata kernel_io_metadata;
+        for (int arg_id : *kernel_call->args()) {
+          const executorch_flatbuffer::EValue* arg =
+              plan->values()->Get(arg_id);
+          if (arg->val_type() == executorch_flatbuffer::KernelTypes::Tensor) {
+            kernel_io_metadata.push_back(IOMetaData(arg->val_as_Tensor()));
+          } else if (
+              arg->val_type() ==
+              executorch_flatbuffer::KernelTypes::TensorList) {
+            if (arg->val_as_TensorList()->items()->size() == 0) {
+              // treat empty tensor list as null type since we can not get
+              // metadata from it.
+              kernel_io_metadata.push_back(
+                  IOMetaData(executorch_flatbuffer::KernelTypes::Null));
+            } else {
+              // all eles in TensorList are tensor and share same tensor
+              // metadata. use the metadata of first element as the metadata for
+              // whole list.
+              const executorch_flatbuffer::Tensor* tensor_arg =
+                  plan->values()
+                      ->Get(arg->val_as_TensorList()->items()->Get(0))
+                      ->val_as_Tensor();
+              kernel_io_metadata.push_back(IOMetaData(tensor_arg));
+            }
+          } else if (
+              arg->val_type() ==
+              executorch_flatbuffer::KernelTypes::OptionalTensorList) {
+            // all eles in OptionalTensorList are either tensor or null, and all
+            // tensors share same metadata. Use the metadata of first tensor
+            // element as the metadata for whole list. If no tensor exists (e.g.
+            // each element is None), treat the whole list as a single null
+            // element.
+            const executorch_flatbuffer::OptionalTensorList* opt_tensor_list =
+                arg->val_as_OptionalTensorList();
+
+            // Find one non-null tensor
+            bool found_tensor_element = false;
+            for (size_t i = 0; i < opt_tensor_list->items()->size(); i++) {
+              // We now adopt both index == -1 and actually serialize a null
+              // type EValue to represent a null data.
+              if (opt_tensor_list->items()->Get(i) != -1 &&
+                  plan->values()
+                          ->Get(opt_tensor_list->items()->Get(i))
+                          ->val_type() ==
+                      executorch_flatbuffer::KernelTypes::Tensor) {
+                const executorch_flatbuffer::Tensor* tensor_arg =
+                    plan->values()
+                        ->Get(opt_tensor_list->items()->Get(i))
+                        ->val_as_Tensor();
+                kernel_io_metadata.push_back(IOMetaData(tensor_arg));
+                found_tensor_element = true;
+                break;
+              }
+            }
+            if (!found_tensor_element) {
+              kernel_io_metadata.push_back(
+                  IOMetaData(executorch_flatbuffer::KernelTypes::Null));
+            }
+          } else {
+            kernel_io_metadata.push_back(IOMetaData(arg->val_type()));
+          }
+        }
+        op_io_metadata[op_overload_name].insert(kernel_io_metadata);
+      }
+    }
+  }
+  return op_io_metadata;
+}
+} // namespace
+
+const executorch_flatbuffer::Program* _get_program_from_buffer(
+    const py::bytes& buffer) {
+  return executorch_flatbuffer::GetProgram(
+      buffer.cast<std::string_view>().data());
+}
+
+py::list _get_program_operators(const executorch_flatbuffer::Program* program) {
+  const auto& plans = *program->execution_plan();
+  std::vector<std::string> op_names;
+  for (const auto& plan : plans) {
+    auto plan_ops = get_operators_from_execution_plan(*plan);
+    if (!plan_ops.empty()) {
+      op_names.insert(op_names.end(), plan_ops.begin(), plan_ops.end());
+    }
+  }
+  return py::cast(op_names);
+}
+
+// expose IO metadatas for all operators in given program
+py::dict _get_io_metadata_for_program_operators(
+    const executorch_flatbuffer::Program* program) {
+  const auto& plans = *program->execution_plan();
+  std::map<std::string, OpIOMetaData> program_op_io_metadata;
+
+  // aggregrate op metadata from different execution plan.
+  for (const executorch_flatbuffer::ExecutionPlan* plan : plans) {
+    std::map<std::string, OpIOMetaData> plan_op_io_metadata =
+        get_kernel_tensor_metadatas_from_execution_plan(plan);
+
+    for (const auto& op_io_metadata : plan_op_io_metadata) {
+      std::string op_name = op_io_metadata.first;
+      if (program_op_io_metadata.count(op_name) == 0) {
+        program_op_io_metadata.insert(std::make_pair(op_name, OpIOMetaData()));
+      }
+      program_op_io_metadata[op_name].insert(
+          plan_op_io_metadata[op_name].begin(),
+          plan_op_io_metadata[op_name].end());
+    }
+  }
+
+  // convert program_op_io_metadata to py data structure.
+  py::dict py_program_op_io_metadata;
+  for (const auto& op_io_meta : program_op_io_metadata) {
+    py::set py_op_io_meta;
+    for (const auto& io_metas : op_io_meta.second) {
+      py::list py_io_metadatas;
+      for (const auto& io_metadata : io_metas) {
+        py_io_metadatas.append(io_metadata);
+      }
+      py_op_io_meta.add(py::tuple(py_io_metadatas));
+    }
+    py_program_op_io_metadata[op_io_meta.first.data()] = py_op_io_meta;
+  }
+
+  return py_program_op_io_metadata;
+}
+
+PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
+  py::class_<executorch_flatbuffer::Program>(m, "_Program");
+
+  m.def(
+      "_get_program_from_buffer",
+      &_get_program_from_buffer,
+      py::return_value_policy::reference);
+
+  m.def(
+      "_get_program_operators",
+      &_get_program_operators,
+      py::return_value_policy::copy);
+
+  m.def(
+      "_get_io_metadata_for_program_operators",
+      &_get_io_metadata_for_program_operators,
+      py::return_value_policy::copy);
+
+  py::class_<IOMetaData>(m, "_IOMetaData")
+      .def_readwrite("kernel_type", &IOMetaData::kernel_type)
+      .def_readwrite("dtype", &IOMetaData::dtype)
+      .def_readwrite("dim_order", &IOMetaData::dim_order);
+}
+
+} // namespace executor
+} // namespace torch

codegen/tools/selective_build.pyi

@@ -0,0 +1,23 @@
+# 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.
+
+from typing import Any, Dict, List
+
+class _Program: ...
+
+class _IOMetaData:
+    @property
+    def kernel_type(self) -> int: ...
+    @property
+    def dtype(self) -> int: ...
+    @property
+    def dim_order(self) -> List[int]: ...
+
+def _get_program_from_buffer(buffer: bytes) -> _Program: ...
+def _get_program_operators(program: _Program) -> List[str]: ...
+def _get_io_metadata_for_program_operators(
+    program: _Program,
+) -> Dict[str, Any]: ...