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

CMakeLists.txt

@@ -48,8 +48,6 @@
 cmake_minimum_required(VERSION 3.24)
 project(executorch)
 
-# MARK: - Start EXECUTORCH_H12025_BUILD_MIGRATION
-
 include(${PROJECT_SOURCE_DIR}/tools/cmake/common/preset.cmake)
 include(${PROJECT_SOURCE_DIR}/tools/cmake/Utils.cmake)
 include(CMakeDependentOption)
@@ -82,6 +80,7 @@ announce_configured_options(BUCK2)
 
 announce_configured_options(CMAKE_CXX_COMPILER_ID)
 announce_configured_options(CMAKE_TOOLCHAIN_FILE)
+announce_configured_options(BUILD_TESTING)
 
 load_build_preset()
 include(${PROJECT_SOURCE_DIR}/tools/cmake/preset/default.cmake)
@@ -97,11 +96,6 @@ else()
 endif()
 announce_configured_options(CCACHE_PROGRAM)
 
-# Print all the configs that were called with announce_configured_options.
-print_configured_options()
-
-# MARK: - End EXECUTORCH_H12025_BUILD_MIGRATION
-
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 # Setup RPATH. See
@@ -750,3 +744,6 @@ if(EXECUTORCH_BUILD_ANDROID_JNI)
 endif()
 
 include(Test.cmake)
+
+# Print all the configs that were called with announce_configured_options.
+print_configured_options()

CMakePresets.json

