Update on "[ET-VK][Ops] torchao.quantize_affine vulkan impl and shader and cleanup"

# Changes
* Implement `torchao.quantize_affine` operator in Vulkan backend with comprehensive texture and buffer storage support
* Add block-wise quantization mode in `quantize_texture.glsl` and `quantize_buffer.glsl` shaders for configurable tensor block quantization
* Introduce comprehensive test suite in `affine_test.cpp` with multi-dimensional tensor validation and reference implementation
* Extend quantization infrastructure in `Quantize.cpp` to handle affine transformations with configurable block sizes and quantization parameters

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

NOTE: I delegated the quantize_affine and future affine operators through a new custom test file denoted as `affine_test.cpp` as the other quantization testing framework was getting a little large, and it makes more sense to separate the namespace between torchao and quantized_decomposed. I believe the _decomposed namespace is getting phased out in favor of this affine operator so deprecation will be easier in the future.

# Motivation
The existing Vulkan quantization infrastructure lacked support for the `torchao.quantize_affine` operator, which is essential for enabling dynamic quantization efficiently. The `quantize_affine` operator provides flexible block-wise quantization that allows different scale and zero-point values for tensor blocks, enabling:

* **Block-wise Quantization**: Applies quantization parameters to configurable tensor blocks rather than entire tensors, improving quantization accuracy for heterogeneous data distributions
* **Affine Transformation**: Uses the formula `qvalue = clamp(round(value / scale) + zero_point, quant_min, quant_max)` for precise floating-point to integer mapping

# Operator Description
The `quantize_affine` operator converts floating-point tensor values to n-bit integer representations using pre-computed quantization parameters (scale and zero_point) applied to configurable tensor blocks. Block-wise quantization divides tensors into blocks and applies separate quantization parameters to each block, allowing fine-grained control over quantization precision.

The quantization formula is: `qvalue = clamp(round(value / scale) + zero_point, quant_min, quant_max)`

**Storage Requirements**: Scale and zero_point tensors must use buffer storage with width-packed layout. Input/output tensors support both buffer and texture storage with standard axis mapping.

# Block-wise Quantization Implementation
Block-wise quantization enables fine-grained quantization by dividing tensors into blocks and applying separate quantization parameters to each block. The implementation uses several key data structures computed in `Quantize.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 `tensor_size_whcn / block_size_vec`
* **`block_stride_vec`**: Pre-computed linear strides for block grid indexing `{1, #W, #W*#H, #W*#H*#C}` to enable efficient block ID calculation

The block coordinate calculation uses: `bcoord = tidx / blockSize` where `tidx` is the tensor coordinate in WHCN layout, then the linear block ID is computed as: `block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w`

# Shader Algorithm Overview

## Texture Storage Implementation (`quantize_texture.glsl`)

**Workgroup Configuration**:
- **Global WG Size**: Default sizing based on texture dimensions
- **Local WG Size**: Default with special handling for batch dimension quantization (Z dimension set to 1 for proper workgroup dispatching when `global_workgroup_size[2] > 1`)

**Block-wise Mode Algorithm**:
The shader processes 3D texture positions where each position represents a texel containing 4 width-packed components. For each texel at position `pos`, it calculates a base tensor index `base_tidx = ivec4(pos.x * 4, pos.y, pos.z, 0)` to account for width-packing.

For each of the 4 components in the texel, it computes the actual tensor coordinate: `tidx = ivec4(base_tidx.x + i, base_tidx.y, (foldedZ % C_total), (foldedZ / C_total))` where `foldedZ = pos.z` handles batch-channel folding in 4D tensors and `C_total = numBlocks.z * blockSize.z` represents the total channel dimension.

The block coordinate is calculated using integer division: `bcoord = tidx / blockSize`, then the linear block ID uses pre-computed strides: `block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w`.

Each component is quantized using its corresponding block's parameters: `qvalue = quantize_val(value, t_scale[block_id], t_zero_point[block_id])` and written to the output texel.

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

**Workgroup Configuration**:
- **Global WG Size**: Default sizing based on buffer element count
- **Local WG Size**: Default sizing without special constraints

**Block-wise Mode Algorithm**:
The shader processes linear buffer indices using `gl_GlobalInvocationID.x` as the output buffer index. It converts this to tensor coordinates using `bufi_to_tidx(out_bufi, t_out_strides, out_dim_order)` which handles the buffer-to-tensor index mapping with proper stride calculations.

For each element, it computes the block coordinate directly: `bcoord = out_tidx / blockSize` where `out_tidx` is the 4D tensor coordinate in WHCN layout. The linear block ID calculation uses the same pre-computed stride approach: `block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w`.

The element value is loaded using the corresponding input buffer index: `value = t_in[in_bufi]` where `in_bufi = tidx_to_bufi(out_tidx, t_in_strides)`. Quantization applies the block-specific parameters: `qvalue = quantize_val(value, t_scale[block_id], t_zero_point[block_id])`.

**Future Improvements**: Dynamic workgroup sizing based on block dimensions, there is likely a better method to making it better than what it is currently.

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

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))