Skip to content

Commit a208948

Browse files
author
Jorge Pineda
committed
Update on "[ET-VK][Ops] aten.index_select"
## The Operator `nn.Module` invocations of [`torch.index_select`](https://pytorch.org/docs/stable/generated/torch.index_select.html) get compiled to `aten.index_select.default` in the Edge Dialect, which carries the following signature. ``` - func: index_select(Tensor self, int dim, Tensor index) -> Tensor ``` ## Implementation This is a C-packing-only implementation. It is very similar to `aten.slice`: #3171 ``` - func: slice.Tensor(Tensor(a) self, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor(a) ``` It features a similar split between a shader for N,H,W and a shader for C, because copying from the C-dimension is more difficult due to C-packing. Both `index_select` and `slice` copy specific indices across 1 dimension. The difference is in the way these indices are specified. - `slice` uses `start=1`/`end=5`/`step=2` as three scalars for indices `1,3`. - `index_select` lists the exact indices inside a tensor e.g. `index=torch.tensor([1,3])`. Hence, `slice` uses a `offset=1` and `step=2` to compute input position. In `index_select`, we read the index tensor to compute input position. Differential Revision: [D57745489](https://our.internmc.facebook.com/intern/diff/D57745489/) [ghstack-poisoned]
2 parents 7b36eff + 10ec068 commit a208948

File tree

26 files changed

+522
-75
lines changed

26 files changed

+522
-75
lines changed

‎.ci/scripts/setup-macos.sh‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -77,16 +77,18 @@ install_sccache() {
7777

7878
export PATH="${SCCACHE_PATH}:${PATH}"
7979

80-
# Create temp directory for sccache shims
81-
TMP_DIR=$(mktemp -d)
82-
trap 'rm -rfv ${TMP_DIR}' EXIT
80+
# Create temp directory for sccache shims if TMP_DIR doesn't exist
81+
if [ -z "${TMP_DIR:-}" ]; then
82+
TMP_DIR=$(mktemp -d)
83+
trap 'rm -rfv ${TMP_DIR}' EXIT
84+
export PATH="${TMP_DIR}:$PATH"
85+
fi
8386

8487
write_sccache_stub "${TMP_DIR}/c++"
8588
write_sccache_stub "${TMP_DIR}/cc"
8689
write_sccache_stub "${TMP_DIR}/clang++"
8790
write_sccache_stub "${TMP_DIR}/clang"
8891

89-
export PATH="${TMP_DIR}:$PATH"
9092
sccache --zero-stats || true
9193
}
9294

‎.ci/scripts/utils.sh‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#!/bin/bash
22
# Copyright (c) Meta Platforms, Inc. and affiliates.
33
# All rights reserved.
4+
# Copyright 2024 Arm Limited and/or its affiliates.
45
#
56
# This source code is licensed under the BSD-style license found in the
67
# LICENSE file in the root directory of this source tree.
@@ -89,6 +90,17 @@ install_flatc_from_source() {
8990
popd || return
9091
}
9192