@@ -8,7 +8,7 @@
     },
     {
       "name": "macos",
-      "displayName": "Build everything buildable on macOS",
+      "displayName": "Build ExecuTorch for macOS",
       "inherits": ["common"],
       "generator": "Xcode",
       "cacheVariables": {
@@ -25,7 +25,7 @@
     },
     {
       "name": "ios",
-      "displayName": "Build everything buildable on iOS",
+      "displayName": "Build ExecuTorch for iOS",
       "inherits": ["common"],
       "generator": "Xcode",
       "cacheVariables": {
@@ -42,7 +42,7 @@
     },
     {
       "name": "ios-simulator",
-      "displayName": "Build everything buildable on iOS simulator",
+      "displayName": "Build ExecuTorch for iOS Simulator",
       "inherits": ["common"],
       "generator": "Xcode",
       "cacheVariables": {
@@ -59,7 +59,7 @@
     },
     {
       "name": "linux",
-      "displayName": "Build everything buildable on Linux",
+      "displayName": "Build ExecuTorch for Linux",
       "inherits": ["common"],
       "cacheVariables": {
         "CMAKE_SYSTEM_NAME": "Linux",
@@ -88,29 +88,21 @@
     {
         "name": "llm",
         "displayName": "Build LLM libraries",
-        "inherits": [
-            "common"
-        ],
+        "inherits": ["common"],
         "cacheVariables": {
             "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/llm.cmake",
             "CMAKE_OSX_DEPLOYMENT_TARGET": "12.0"
         },
         "condition": {
             "type": "inList",
             "string": "${hostSystemName}",
-            "list": [
-                "Darwin",
-                "Linux",
-                "Windows"
-            ]
+            "list": ["Darwin", "Linux", "Windows"]
         }
     },
     {
         "name": "zephyr",
-        "displayName": "Build everything buildable on Zephyr RTOS",
-        "inherits": [
-            "common"
-        ],
+        "displayName": "Build ExecuTorch for Zephyr RTOS",
+        "inherits": ["common"],
         "cacheVariables": {
             "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/zephyr.cmake",
             "CMAKE_TOOLCHAIN_FILE": "${sourceDir}/examples/zephyr/x86_64-linux-arm-zephyr-eabi-gcc.cmake"

backends/nxp/nxp_backend.py

@@ -174,7 +174,8 @@ def preprocess(
             # Otherwise, we get violation that this op is not part of ATen Core ops.
             edge_program._verifiers = [
                 EXIREdgeDialectVerifier(
-                    class_only=True, core_aten_ops_exception_list=[torch.ops.aten.max_pool2d.default]
+                    class_only=True,
+                    core_aten_ops_exception_list=[torch.ops.aten.max_pool2d.default],
                 )
             ]
 

backends/qualcomm/qnn_preprocess.py

@@ -178,6 +178,11 @@ def preprocess_multimethod(
 
             if len(py_op_wrapper_list) == len(edge_programs.values()):
                 qnn_context_binary = qnn_manager.Compile(graph_name, py_op_wrapper_list)
+                if option.saver:
+                    # TODO: Currently, only the first method is saved. Update this logic if saving multiple methods becomes necessary in the future.
+                    exit(
+                        f"Record all QNN API calls from saver backend at: {option.saver_output_dir}"
+                    )
                 assert (
                     len(qnn_context_binary) != 0
                 ), "Failed to generate Qnn context binary."

backends/qualcomm/tests/test_qnn_delegate.py

@@ -3384,6 +3384,38 @@ def test_qnn_backend_rewrite_prepared_observer(self):
         quantized_module = convert_pt2e(prepared)
         self.lower_module_and_test_output(quantized_module, sample_input)
 
+    def test_qnn_backend_saver_backend(self):
+        backend_options = generate_htp_compiler_spec(use_fp16=False)
+        TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec(
+            soc_model=self.chipset_table[TestQNN.model],
+            backend_options=backend_options,
+            saver=True,
+        )
+        module = Relu()  # noqa: F405
+        sample_input = (torch.randn([2, 5, 1, 3]),)
+        module = self.get_qdq_module(module, sample_input)
+
+        from executorch.backends.qualcomm.serialization.qc_schema_serialize import (
+            flatbuffer_to_option,
+            option_to_flatbuffer,
+        )
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            option = flatbuffer_to_option(TestQNN.compiler_specs[0].value)
+            option.saver_output_dir = f"{tmp_dir}/saver_output"
+            TestQNN.compiler_specs[0].value = option_to_flatbuffer(option)
+
+            with self.assertRaises(SystemExit):
+                self.lower_module_and_test_output(module, sample_input)
+            self.assertTrue(
+                os.path.isfile(f"{tmp_dir}/saver_output/params.bin"),
+                "failed to find params.bin",
+            )
+            self.assertTrue(
+                os.path.isfile(f"{tmp_dir}/saver_output/saver_output.c"),
+                "failed to find saver_output.c",
+            )
+
     def test_qnn_backend_skip_node_id_partitioner(self):
         module = SimpleModel()  # noqa: F405
         sample_input = (torch.ones(1, 32, 28, 28), torch.ones(1, 32, 28, 28))
@@ -5022,6 +5054,40 @@ def test_swin_transformer(self):
                 self.assertGreaterEqual(msg["top_1"], 60)
                 self.assertGreaterEqual(msg["top_5"], 80)
 
+    def test_t5(self):
+        if not self.required_envs([self.qa_dataset]):
+            self.skipTest("missing required envs")
+        cmds = [
+            "python",
+            f"{self.executorch_root}/examples/qualcomm/oss_scripts/t5/t5.py",
+            "--dataset",
+            self.sentence_dataset,
+            "--artifact",
+            self.artifact_dir,
+            "--build_folder",
+            self.build_folder,
+            "--device",
+            self.device,
+            "--model",
+            self.model,
+            "--ip",
+            self.ip,
+            "--port",
+            str(self.port),
+        ]
+        if self.host:
+            cmds.extend(["--host", self.host])
+
+        p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL)
+        with Listener((self.ip, self.port)) as listener:
+            conn = listener.accept()
+            p.communicate()
+            msg = json.loads(conn.recv())
+            if "Error" in msg:
+                self.fail(msg["Error"])
+            else:
+                self.assertGreaterEqual(msg["f1"], 0.7)
+
     def test_whisper(self):
         if not self.required_envs():
             self.skipTest("missing required envs")

backends/qualcomm/tests/utils.py

@@ -183,6 +183,7 @@ class TestQNN(unittest.TestCase):
     executorch_root: str = ""
     artifact_dir: str = ""
     image_dataset: str = ""
+    qa_dataset: str = ""
     sentence_dataset: str = ""
     pretrained_weight: str = ""
     enable_profile: bool = False

backends/vulkan/op_registry.py

@@ -693,6 +693,7 @@ def register_transfer_ops(features: OpFeatures):
         exir_ops.edge.aten.full_like.default,
         exir_ops.edge.aten.ones.default,
         exir_ops.edge.aten.ones_like.default,
+        exir_ops.edge.aten.scalar_tensor.default,
         exir_ops.edge.aten.upsample_nearest2d.vec,
         exir_ops.edge.aten.upsample_bilinear2d.vec,
         exir_ops.edge.aten.zeros.default,

backends/vulkan/quantizer/TARGETS

@@ -3,24 +3,18 @@ load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
 oncall("executorch")
 
 python_library(
-    name = "vulkan_quantizer_utils",
-    srcs = [
-        "vulkan_quantizer_utils.py",
-    ],
+    name = "vulkan_quantizer",
+    srcs = ["vulkan_quantizer.py"],
     deps = [
+        ":vulkan_quantizer_utils",
         "//caffe2:torch",
-        "//pytorch/ao:torchao",  # @manual
     ],
 )
 
 python_library(
-    name = "vulkan_quantizer",
-    srcs = [
-        "vulkan_quantizer.py",
-    ],
+    name = "vulkan_quantizer_utils",
+    srcs = ["vulkan_quantizer_utils.py"],
     deps = [
-        ":vulkan_quantizer_utils",
         "//caffe2:torch",
-        "//pytorch/ao:torchao",  # @manual
     ],
 )

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -273,6 +273,14 @@ vkapi::ScalarType ComputeGraph::dtype_of(const ValueRef idx) const {
     return val.toConstTensor().dtype();
   } else if (val.isTensorRef()) {
     return val.toConstTensorRef().dtype;
+  } else if (val.isBool()) {
+    return vkapi::ScalarType::Bool;
+  } else if (val.isDouble()) {
+    // We downcast anyway in the shader and we want to avoid having to
+    // write special cases there.
+    return vkapi::ScalarType::Float;
+  } else if (val.isInt()) {
+    return vkapi::ScalarType::Int;
   }
   VK_THROW("Could not get dtype of value with type ", val.type());
 }

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

@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#version 450 core
+
+#define PRECISION ${PRECISION}
+
+#define BUF_T ${buffer_scalar_type(DTYPE)}
+#define VEC4_T ${texel_type(DTYPE)}
+
+${define_active_storage_type(STORAGE)}
+${define_required_extensions(DTYPE)}
+${define_required_extensions(SCALAR_VALUE_TYPE)}
+
+#include "indexing_utils.h"
+
+layout(std430) buffer;
+
+${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)}
+${layout_declare_ubo(B, buffer_scalar_type(SCALAR_VALUE_TYPE), "scalar_value")}
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+#ifdef USING_BUFFER
+
+void main() {
+  const int i = int(gl_GlobalInvocationID.x);
+
+  if (i > 0) {
+    return;
+  }
+
+  t_out[i] = BUF_T(scalar_value);
+}
+
+# else // !USING_BUFFER
+
+void main() {
+  const ivec3 pos = ivec3(gl_GlobalInvocationID);
+
+  // Scalar tensor is a special case where the packed dim is always 1.
+  if (any(greaterThanEqual(pos, ivec3(1)))) {
+    return;
+  }
+
+  VEC4_T outtex = VEC4_T(scalar_value);
+  write_texel(t_out, pos, outtex);
+}
+
+#endif // !USING_BUFFER