Update on "[ET Device Support] Schema changes: device info on Tensor and buffer-level device array"

This diff adds device placement information to the ExecuTorch schema to support representing tensor-level device type information, which will be the basic requirement for the following tensor_parser updates.

This is part of the Phase 1 implementation to make ET device type work E2E without user-specified device placement.

Design doc: https://docs.google.com/document/d/1lwd9BlohmwkN5EEvRulO_b-XnZBwv1nMb5l2K3jfuwA/edit?tab=t.0#heading=h.o6anuvkix4bu

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

[ghstack-poisoned]

Author: Gasoonjia

.ci/scripts/test_backend.sh

@@ -46,7 +46,7 @@ if [[ "$FLOW" == *qnn* ]]; then
     export LD_LIBRARY_PATH"=$QNN_X86_LIB_DIR:$QNN_SDK_ROOT/lib/x86_64-linux-clang/:${LD_LIBRARY_PATH:-}"
 
     # TODO Get SDK root from install scripts
-    EXTRA_BUILD_ARGS+=" -DEXECUTORCH_BUILD_QNN=ON -DQNN_SDK_ROOT=$QNN_SDK_ROOT"
+    EXTRA_BUILD_ARGS+=" -DEXECUTORCH_BUILD_QNN=ON -DQNN_SDK_ROOT=$QNN_SDK_ROOT -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON"
 fi
 
 if [[ "$FLOW" == *vulkan* ]]; then

.github/workflows/_unittest.yml

@@ -58,10 +58,11 @@ jobs:
     if: ${{ inputs.build-tool == 'cmake' }}
     uses: pytorch/test-infra/.github/workflows/windows_job.yml@main
     with:
-      submodules: 'recursive'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       timeout: 120
       script: |
+        git config --global http.sslBackend openssl
+        git submodule update --init --recursive 
         conda init powershell
 
         powershell -Command "& {

.github/workflows/build-presets.yml

@@ -113,10 +113,11 @@ jobs:
     with:
       job-name: build
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-      submodules: recursive
       timeout: 90
       script: |
         set -eux
+        git config --global http.sslBackend openssl
+        git submodule update --init --recursive
         conda init powershell
         powershell -Command "& {
           Set-PSDebug -Trace 1

.github/workflows/cuda-windows.yml

@@ -127,10 +127,11 @@ jobs:
       runner: windows.g5.4xlarge.nvidia.gpu
       gpu-arch-type: cuda
       gpu-arch-version: 12.8
-      submodules: recursive
       download-artifact: ${{ matrix.model_repo }}-${{ matrix.model_name }}-cuda-windows-${{ matrix.quant }}
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       script: |
+        git config --global http.sslBackend openssl
+        git submodule update --init --recursive
         conda init powershell
         powershell -Command "& {
           Set-PSDebug -Trace 1

.github/workflows/trunk.yml

@@ -1123,10 +1123,11 @@ jobs:
         model: [mv3, resnet50, vit, mobilebert, emformer_transcribe]
         backend: [portable, xnnpack-q8]
     with:
-      submodules: 'recursive'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       timeout: 60
       script: |
+        git config --global http.sslBackend openssl
+        git submodule update --init --recursive
         conda init powershell
 
         powershell -Command "& {

.github/workflows/windows-msvc.yml

@@ -22,10 +22,11 @@ jobs:
     name: build-windows-msvc
     uses: pytorch/test-infra/.github/workflows/windows_job.yml@main
     with:
-      submodules: 'recursive'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       timeout: 60
       script: |
+        git config --global http.sslBackend openssl
+        git submodule update --init --recursive
         conda init powershell
         powershell -Command "& {
           Set-PSDebug -Trace 1

backends/arm/_passes/__init__.py

@@ -52,6 +52,7 @@
 from .decompose_glu_pass import DecomposeGluPass  # noqa
 from .decompose_grouped_conv_pass import DecomposeGroupedConvPass  # noqa
 from .decompose_groupnorm_pass import DecomposeGroupNormPass  # noqa
+from .decompose_index_copy_pass import DecomposeIndexCopyPass  # noqa
 from .decompose_index_select_to_gather_pass import (  # noqa
     DecomposeIndexSelectToGatherPass,
 )
@@ -94,6 +95,7 @@
 from .decompose_tril_pass import DecomposeTrilPass  # noqa
 from .decompose_unfold_to_gather_pass import DecomposeUnfoldToGatherPass  # noqa
 from .decompose_var_pass import DecomposeVarPass  # noqa
