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

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