Update base for Update on "[ET-VK][Ops] torchao.choose_qparams_affine vulkan impl and shader (buffer only) and cleanup"

# Changes
* Implement `torchao.choose_qparams_affine` operator in Vulkan backend with comprehensive buffer storage support
* Add block-wise quantization parameter computation in `choose_qparams_buffer.glsl` shader for configurable tensor block analysis
* Extend quantization parameter infrastructure in `ChooseQParams.cpp` to handle affine transformations with configurable block sizes and multiple mapping types
* Support three quantization mapping strategies: ASYMMETRIC, SYMMETRIC, and SYMMETRIC_NO_CLIPPING_ERR for optimal parameter selection
* Consolidated the logic for choosing scale and zero point between affine cases and regular quantized_decomposed cases.

BE: Improved the documentation in the shader logic which is more detailed and clear

# Motivation
The existing Vulkan quantization infrastructure lacked support for the `torchao.choose_qparams_affine` operator, which is essential for computing optimal quantization parameters in dynamic quantization workflows. The `choose_qparams_affine` operator provides flexible block-wise parameter computation that analyzes statistical distributions within tensor blocks, enabling:

* **Block-wise Parameter Computation**: Analyzes configurable tensor blocks to compute optimal scale and zero-point values, improving quantization accuracy for heterogeneous data distributions
* **Multiple Mapping Types**: Supports ASYMMETRIC, SYMMETRIC, and SYMMETRIC_NO_CLIPPING_ERR quantization strategies for different precision-performance trade-offs

# Operator Description
The `choose_qparams_affine` operator computes optimal quantization parameters (scale and zero_point) from floating-point tensor blocks using statistical analysis of data distributions. Block-wise parameter computation divides tensors into blocks and analyzes each block independently to determine the best quantization mapping for subsequent quantization operations.

The parameter calculation varies by mapping type:
- **ASYMMETRIC**: `scale = (max - min) / (quant_max - quant_min)`, `zero_point = quant_min - round(min / scale)`
- **SYMMETRIC**: `scale = max_abs / ((quant_max - quant_min) / 2)`, `zero_point = midpoint`
- **SYMMETRIC_NO_CLIPPING_ERR**: `scale = max(abs(min)/abs(quant_min), max/quant_max)`, `zero_point = midpoint`

**Storage Requirements**: Input tensors must be floating-point (kFloat) with width-packed layout. Output scale/zero_point tensors use buffer storage.

NOTE: Texture storage implementation is not supported due to complexity of block-wise coordinate mapping in 3D texture space. This will likely be necessary for better efficiency in the future.

# Block-wise Parameter Computation Implementation
Block-wise parameter computation enables fine-grained quantization analysis by dividing tensors into blocks and computing separate scale/zero_point parameters for each block. The implementation uses several key data structures computed in `ChooseQParams.cpp`:

* **`block_size_vec`**: WHCN-ordered block dimensions converted from PyTorch NCHW layout (e.g., [3,3,2,1] for 3×3×2×1 blocks)
* **`tensor_size_whcn`**: Input tensor dimensions converted to WHCN layout using `utils::make_whcn_ivec4()`
* **`num_blocks_vec`**: Number of blocks per dimension calculated as `ceil(tensor_size_whcn / block_size_vec)` to handle non-divisible dimensions
* **`block_stride_vec`**: Pre-computed linear strides for block grid indexing `{1, #W, #W*#H, #W*#H*#C}` to enable efficient block ID calculation
* **`mapping_type`**: Integer encoding of quantization strategy (0=ASYMMETRIC, 1=SYMMETRIC, 2=SYMMETRIC_NO_CLIPPING_ERR)

The block coordinate calculation uses: `block_coord = block_id_to_coord(block_id)` which converts linear block IDs back to 4D WHCN coordinates, then computes element ranges: `t0 = block_coord * blockSize` and `tEnd = t0 + blockSize` for nested loop iteration.

# Shader Algorithm Overview

