Update

[ghstack-poisoned]

Author: swolchok

.github/workflows/pull.yml

@@ -762,3 +762,66 @@ jobs:
 
         # Test selective build
         PYTHON_EXECUTABLE=python bash examples/wasm/test_build_wasm.sh
+
+  unittest-nxp-neutron:
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    permissions:
+      id-token: write
+      contents: read
+    with:
+      runner: linux.2xlarge
+      docker-image: executorch-ubuntu-22.04-clang12
+      submodules: 'recursive'
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      timeout: 90
+      script: |
+        set -eux
+
+        # 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}"
+
+        # Build and install Executorch
+        PYTHON_EXECUTABLE=python \
+        CMAKE_ARGS="-DEXECUTORCH_BUILD_NXP_NEUTRON=ON" \
+        .ci/scripts/setup-linux.sh --build-tool "cmake"
+
+        # Install test requirements
+        pip install -r backends/nxp/requirements-tests.txt
+
+        # Run pytest
+        PYTHON_EXECUTABLE=python bash backends/nxp/run_unittests.sh
+                
+        # Run aot example: 
+        PYTHON_EXECUTABLE=python bash examples/nxp/run_aot_example.sh
+
+
+  nxp-build-test:
+    name: nxp-build-test
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    permissions:
+      id-token: write
+      contents: read
+    with:
+      runner: linux.2xlarge
+      docker-image: executorch-ubuntu-22.04-arm-sdk
+      submodules: 'recursive'
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      timeout: 90
+      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}"
+
+        # Build
+        cmake -DEXECUTORCH_BUILD_NXP_NEUTRON=ON -Bcmake-out .
+        cmake --build cmake-out --target executorch_delegate_neutron --config Release
+
+        # Build check for the neutron backend library
+        lib_neutron="cmake-out/backends/nxp/libexecutorch_delegate_neutron.a"
+        if [ -f $lib_neutron ]; then
+            echo "Neutron backend library built."
+        else
+            echo "Neutron backend library not found!"
+            exit 1
+        fi

.github/workflows/trunk.yml

@@ -302,36 +302,6 @@ jobs:
           exit 1
         fi
 
-  nxp-build-test:
-    name: nxp-build-test
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
-    permissions:
-      id-token: write
-      contents: read
-    with:
-      runner: linux.2xlarge
-      docker-image: executorch-ubuntu-22.04-arm-sdk
-      submodules: 'recursive'
-      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-      timeout: 90
-      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}"
-
-        # Build
-        cmake -DEXECUTORCH_BUILD_NXP_NEUTRON=ON -Bcmake-out .
-        cmake --build cmake-out --target executorch_delegate_neutron --config Release
-
-        # Build check for the neutron backend library
-        lib_neutron="cmake-out/backends/nxp/libexecutorch_delegate_neutron.a"
-        if [ -f $lib_neutron ]; then
-            echo "Neutron backend library built."
-        else
-            echo "Neutron backend library not found!"
-            exit 1
-        fi
-
   test-coreml-delegate:
     name: test-coreml-delegate
     uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
@@ -771,35 +741,3 @@ jobs:
       build-mode: Release
       build-tool: cmake
       docker-image: executorch-ubuntu-22.04-clang12