93+
install_arm() {
94+
# NB: This function could be used to install Arm dependencies
95+
# Setup arm example environment (including TOSA tools)
96+
git config --global user.email "[email protected]"
97+
git config --global user.name "Github Executorch"
98+
bash examples/arm/setup.sh --i-agree-to-the-contained-eula
99+
100+
# Test tosa_reference flow
101+
source examples/arm/ethos-u-scratch/setup_path.sh
102+
}
103+
92104
build_executorch_runner_buck2() {
93105
# Build executorch runtime with retry as this step is flaky on macos CI
94106
retry buck2 build //examples/portable/executor_runner:executor_runner

‎.github/workflows/_unittest.yml‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,6 @@ jobs:
4141
4242
macos:
4343
uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
44-
strategy:
45-
matrix:
46-
include:
47-
- build-tool: buck2
4844
with:
4945
runner: macos-m1-stable
5046
python-version: '3.11'
@@ -53,18 +49,21 @@ jobs:
5349
script: |
5450
set -eux
5551
56-
BUILD_TOOL=${{ matrix.build-tool }}
57-
5852
bash .ci/scripts/setup-conda.sh
5953
54+
# Create temp directory for sccache shims
55+
export TMP_DIR=$(mktemp -d)
56+
export PATH="${TMP_DIR}:$PATH"
57+
trap 'rm -rfv ${TMP_DIR}' EXIT
58+
6059
# Setup MacOS dependencies as there is no Docker support on MacOS atm
6160
PYTHON_EXECUTABLE=python \
6261
EXECUTORCH_BUILD_PYBIND=ON \
6362
CMAKE_ARGS="-DEXECUTORCH_BUILD_COREML=ON -DEXECUTORCH_BUILD_MPS=ON -DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON" \
6463
${CONDA_RUN} --no-capture-output \
65-
.ci/scripts/setup-macos.sh "${BUILD_TOOL}"
64+
.ci/scripts/setup-macos.sh cmake
6665
6766
# Run pytest with coverage
6867
${CONDA_RUN} pytest -n auto --cov=./ --cov-report=xml
6968
# Run gtest
70-
${CONDA_RUN} buck2 test runtime/core/... runtime/platform/...
69+
${CONDA_RUN} test/run_oss_cpp_tests.sh

‎.github/workflows/pull.yml‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -307,3 +307,37 @@ jobs:
307307
uses: ./.github/workflows/_unittest.yml
308308
with:
309309
docker-image: executorch-ubuntu-22.04-clang12
310+
311+
unittest-arm:
312+
uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
313+
strategy:
314+
matrix:
315+
include:
316+
- build-tool: buck2
317+
with:
318+
runner: linux.2xlarge
319+
docker-image: executorch-ubuntu-22.04-arm-sdk
320+
submodules: 'true'
321+
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
322+
timeout: 90
323+
script: |
324+
set -eux
325+
326+
# The generic Linux job chooses to use base env, not the one setup by the image
327+
CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
328+
conda activate "${CONDA_ENV}"
329+
330+
BUILD_TOOL=${{ matrix.build-tool }}
331+
332+
# Setup MacOS dependencies as there is no Docker support on MacOS atm
333+
PYTHON_EXECUTABLE=python \
334+
EXECUTORCH_BUILD_PYBIND=ON \
335+
EXECUTORCH_BUILD_ARM_BAREMETAL=ON \
336+
.ci/scripts/setup-linux.sh "${BUILD_TOOL}"
337+
338+
source .ci/scripts/utils.sh
339+
# Install Arm dependencies
340+
install_arm
341+
342+
# Run pytest with coverage
343+
pytest -c /dev/null -v -n auto --cov=./ --cov-report=xml backends/arm/test

‎.github/workflows/trunk.yml‎

Lines changed: 2 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -142,17 +142,13 @@ jobs:
142142
install_flatc_from_source
143143
install_executorch
144144
145-
# Setup arm example environment (including TOSA tools)
146-
git config --global user.email "[email protected]"
147-
git config --global user.name "Github Executorch"
148-
bash examples/arm/setup.sh --i-agree-to-the-contained-eula
145+
install_arm
149146
150147
# Increase number of files user can monitor to bypass buck failures.
151148
# Hopefully this is high enough for this setup.
152149
sudo sysctl fs.inotify.max_user_watches=1048576 # 1024 * 1024
153150
154151
# Test ethos-u delegate examples with run.sh
155-
source examples/arm/ethos-u-scratch/setup_path.sh
156152
PYTHON_EXECUTABLE=python bash examples/arm/run.sh examples/arm/ethos-u-scratch/ buck2
157153
158154
test-arm-reference-delegation:
@@ -172,20 +168,11 @@ jobs:
172168
install_flatc_from_source
173169
install_executorch
174170
175-
# Setup arm example environment (including TOSA tools)
176-
git config --global user.email "[email protected]"
177-
git config --global user.name "Github Executorch"
178-
bash examples/arm/setup.sh --i-agree-to-the-contained-eula
171+
install_arm
179172
180173
# Test tosa_reference flow
181-
source examples/arm/ethos-u-scratch/setup_path.sh
182174
PYTHON_EXECUTABLE=python bash backends/arm/test/run_tosa_reference.sh
183175
184-
# Run Arm specific unit-tests
185-
# Run pytest only on specified folders. These test should migrate into
186-
# the _unittest.yml once that test env is fixed
187-
pytest -c /dev/null -v -n auto --cov=./ --cov-report=xml backends/arm/test/ops/ backends/arm/test/models
188-
189176
test-coreml-delegate:
190177
name: test-coreml-delegate
191178
uses: pytorch/test-infra/.github/workflows/macos_job.yml@main

‎.lintrunner.toml‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,3 +156,33 @@ command = [
156156
'--',
157157
'@{{PATHSFILE}}',
158158
]
159+
160+
[[linter]]
161+
code = 'NEWLINE'
162+
include_patterns = ['**']
163+
exclude_patterns = [
164+
'third-party/**',
165+
'**/third-party/**',
166+
'**/*.png',
167+
'**/*.webp',
168+
'**/*.jpeg',
169+
'**/*.mp4',
170+
'**/*.pte',
171+
'**/*.pth',
172+
'**/*.bin',
173+
'**/*.patch',
174+
'**/*.svg',
175+
'**/*.bat',
176+
'**/*.jpg',
177+
'**/*.jar',
178+
]
179+
command = [
180+
'python',
181+
'-m',
182+
'lintrunner_adapters',
183+
'run',
184+
'newlines_linter',
185+
'--',
186+
'@{{PATHSFILE}}',
187+
]
188+
is_formatter = true

‎backends/apple/coreml/.clang-format‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,4 +25,3 @@ BraceWrapping:
2525

2626
# Options for aligning backslashes in escaped newlines.
2727
AlignEscapedNewlines: Left
28-

‎backends/apple/coreml/runtime/delegate/backend_delegate.mm‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -200,4 +200,3 @@ bool purge_models_cache() const noexcept override {
200200
return std::make_shared<BackendDelegateImpl>(config);
201201
}
202202
} //namespace executorchcoreml
203-

‎backends/apple/coreml/runtime/delegate/coreml_backend_delegate.mm‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -235,4 +235,3 @@ ModelLoggingOptions get_logging_options(BackendExecutionContext& context) {
235235

236236
} // namespace executor
237237
} // namespace torch
238-

‎backends/vulkan/runtime/graph/ops/glsl/index_select.yaml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,5 +7,6 @@ index_select:
77
DTYPE:
88
- VALUE: half
99
- VALUE: float
10+
- VALUE: int
1011
shader_variants:
1112
- NAME: index_select

0 commit comments

Comments
 (0)