## Buffer Storage Implementation (`choose_qparams_buffer.glsl`)

**Workgroup Configuration**:
- **Global WG Size**: `{nBlocks, 1u, 1u}` where `nBlocks = total number of blocks` computed from `ceil(tensor_size / block_size)` for each dimension
- **Local WG Size**: `{1u, 1u, 1u}` (single thread per block for simplicity, though could be optimized for larger blocks)

**Block-wise Mode Algorithm**:
The shader uses a sophisticated multi-level nested approach to process tensor blocks efficiently. Each thread is assigned multiple blocks using strided access: `for (uint block_id = gl_GlobalInvocationID.x; block_id < TOTAL_BLOCKS; block_id += STRIDE)` where `STRIDE = gl_WorkGroupSize.x * gl_NumWorkGroups.x`.

For each assigned block, the algorithm performs several key steps:

**1. Block Coordinate Conversion**:
The `block_id_to_coord(block_id)` function converts linear block IDs to 4D WHCN coordinates using modular arithmetic.

**2. Element Range Calculation**: Computes the inclusive start coordinate `t0 = bc * blockSize` and exclusive end coordinate `tEnd = t0 + blockSize` to define the block's element boundaries in tensor space.

**3. Nested Loop Min/Max Scan**: Uses four nested loops to iterate through all elements within the block:

`for (int n = t0.w; n < tEnd.w; ++n) for (int c = t0.z; c < tEnd.z; ++c) for (int h = t0.y; h < tEnd.y; ++h) for (int w = t0.x; w < tEnd.x; ++w)`

Each element is accessed using `tidx_to_bufi(ivec4(w,h,c,n), t_in_strides)` to convert 4D tensor coordinates to linear buffer indices with proper stride handling.

**4. Parameter Calculation**: Calls `calc_scale_zp(lo, hi, quant_min, quant_max, mapping_type, eps, scale, zp)` which implements the three mapping strategies:

*   **ASYMMETRIC (mapping_type=0)**: Maps full range [min, max] to [quant_min, quant_max] preserving data distribution
*   **SYMMETRIC (mapping_type=1)**: Centers around zero using `max_abs = max(abs(min), abs(max))` for balanced quantization
*   **SYMMETRIC_NO_CLIPPING_ERR (mapping_type=2)**: Computes separate scales for positive/negative ranges and uses the maximum to prevent clipping

**Future Improvements**: Implement workgroup-level reduction for large blocks, optimize memory access patterns for better cache utilization, and explore texture storage implementation with simplified block alignment constraints.

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

cc SS-JIA manuelcandales cbilgin

[ghstack-poisoned]

Author: morelos

.github/workflows/trunk.yml

@@ -302,6 +302,37 @@ jobs:
           exit 1
         fi
 
+  test-arm-ootb-linux:
+    name: test-arm-ootb-linux
+    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}"
+
+        # Follow the steps required before running the notebooks
+        # Try to mirror these as closely as possible
+        source .ci/scripts/utils.sh
+        install_executorch "--use-pt-pinned-commit"
+
+        .ci/scripts/setup-arm-baremetal-tools.sh
+        source examples/arm/ethos-u-scratch/setup_path.sh
+
+        # Install requirements for converting notebooks
+        pip install notebook
+
+        # Run OOTB tests
+        backends/arm/test/test_arm_ootb.sh
+
   test-coreml-delegate:
     name: test-coreml-delegate
     uses: pytorch/test-infra/.github/workflows/macos_job.yml@main

CMakePresets.json

@@ -15,7 +15,8 @@
         "CMAKE_TOOLCHAIN_FILE": "${sourceDir}/third-party/ios-cmake/ios.toolchain.cmake",
         "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/macos.cmake",
         "PLATFORM": "MAC_ARM64",