+from .decompose_where_scalar_other_pass import DecomposeWhereScalarOtherPass  # noqa
 from .decorate_fp32_to_int32_casting_pass import DecorateFp32toInt32CastingPass  # noqa
 from .fold_qdq_with_annotated_qparams_pass import (  # noqa
     FoldAndAnnotateQParamsPass,

backends/arm/_passes/arm_pass_manager.py

@@ -59,6 +59,7 @@
     DecomposeGluPass,
     DecomposeGroupedConvPass,
     DecomposeGroupNormPass,
+    DecomposeIndexCopyPass,
     DecomposeIndexSelectToGatherPass,
     DecomposeIndexTensorToGatherPass,
     DecomposeIntPowPass,
@@ -92,6 +93,7 @@
     DecomposeTrilPass,
     DecomposeUnfoldToGatherPass,
     DecomposeVarPass,
+    DecomposeWhereScalarOtherPass,
     DecorateFp32toInt32CastingPass,
     FoldAndAnnotateQParamsPass,
     FuseBatchNorm2dPass,
@@ -320,6 +322,7 @@ def _tosa_pipeline(
             [
                 ReplaceScalarWithTensorByProfilePass(),
                 RewriteLeLtToGeGtPass(),
+                DecomposeLeakyReLUPass(),  # Emits full_like so before ConvertFullLikeToFullPass
                 ConvertFullLikeToFullPass(),
                 MatchArgDtypePass(),
                 UnsqueezeScalarPlaceholdersPass(exported_program),
@@ -340,7 +343,6 @@ def _tosa_pipeline(
                 FuseBatchNorm2dPass(exported_program),
                 ConvertMmToBmmPass(),
                 DecomposeGluPass(),
-                DecomposeLeakyReLUPass(),
                 DecomposeDivPass(),
                 # _safe_softmax results in a ReduceMax
                 # which is not currently supported by TOSA in U55
@@ -418,6 +420,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
         # Transformation passes (pre scalar -> tensor)
         self.add_passes(
             [
+                DecomposeIndexCopyPass(tfa_pass=True),
                 DecomposeSelectScatterPass(tfa_pass=True),
                 DecomposeSliceScatterPass(tfa_pass=True),
                 ConvertInt64ConstOpsToInt32Pass(tfa_pass=True),
@@ -434,6 +437,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                 DecomposeRemainderPass(tfa_pass=True),
                 DecomposeFloorDividePass(tfa_pass=True),
                 DecomposeDivTensorModePass(tfa_pass=True),
+                DecomposeWhereScalarOtherPass(tfa_pass=True),
             ]
         )
 

backends/arm/_passes/decompose_index_copy_pass.py

@@ -0,0 +1,33 @@
+# Copyright 2026 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.
+
+from typing import Set, Type
+
+import torch
+from executorch.backends.arm._passes.get_decomposition_pass import GetDecompositionPass
+from executorch.backends.arm._passes.insert_int32_casts_after_int64_placeholders import (
+    InsertInt32CastsAfterInt64PlaceholdersPass,
+)
+from executorch.exir.pass_base import ExportPass
+
+
+class DecomposeIndexCopyPass(GetDecompositionPass):
+    """Decomposes aten.index_copy into aten.index_put, as well as it's
+    surrounding operators.
+
+    This pass is intended to be called in transform_for_annotation to prepare
+    the graph for quantization. After quantization, this operator will be
+    prepared for lowering to TOSA using the RewriteIndexPut pass
+
+    """
+
+    _passes_required_after: Set[Type[ExportPass]] = {
+        InsertInt32CastsAfterInt64PlaceholdersPass
+    }
+
+    targeted_ops = [
+        torch.ops.aten.index_copy.default,
+        torch.ops.aten.index_copy_.default,
+    ]

backends/arm/_passes/decompose_leaky_relu_pass.py

@@ -9,6 +9,9 @@
 
 import torch
 from executorch.backends.arm._passes import ArmPass
+from executorch.backends.arm._passes.convert_full_like_to_full_pass import (
+    ConvertFullLikeToFullPass,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
@@ -20,14 +23,14 @@ def _get_leaky_relu_ops(op) -> tuple:
     if op in edge_ops:
         return (
             exir_ops.edge.aten.clamp.default,
-            exir_ops.edge.aten.full.default,
+            exir_ops.edge.aten.full_like.default,
             exir_ops.edge.aten.mul.Tensor,
             exir_ops.edge.aten.add.Tensor,
         )
     elif op in torch_ops:
         return (
             torch.ops.aten.clamp.default,
-            torch.ops.aten.full.default,
+            torch.ops.aten.full_like.default,
             torch.ops.aten.mul.Tensor,
             torch.ops.aten.add.Tensor,
         )
@@ -42,33 +45,31 @@ class DecomposeLeakyReLUPass(ArmPass):
     Example:
         %op1 = clamp(x,0,None) (equivalent to max(0,x))
         %op2 = clamp(x,None,0) (equivalent to min(0,x))
-        %op3 = full(x.shape,slope)
+        %op3 = full_like(x,slope)
         %op4 = mul(%op3,%op2)
         %op5 = add(%op1,%op4)
 
     """
 
-    _passes_required_after: Set[Type[ExportPass]] = set()
+    _passes_required_after: Set[Type[ExportPass]] = {ConvertFullLikeToFullPass}
 
     def call_operator(self, op, args, kwargs, meta):
         if op not in (edge_ops + torch_ops) or not self.allowed_to_transform(meta):
             return super().call_operator(op, args, kwargs, meta)
 
         x = args[0]
         slope = args[1] if len(args) > 1 else 0.01
-        dtype = x.node.meta["val"].dtype
-        device = x.node.meta["val"].device
-        clamp, full, mul, add = _get_leaky_relu_ops(op)
+        clamp, full_like, mul, add = _get_leaky_relu_ops(op)
         op1 = super().call_operator(
             op=clamp, args=(x, 0, None), kwargs=kwargs, meta=meta
         )
         op2 = super().call_operator(
             op=clamp, args=(x, None, 0), kwargs=kwargs, meta=meta
         )
         op3 = super().call_operator(
-            op=full,
-            args=(x.node.meta["val"].shape, slope),
-            kwargs={"dtype": dtype, "device": device},
+            op=full_like,
+            args=(x, slope),
+            kwargs={},
             meta=meta,
         )
         op4 = super().call_operator(op=mul, args=(op3, op2), kwargs=kwargs, meta=meta)