Make _get_program_from_buffer work for bundled programs

Add some cmake to only do this if executorch is built with
bundleio.

codegen/tools subdirectory needs to be moved in top-level CmakeLists.txt
to have access to the bundled_program target.

Signed-off-by: Erik Lundell <[email protected]>
Change-Id: Ic953bdfd5223eb835e46150bde19b08b082e8444

Author: Erik-Lundell

CMakeLists.txt

@@ -738,9 +738,6 @@ 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()
@@ -749,6 +746,9 @@ if(EXECUTORCH_BUILD_PYBIND)
     add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/devtools)
   endif()
 
+  # Add codegen tools subdirectory for selective_build pybind module
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/codegen/tools)
+
   # Create bundled_module target only for pybindings when bundled_program exists
   # This target has hard dependencies on devtools generated headers
   if(TARGET bundled_program)

codegen/tools/CMakeLists.txt

@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2025 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -28,6 +29,10 @@ target_compile_options(
 )
 
 # Link against required libraries
+if(TARGET bundled_program)
+  target_compile_definitions(selective_build PRIVATE -DET_BUNDLE_IO)
+  target_link_libraries(selective_build PRIVATE bundled_program)
+endif()
 target_link_libraries(selective_build PRIVATE executorch_core program_schema)
 
 # Install the module

codegen/tools/selective_build.cpp

@@ -1,16 +1,21 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
+ * Copyright 2025 Arm Limited and/or its affiliates.
  *
  * This source code is licensed under the BSD-style license found in the
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <executorch/runtime/platform/assert.h>
+#include <executorch/schema/program_generated.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
-#include <executorch/runtime/platform/assert.h>
-#include <executorch/schema/program_generated.h>
+#ifdef ET_BUNDLE_IO
+#include <executorch/devtools/bundled_program/bundled_program.h>
+#include <stdexcept>
+#endif
 
 namespace py = pybind11;
 
@@ -186,8 +191,39 @@ get_kernel_tensor_metadatas_from_execution_plan(
 
 const executorch_flatbuffer::Program* _get_program_from_buffer(
     const py::bytes& buffer) {
+  // Access the Python bytes without copying and get raw pointer/size.
+  const std::string_view sv = buffer.cast<std::string_view>();
+  void* buf_ptr = const_cast<void*>(static_cast<const void*>(sv.data()));
+  const size_t buf_len = sv.size();
+#ifdef ET_BUNDLE_IO
+
+  // If this is a bundled program, extract the inner ExecuTorch program bytes.
+  if (executorch::bundled_program::is_bundled_program(buf_ptr, buf_len)) {
+    const void* program_data = nullptr;
+    size_t program_size = 0;
+
+    const auto status = executorch::bundled_program::get_program_data(
+        buf_ptr, // serialized BundledProgram start
+        buf_len, // total size of the BundledProgram blob
+        &program_data, // [out] pointer to inner .pte bytes
+        &program_size // [out] size of inner .pte bytes
+    );
+
+    if (status != ::executorch::runtime::Error::Ok || program_data == nullptr ||
+        program_size == 0) {
+      throw std::runtime_error(
+          "bundled_program::get_program_data() failed or returned empty data");
+    }
+
+    // program_data points directly at the flatbuffer-encoded Program region.
+    return executorch_flatbuffer::GetProgram(
+        reinterpret_cast<const uint8_t*>(program_data));
+  }
+#endif
+  // Otherwise treat the buffer as a raw .pte (flatbuffer Program with optional
+  // extended header).
   return executorch_flatbuffer::GetProgram(
-      buffer.cast<std::string_view>().data());
+      reinterpret_cast<const uint8_t*>(sv.data()));
 }
 
 py::list _get_program_operators(const executorch_flatbuffer::Program* program) {