Update on "[executorch][runtime] Add NamedDataMap to method load"

Add NamedDataMap as an arg to:
- Method
- load_method
- parseTensor

Use NamedDataMap to resolve external tensors in parseTensor.

Test that the PTE + PTD file run well inside method_test.

Differential Revision: [D67127327](https://our.internmc.facebook.com/intern/diff/D67127327/)

[ghstack-poisoned]

Author: lucylq

.buckconfig

@@ -33,3 +33,6 @@
       **/.git, \
       cmake-out, \
       pip-out
+
+[buck2]
+restarter=true

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-0a94bb432ed75cc2d950d81b2921363218a7e459
+27e35de6c288bffad1b4d18b393579c1d1a95547

.ci/docker/conda-env-ci.txt

@@ -1,4 +1,5 @@
 cmake=3.22.1
 ninja=1.10.2
 libuv
+llvm-openmp
 pkg-config

.ci/scripts/gather_benchmark_configs.py

@@ -238,7 +238,7 @@ def set_output(name: str, val: Any) -> None:
     try:
         with open(github_output, "a") as env:
             env.write(f"{name}={val}\n")
-    except PermissionError:
+    except (PermissionError, FileNotFoundError):
         # Fall back to printing in case of permission error in unit tests
         print(f"::set-output name={name}::{val}")
 

.ci/scripts/setup-linux.sh

@@ -23,4 +23,7 @@ fi
 # of nightly. This allows CI to test against latest commits from PyTorch
 install_executorch "use-pt-pinned-commit"
 build_executorch_runner "${BUILD_TOOL}"
-do_not_use_nightly_on_ci
+
+if [[ "${GITHUB_BASE_REF:-}" == *main* || "${GITHUB_BASE_REF:-}" == *gh* ]]; then
+  do_not_use_nightly_on_ci
+fi

.ci/scripts/setup-macos.sh

@@ -121,6 +121,7 @@ setup_macos_env_variables
 # NB: we need buck2 in all cases because cmake build also depends on calling
 # buck2 atm
 install_buck
+brew install libomp
 install_pip_dependencies
 
 # TODO(huydhn): Unlike our self-hosted runner, GitHub runner doesn't have access
@@ -136,4 +137,7 @@ install_pytorch_and_domains
 # the pinned commit from PyTorch
 install_executorch "use-pt-pinned-commit"
 build_executorch_runner "${BUILD_TOOL}"
-do_not_use_nightly_on_ci
+
+if [[ "${GITHUB_BASE_REF:-}" == *main* || "${GITHUB_BASE_REF:-}" == *gh* ]]; then
+  do_not_use_nightly_on_ci
+fi

.github/workflows/pull.yml

@@ -212,17 +212,14 @@ jobs:
       docker-image: executorch-ubuntu-22.04-clang12
       submodules: 'true'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-      timeout: 90
+      timeout: 180
       script: |
         # The generic Linux job chooses to use base env, not the one setup by the image
         CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
         conda activate "${CONDA_ENV}"
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "cmake"
 
-        # install pybind
-        bash install_executorch.sh --pybind xnnpack --use-pt-pinned-commit
-
         # install Llava requirements
         bash examples/models/llama/install_requirements.sh
         bash examples/models/llava/install_requirements.sh
@@ -483,9 +480,6 @@ jobs:
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "cmake"
 
-        # install pybind
-        bash install_executorch.sh --pybind xnnpack --use-pt-pinned-commit
-
         # install phi-3-mini requirements
         bash examples/models/phi-3-mini/install_requirements.sh
 
@@ -513,9 +507,6 @@ jobs:
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "cmake"
 
-        # install pybind
-        bash install_executorch.sh --pybind xnnpack --use-pt-pinned-commit
-
         # install llama requirements
         bash examples/models/llama/install_requirements.sh
 
@@ -535,17 +526,14 @@ jobs:
       docker-image: executorch-ubuntu-22.04-clang12
       submodules: 'true'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-      timeout: 90
+      timeout: 180
       script: |
         # The generic Linux job chooses to use base env, not the one setup by the image
         CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
         conda activate "${CONDA_ENV}"
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "cmake"
 
-        # install pybind
-        bash install_executorch.sh --pybind xnnpack --use-pt-pinned-commit
-
         # install llama requirements
         bash examples/models/llama/install_requirements.sh
 
@@ -573,9 +561,6 @@ jobs:
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "cmake"
 
-        # install pybind
-        bash install_executorch.sh --pybind xnnpack --use-pt-pinned-commit
-
         # install llama requirements
         bash examples/models/llama/install_requirements.sh
 

.gitmodules

@@ -64,3 +64,9 @@
 [submodule "third-party/ao"]
 	path = third-party/ao
 	url = https://github.com/pytorch/ao.git
+[submodule "backends/cadence/utils/FACTO"]
+	path = backends/cadence/utils/FACTO
+	url = https://github.com/pytorch-labs/FACTO.git
+[submodule "third-party/pocketfft"]
+	path = third-party/pocketfft
+	url = https://github.com/mreineck/pocketfft

.mypy.ini

@@ -77,6 +77,9 @@ ignore_missing_imports = True
 [mypy-ruamel]
 ignore_missing_imports = True
 
+[mypy-serializer.*]
+ignore_missing_imports = True
+
 [mypy-setuptools.*]
 ignore_missing_imports = True
 

CMakeLists.txt

@@ -240,6 +240,13 @@ cmake_dependent_option(
   "NOT EXECUTORCH_BUILD_ARM_BAREMETAL" OFF
 )
 
+
+if(EXECUTORCH_BUILD_EXTENSION_TRAINING)
+  set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON)
+  set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER ON)
+  set(EXECUTORCH_BUILD_EXTENSION_MODULE ON)
+endif()
+
 if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
   set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON)
   set(EXECUTORCH_BUILD_KERNELS_CUSTOM ON)
@@ -791,6 +798,35 @@ if(EXECUTORCH_BUILD_PYBIND)
   install(TARGETS portable_lib
           LIBRARY DESTINATION executorch/extension/pybindings
   )
+
+  if(EXECUTORCH_BUILD_EXTENSION_TRAINING)
+
+    set(_pybind_training_dep_libs
+        ${TORCH_PYTHON_LIBRARY}
+        etdump
+        executorch
+        util
+        torch
+        extension_training
+    )
+
+    if(EXECUTORCH_BUILD_XNNPACK)
+      # need to explicitly specify XNNPACK and microkernels-prod
+      # here otherwise uses XNNPACK and microkernel-prod symbols from libtorch_cpu
+      list(APPEND _pybind_training_dep_libs xnnpack_backend XNNPACK microkernels-prod)
+    endif()
+
+    # pybind training
+    pybind11_add_module(_training_lib SHARED extension/training/pybindings/_training_lib.cpp)
+
+    target_include_directories(_training_lib PRIVATE ${TORCH_INCLUDE_DIRS})
+    target_compile_options(_training_lib PUBLIC ${_pybind_compile_options})
+    target_link_libraries(_training_lib PRIVATE ${_pybind_training_dep_libs})
+
+    install(TARGETS _training_lib
+            LIBRARY DESTINATION executorch/extension/training/pybindings
+    )
+  endif()
 endif()
 
 if(EXECUTORCH_BUILD_KERNELS_CUSTOM)