-
-  unittest-nxp-neutron:
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
-    permissions:
-      id-token: write
-      contents: read
-    with:
-      runner: linux.2xlarge
-      docker-image: executorch-ubuntu-22.04-clang12
-      submodules: 'recursive'
-      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-      timeout: 90
-      script: |
-        set -eux
-
-        # 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}"
-
-        # Build and install Executorch
-        PYTHON_EXECUTABLE=python \
-        CMAKE_ARGS="-DEXECUTORCH_BUILD_NXP_NEUTRON=ON" \
-        .ci/scripts/setup-linux.sh --build-tool "cmake"
-
-        # Install test requirements
-        pip install -r backends/nxp/requirements-tests.txt
-
-        # Run pytest
-        PYTHON_EXECUTABLE=python bash backends/nxp/run_unittests.sh
-        
-        # Run aot example: 
-        PYTHON_EXECUTABLE=python bash examples/nxp/run_aot_example.sh

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -424,6 +424,12 @@ class ComputeGraph final {
   // Scalar Value Extraction
   //
 
+  bool is_scalar_or_none(const ValueRef idx) const {
+    const Value& value = values_.at(idx);
+    return value.isInt() || value.isDouble() || value.isBool() ||
+        value.isNone();
+  }
+
   template <typename T>
   T extract_scalar(const ValueRef idx) {
     Value& value = values_.at(idx);
@@ -439,6 +445,15 @@ class ComputeGraph final {
     VK_THROW("Cannot extract scalar from Value with type ", value.type());
   }
 
+  template <typename T>
+  T extract_scalar_or(const ValueRef idx, const T default_value) {
+    Value& value = values_.at(idx);
+    if (value.isNone()) {
+      return default_value;
+    }
+    return extract_scalar<T>(idx);
+  }
+
   template <typename T>
   std::optional<T> extract_optional_scalar(const ValueRef idx) {
     if (val_is_none(idx)) {

backends/vulkan/runtime/graph/ops/glsl/transfer_buffer.glsl

@@ -9,6 +9,7 @@
 #version 450 core
 
 #define PRECISION ${PRECISION}
+#define UBO_PARAMS ${UBO_PARAMS}
 
 #define VEC4_T ${texel_type(DTYPE)}
 #define T ${buffer_scalar_type(DTYPE)}
@@ -22,19 +23,27 @@ layout(std430) buffer;
 ${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")}
 ${layout_declare_tensor(B, "r", "t_in", DTYPE, "buffer")}
 
-$if OP_NAME == "slice":
-  ${layout_declare_ubo(B, "int", "start")}
-  ${layout_declare_ubo(B, "int", "step")}
+$if UBO_PARAMS:
+  $if OP_NAME == "slice":
+    ${layout_declare_ubo(B, "int", "start")}
+    ${layout_declare_ubo(B, "int", "step")}
 
-$if OP_NAME == "select":
-  ${layout_declare_ubo(B, "int", "index")}
+  $if OP_NAME == "select":
+    ${layout_declare_ubo(B, "int", "index")}
 
 layout(push_constant) uniform restrict Block {
   ivec4 in_sizes;
   ivec4 out_strides;
   ivec4 in_strides;
   int out_numel;
   int selected_dim;
+  $if not UBO_PARAMS:
+    $if OP_NAME == "slice":
+      int start;
+      int step;
+
+    $if OP_NAME == "select":
+      int index;
 };
 
 ${layout_declare_spec_const(C, "int", "out_layout", "DEFAULT_LAYOUT")}

backends/vulkan/runtime/graph/ops/glsl/transfer_buffer.yaml

@@ -2,6 +2,7 @@ transfer_buffer:
   parameter_names_with_default_values:
     DTYPE: float
     OP_NAME: select
+    UBO_PARAMS: False
   generate_variant_forall:
     DTYPE:
       - VALUE: half
@@ -11,3 +12,9 @@ transfer_buffer:
       OP_NAME: select
     - NAME: slice_buffer
       OP_NAME: slice
+    - NAME: select_ubo_buffer
+      OP_NAME: select
+      UBO_PARAMS: True
+    - NAME: slice_ubo_buffer
+      OP_NAME: slice
+      UBO_PARAMS: True

backends/vulkan/runtime/graph/ops/glsl/transfer_texture.glsl

@@ -9,6 +9,7 @@
 #version 450 core
 
 #define PRECISION ${PRECISION}
+#define UBO_PARAMS ${UBO_PARAMS}
 
 #define VEC4_T ${texel_type(DTYPE)}
 #define T ${buffer_scalar_type(DTYPE)}
@@ -23,17 +24,25 @@ layout(std430) buffer;
 ${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")}
 ${layout_declare_tensor(B, "r", "t_in", DTYPE, "texture3d")}
 
-$if OP_NAME == "slice":
-  ${layout_declare_ubo(B, "int", "start")}
-  ${layout_declare_ubo(B, "int", "step")}
+$if UBO_PARAMS:
+  $if OP_NAME == "slice":
+    ${layout_declare_ubo(B, "int", "start")}
+    ${layout_declare_ubo(B, "int", "step")}
 
-$if OP_NAME == "select":
-  ${layout_declare_ubo(B, "int", "index")}
+  $if OP_NAME == "select":
+    ${layout_declare_ubo(B, "int", "index")}
 
 layout(push_constant) uniform restrict Block {
   ivec4 out_sizes;
   ivec4 in_sizes;
   int selected_dim;
+  $if not UBO_PARAMS:
+    $if OP_NAME == "slice":
+      int start;
+      int step;
+
+    $if OP_NAME == "select":
+      int index;
 };
 
 ${layout_declare_spec_const(C, "int", "out_layout", "DEFAULT_LAYOUT")}

backends/vulkan/runtime/graph/ops/glsl/transfer_texture.yaml

@@ -2,6 +2,7 @@ transfer_texture:
   parameter_names_with_default_values:
     DTYPE: float
     OP_NAME: select
+    UBO_PARAMS: False
   generate_variant_forall:
     DTYPE:
       - VALUE: half
@@ -11,3 +12,9 @@ transfer_texture:
       OP_NAME: select
     - NAME: slice_texture3d
       OP_NAME: slice
+    - NAME: select_ubo_texture3d
+      OP_NAME: select
+      UBO_PARAMS: True
+    - NAME: slice_ubo_texture3d
+      OP_NAME: slice
+      UBO_PARAMS: True

backends/vulkan/runtime/graph/ops/impl/Transfer.cpp

@@ -40,34 +40,52 @@ void add_transfer_copy_node(
 
   int64_t dim_whcn = nchw_dim_to_whcn_dim(dim, ndim);
 
+  struct TransferParams {
+    int32_t dim;
+    int32_t index_or_start_ref;
+    int32_t step_ref;
+  } transfer_params{static_cast<int32_t>(dim_whcn), 0, 0};
+
+  const bool param_is_scalar = graph.is_scalar_or_none(index_or_start_ref) &&
+      (transfer_type == TransferType::SELECT ||
+       graph.is_scalar_or_none(step_ref));
+
   vkapi::ParamsBindList param_buffers;
-  if (transfer_type == TransferType::SELECT) {
-    param_buffers = {
-        graph.get_or_create_int_param_buffer(index_or_start_ref, 0)};
-  } else { // TransferType::SLICE
-    param_buffers = {
-        graph.get_or_create_int_param_buffer(index_or_start_ref, 0),
-        graph.get_or_create_int_param_buffer(step_ref, 1)};
+  if (!param_is_scalar) {
+    if (transfer_type == TransferType::SELECT) {
+      param_buffers = {
+          graph.get_or_create_int_param_buffer(index_or_start_ref, 0)};
+    } else { // TransferType::SLICE
+      param_buffers = {
+          graph.get_or_create_int_param_buffer(index_or_start_ref, 0),
+          graph.get_or_create_int_param_buffer(step_ref, 1)};
+    }
+  } else {
+    transfer_params.index_or_start_ref =
+        graph.extract_scalar_or<int32_t>(index_or_start_ref, 0);
+    if (transfer_type != TransferType::SELECT) {
+      transfer_params.step_ref = graph.extract_scalar_or<int32_t>(step_ref, 1);
+    }
   }
 
-  const struct TransferParams {
-    const int32_t dim;
-  } transfer_params{static_cast<int32_t>(dim_whcn)};
-
   std::vector<PushConstantDataInfo> push_constants;
+  push_constants.reserve(graph.is_buffer_storage(out) ? 5 : 3);
 
   if (graph.is_buffer_storage(out)) {
-    push_constants = {
-        graph.sizes_pc_of(in),
-        graph.strides_pc_of(out),
-        graph.strides_pc_of(in),
-        graph.numel_pc_of(out),
-        PushConstantDataInfo(&transfer_params, sizeof(transfer_params))};
+    push_constants.emplace_back(graph.sizes_pc_of(in));
+    push_constants.emplace_back(graph.strides_pc_of(out));
+    push_constants.emplace_back(graph.strides_pc_of(in));
+    push_constants.emplace_back(graph.numel_pc_of(out));
   } else {
-    push_constants = {
-        graph.sizes_pc_of(out),
-        graph.sizes_pc_of(in),
-        PushConstantDataInfo(&transfer_params, sizeof(transfer_params))};
+    push_constants.emplace_back(graph.sizes_pc_of(out));
+    push_constants.emplace_back(graph.sizes_pc_of(in));
+  }
+
+  if (param_is_scalar) {
+    push_constants.emplace_back(&transfer_params, sizeof(transfer_params));
+  } else {
+    push_constants.emplace_back(
+        &transfer_params.dim, sizeof(transfer_params.dim));
   }
 
   vkapi::SpecVarList spec_vars = {
@@ -82,6 +100,9 @@ void add_transfer_copy_node(
   } else { // TransferType::SLICE
     kernel_name = "slice";
   }
+  if (!param_is_scalar) {
+    kernel_name += "_ubo";
+  }
   add_storage_type_suffix(kernel_name, graph.storage_type_of(out));
   add_dtype_suffix(kernel_name, graph.dtype_of(out));
 

backends/xnnpack/cmake/Dependencies.cmake

@@ -83,10 +83,10 @@ if(EXECUTORCH_XNNPACK_ENABLE_KLEIDI)
     KLEIDIAI_LIBRARY kleidiai
     PATHS "${CMAKE_CURRENT_BINARY_DIR}/XNNPACK/kleidiai-source"
   )
-  if(not KLEIDIAI_LIBRARY)
-    message(FATAL_ERROR "Can't find KleidiAI")
+  if(KLEIDIAI_LIBRARY)
+    message(WARNING "Can't find KleidiAI, skipping install")
+    install(FILES ${KLEIDIAI_LIBRARY} PUBLIC_HEADER
+            DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
+    )
   endif()
-  install(FILES ${KLEIDIAI_LIBRARY} PUBLIC_HEADER
-          DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
-  )
 endif()