Introduce APIs In Bundled Program to Not Take **Method** When Loading Input

[zhenyan-zhang-meta]

Context

This issue is the first step of #9638. In the discussion, we want to unblock having extension/module as the single source of implementation, which means that pybindings/pybindings should use extension/module rather than its own.

Specifically, the thing we need to unblock is that pybindings/pybindings needs to load bundled program, and the current BundledProgram APIs take method module as input (e.g. executorch::bundled_program::load_bundled_input), where Extension.module does not expose method class. Therefore we need to introduce new APIs to do the same job without exposing method, wrap as a form of BundledModule.

As a first step, we want to introduce a new API in BundledProgram to get bundled program input out of the program without taking method.

Proposal

The proposed change includes a utility function to the bundled program.

// in bundled_program.bundled_program.h

/**
  * Load testset_idx-th bundled input of method_name into given data_ptr 
  */
ET_NODISCARD ::executorch::runtime::Error get_bundled_input(
    SerializedBundledProgram* bundled_program_ptr,
    char* method_name,
    size_t testset_idx,
    void* data_ptr,
    size_t data_length);

With this API, we are able to unblocked implementation the bundled module wrapper for pybindings.

[zhenyan-zhang-meta]

We are closing this issue as we are about to take another path to not make changes to bundled_program but use the existing implementations to the largest extent. We will start from #10364.

We do get learnings from an implementation:

ET_NODISCARD Error get_bundled_input(
    SerializedBundledProgram* bundled_program_ptr,
    const char* method_name,
    size_t testset_idx,
    EValue* data_ptr,
    size_t& data_len) {
  ET_CHECK_OR_RETURN_ERROR(
      bundled_program_flatbuffer::BundledProgramBufferHasIdentifier(
          bundled_program_ptr),
      NotSupported,
      "The input buffer should be a bundled program.");

  auto method_test = get_method_test_suite(
      bundled_program_flatbuffer::GetBundledProgram(bundled_program_ptr),
      method_name);

  if (!method_test.ok()) {
    return method_test.error();
  }

  auto bundled_inputs =
      method_test.get()->test_cases()->Get(testset_idx)->inputs();
  EValue* current_data_ptr = data_ptr;
  data_len = 0;

  for (size_t input_idx = 0; input_idx < bundled_inputs->size(); input_idx++) {
    auto bundled_input = bundled_inputs->GetMutableObject(input_idx);
    EValue e_input;
    size_t e_size;

    switch (bundled_input->val_type()) {
      case bundled_program_flatbuffer::ValueUnion::Tensor: {
        auto bundled_input_tensor =
            static_cast<bundled_program_flatbuffer::Tensor*>(
                bundled_input->mutable_val());

#ifdef USE_ATEN_LIB
        Tensor t = tensor_like(bundled_input_tensor);
#else // !USE_ATEN_LIB
        TensorImpl impl = impl_like(bundled_input_tensor);
        Tensor t = Tensor(&impl);
#endif
        e_input = EValue(t);
        e_size = t.nbytes();
        break;
      }
      case bundled_program_flatbuffer::ValueUnion::Int: {
        auto bundled_input_int = bundled_input->val_as_Int();
        e_input = EValue(bundled_input_int->int_val());
        e_size = sizeof(int64_t);
        break;
      }
      case bundled_program_flatbuffer::ValueUnion::Double: {
        auto bundled_input_double = bundled_input->val_as_Double();
        e_input = EValue(bundled_input_double->double_val());
        e_size = sizeof(double);
        break;
      }
      case bundled_program_flatbuffer::ValueUnion::Bool: {
        auto bundled_input_bool = bundled_input->val_as_Bool();
        e_input = EValue(bundled_input_bool->bool_val());
        e_size = sizeof(bool);
        break;
      }
      default: {
        ET_CHECK_OR_RETURN_ERROR(
            false,
            NotSupported,
            "Data type %hhu not supported",
            static_cast<uint8_t>(bundled_input->val_type()));
      }
    }

    memcpy(current_data_ptr, &e_input, e_size);
    current_data_ptr = (EValue*)((char*)current_data_ptr + e_size);
    data_len += e_size;
  }

  return Error::Ok;
}