-        "DEPLOYMENT_TARGET": "12.0"
+        "DEPLOYMENT_TARGET": "12.0",
+        "CMAKE_MACOSX_BUNDLE": "OFF"
       },
       "condition": {
         "lhs": "${hostSystemName}",

backends/arm/_passes/decompose_asin_pass.py

@@ -85,12 +85,11 @@ def _build_polynomial(
         return result
 
     def call_operator(self, op, args, kwargs, meta):
+        if op not in edge_asin_op:
+            return super().call_operator(op, args, kwargs, meta)
         logging.info(
             f"Approximating asin. This may introduce small numerical errors. For details, see {__file__}."
         )
-        if op not in edge_asin_op:
-            return super().call_operator(op, args, kwargs, meta)
-
         x = args[0]
         half = 0.5
         one = 1.0

backends/arm/test/ops/test_add.py

@@ -205,6 +205,5 @@ def test_add_tensor_vgf_int(test_data: input_t1):
         aten_op,
         exir_op,
         tosa_version="TOSA-1.0+INT",
-        symmetric_io_quantization=True,
     )
     pipeline.run()

backends/arm/test/test_arm_ootb.sh

@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+
+# Copyright 2025 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.
+
+set -e
+
+run_ootb_tests_ethos_u() {
+    echo "$FUNCNAME: Running out-of-the-box tests for Arm Ethos-U"
+    jupyter nbconvert \
+        --to notebook \
+        --execute examples/arm/ethos_u_minimal_example.ipynb
+    echo "${FUNCNAME}: PASS"
+}
+
+run_ootb_tests_ethos_u

backends/arm/test/tester/arm_tester.py

@@ -43,6 +43,7 @@
     EthosUQuantizer,
     get_symmetric_quantization_config,
     TOSAQuantizer,
+    VgfQuantizer,
 )
 from executorch.backends.arm.test.runner_utils import (
     dbg_tosa_fb_to_json,
@@ -332,6 +333,8 @@ def quantize(
                 quantizer = TOSAQuantizer(tosa_spec)
             elif is_ethosu(self.compile_spec):
                 quantizer = EthosUQuantizer(self.compile_spec)
+            elif is_vgf(self.compile_spec):
+                quantizer = VgfQuantizer(self.compile_spec)
             quantize_stage = tester.Quantize(
                 quantizer,
                 get_symmetric_quantization_config(),

backends/arm/test/tester/test_pipeline.py

@@ -861,18 +861,15 @@ def __init__(
         rtol: float = 1e-03,
         qtol: int = 1,
         dynamic_shapes: Optional[Tuple[Any]] = None,
+        transform_passes: Optional[
+            Union[Sequence[PassType], Dict[str, Sequence[PassType]]]
+        ] = None,
     ):
 
         if (
             symmetric_io_quantization or per_channel_quantization
         ) and tosa_version == "TOSA-1.0+FP":
             raise ValueError("Dont configure quantization with FP TOSA profile.")
-        if (
-            symmetric_io_quantization is False
-            and per_channel_quantization is False
-            and tosa_version == "TOSA-1.0+INT"
-        ):
-            raise ValueError("Missing quantization options for INT TOSA profile.")
 
         tosa_profile = TosaSpecification.create_from_string(tosa_version)
         compile_spec = common.get_vgf_compile_spec(
@@ -887,6 +884,7 @@ def __init__(
             exir_op,
             use_to_edge_transform_and_lower,
             dynamic_shapes,
+            transform_passes=transform_passes,
         )
 
         if symmetric_io_quantization or per_channel_quantization:
@@ -900,7 +898,7 @@ def __init__(
         else:
             quant_stage = None
 
-        if quant_stage:
+        if "INT" in tosa_version:
             self.add_stage(self.tester.quantize, quant_stage, pos=0)
 
             self.add_stage_after(

backends/cadence/aot/TARGETS

@@ -539,8 +539,10 @@ python_unittest(
     ],
     typing = True,
     deps = [
+        ":program_builder",
         "//executorch/backends/cadence/aot:graph_builder",
         "//executorch/backends/cadence/aot:ops_registrations",
+        "//executorch/runtime:runtime",
         "//later:lib",
     ],
 )

backends/cadence/aot/memory_planning.py

@@ -19,7 +19,10 @@
     MemoryPlanningAlgo,
     MemoryPlanningState,
 )
-from executorch.backends.cadence.aot.utils import MemoryConfig
+from executorch.backends.cadence.aot.utils import (
+    MemoryConfig,
+    MemoryPlanningAlgoFailure,
+)
 
 from executorch.exir import ExecutorchProgramManager
 from executorch.exir.memory_planning import collect_specs_from_nodes, Verifier
@@ -95,7 +98,9 @@ def plan(
         ):
             self.plan_spec(spec, state, placement_constraints)
             if not state.is_placed(spec):
-                raise MemoryError(f"Cannot fit {spec} in any memory hierarchy")
+                raise MemoryPlanningAlgoFailure(
+                    f"Cannot fit {spec} {spec.allocated_memory=} in any memory hierarchy for {self.memory_config}"
+                )
 
 
 class GreedyWithHeuristic(MemoryPlanningAlgo):
@@ -169,7 +174,9 @@ def plan(
         ):
             self.plan_spec(spec, state, placement_constraints)
             if not state.is_placed(spec):
-                raise MemoryError(f"Cannot fit {spec} in any memory hierarchy")
+                raise MemoryPlanningAlgoFailure(
+                    f"Cannot fit {spec} in any memory hierarchy for {self.memory_config}"
+                )
 
         logging.debug(
             f"greedy by size for offset calculation with hierarchy returns bufsizes: {state.bufsizes}"

backends/cadence/aot/tests/test_fusion_ops_passes.py

@@ -12,7 +12,6 @@
 
 import executorch.backends.cadence.aot.ops_registrations  # noqa
 import torch
-from executorch.backends.cadence.aot import compiler
 from executorch.backends.cadence.aot.fuse_ops import (
     FuseCascadedTransposeOrPermuteOps,
     FuseCascadedViewOps,
@@ -30,7 +29,6 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.pass_base import PassResult, ProxyValue
-from torch import nn
 
 
 class TestFusionPassesBase(unittest.TestCase):
@@ -178,43 +176,6 @@ def test_keep_mm_add_with_multiple_users(self) -> None:
         self.assertEqual(count_node(converted_graph, exir_ops.edge.aten.mm.default), 1)
         self.assertEqual(count_node(converted_graph, exir_ops.edge.aten.add.Tensor), 3)
 
-    # TODO(matthiascremon) -> None: enable that pass with new flow
-    @torch.no_grad()
-    @unittest.expectedFailure
-    def test_legacy_conv_bn_fusion(self) -> None:
-        class ModelConvBN(torch.nn.Module):
-            def __init__(
-                self, in_features: int, out_features: int, kernel_size: int
-            ) -> None:
-                super().__init__()
-                self.conv1d = nn.Conv1d(in_features, out_features, kernel_size)
-                self.bn = nn.BatchNorm1d(out_features)
-
-            def forward(self, x: torch.Tensor) -> torch.Tensor:
-                y = self.conv1d(x)
-                return self.bn(y)
-
-        model = ModelConvBN(64, 1, 2)
-        x = torch.randn(1, 64, 4)
-
-        graph_module = (
-            compiler.export_to_executorch_gen_etrecord(model.eval(), (x,))
-            .exported_program()
-            .graph_module
-        )
-        # Assert that after running the fusion passes, batchnorm was fused with conv1d
-        self.assertEqual(
-            count_node(graph_module, torch.ops.aten.linear.out)
-            + count_node(graph_module, torch.ops.cadence.convolution.out),
-            1,
-        )
-        self.assertEqual(
-            count_node(
-                graph_module, torch.ops.aten._native_batch_norm_legit_no_training.out
-            ),
-            0,
-        )
-
     def test_permute_transpose_fusion(self) -> None:
         builder = GraphBuilder()
         x = builder.placeholder("x", torch.randn(3, 1, 3, 1, 4, dtype=torch.float32))