[ET-VK][Ops] torchao.choose_qparams_affine vulkan impl and shader

Pull Request resolved: #12577

# 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.
ghstack-source-id: 299473615
@exported-using-ghexport

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

Author: morelos

backends/vulkan/op_registry.py

@@ -245,9 +245,9 @@ def register_ephemeral_op(features: OpFeatures):
 
 @update_features(
     [
-        exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
         exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
         exir_ops.edge.quantized_decomposed.quantize_per_tensor.tensor,
+        exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
         exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
         exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor,
         exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
@@ -276,14 +276,32 @@ def register_quantization_op(features: OpFeatures):
     [
         exir_ops.edge.torchao.quantize_affine.default,
         exir_ops.edge.torchao.dequantize_affine.default,
+    ]
+)
+def register_affine_quantization_op(features: OpFeatures):
+    features.texture_impl = TextureImplFeatures(
+        uses_axis_map=False,
+        valid_packed_dims={PackedDim.WIDTH},
+    )
+    features.buffer_impl = True
+    features.resize_fn = True
+    features.optimal_storage = VkStorageType.TEXTURE_3D
+    features.optimal_layout = VkMemoryLayout.TENSOR_WIDTH_PACKED
+    features.handles_own_prepacking = True
+
+    return features
+
+
+@update_features(
+    [
         exir_ops.edge.torchao.choose_qparams_affine.default,
     ]
 )
-def register_torchao_quantization_op(features: OpFeatures):
-    # TorchAO quantization operators - default to per-tensor behavior
-    # Same features as standard quantization ops
+def register_choose_qparams_affine_op(features: OpFeatures):
+    # Currently only created a rudimentary buffer implementation for choose_qparams_affine
+    # since the reduction logic for blocks in texture3d is not trivial to implement in vulkan.
     features.texture_impl = TextureImplFeatures(
-        uses_axis_map=True,
+        uses_axis_map=False,
         valid_packed_dims={
             PackedDim.WIDTH,
         },
@@ -292,37 +310,6 @@ def register_torchao_quantization_op(features: OpFeatures):
     features.resize_fn = True
     features.optimal_storage = VkStorageType.BUFFER
 
-    def check_torchao_quantization_node(node: torch.fx.Node) -> bool:
-        # Only per-tensor quantization is supported by the Vulkan backend.
-        if len(node.args) < 2:
-            return False
-
-        block_size = node.args[1]
-
-        if not isinstance(block_size, (list, tuple)):
-            return False
-
-        input_arg = node.args[0]
-        if not isinstance(input_arg, torch.fx.Node):
-            return False
-
-        input_tensor = input_arg.meta.get("val", None)
-        if not isinstance(input_tensor, FakeTensor):
-            return False
-
-        input_shape = list(input_tensor.shape)
-
-        if len(block_size) != len(input_shape):
-            return False
-
-        # Check if block_size matches input_shape exactly (per-tensor quantization)
-        for i in range(len(block_size)):
-            if block_size[i] != input_shape[i]:
-                return False
-
-        return True
-
-    features.check_node_fn = check_torchao_quantization_node
     return features
 
 

backends/vulkan/runtime/graph/ops/glsl/choose_qparams.glslh

@@ -9,59 +9,67 @@
 #ifndef CHOOSE_QPARAMS_GLSLH
 #define CHOOSE_QPARAMS_GLSLH
 
-// Calculate scale and zero point from min and max values
-void calculate_scale_and_zero_point(
-    float min_val,
-    float max_val,
-    int qmin,
-    int qmax,
-    float eps_threshold,
-    out float scale_val,
-    out int zero_point_val) {
-  // ensure we have zero included in our range
-  min_val = min(min_val, 0.0);
-  max_val = max(max_val, 0.0);
+// mapping_type : 0 = ASYM, 1 = SYM, 2 = SYM_NO_CLIP
+void calc_scale_zp(
+    float lo, float hi,
+    int qmin, int qmax,
+    int mapping_type,
+    float eps,
+    out float scale, out int zp) {
+  // Handle case where lo and hi are +/-INF (no valid values found)
+  if (isinf(lo) || isinf(hi)) {
+    lo = 0.0;
+    hi = 0.0;
+  }
 
-  scale_val = (max_val - min_val) / float(qmax - qmin);
+  float minv = min(lo, 0.0);
+  float maxv = max(hi, 0.0);
 
-  // Handle zero or very small scale
-  if (scale_val == 0.0 || isinf(1.0 / scale_val)) {
-    scale_val = 0.1;
-  }
+  if (mapping_type == 0) { // asymmetric
+    scale = (maxv - minv) / float(qmax - qmin);
+
+    // Handle zero or very small scale
+    if (scale == 0.0 || isinf(1.0/scale)) {
+      scale = eps;
+    }
 
-  // Cut off small scale using the provided eps threshold
-  if (scale_val < eps_threshold) {
-    float org_scale = scale_val;
-    scale_val = eps_threshold;
+    if (scale < eps) {
+      float org_scale = scale;
+      scale = eps;
 
-    // Adjust min and max based on new scale
-    if (min_val == 0.0) {
-      max_val = eps_threshold * float(qmax - qmin);
-    } else if (max_val == 0.0) {
-      min_val = -eps_threshold * float(qmax - qmin);
-    } else {
-      float amplifier = eps_threshold / org_scale;
-      min_val *= amplifier;
-      max_val *= amplifier;
+      // Adjust min and max based on new scale to maintain proper quantization range
+      if (minv == 0.0) {
+        maxv = eps * float(qmax - qmin);
+      } else if (maxv == 0.0) {
+        minv = -eps * float(qmax - qmin);
+      } else {
+        float amplifier = eps / org_scale;
+        minv *= amplifier;
+        maxv *= amplifier;
+      }
+    }
+
+    // Calculate zero_point (matching reference implementation)
+    float initial_zero_point = float(qmin) - round(minv / scale);
+    zp = int(clamp(initial_zero_point, float(qmin), float(qmax)));
+  } else { // symmetric -- centred
+    float scale_sym;
+    if (mapping_type == 1) { // SYM
+      float M = max(abs(minv), abs(maxv));
+      scale_sym = M / (float(qmax - qmin) * 0.5);
+    } else { // SYM_NO_CLIP
+      float smin = abs(minv) / max(abs(float(qmin)), 1.0); // Avoid division by zero
+      float smax = maxv / max(float(qmax), 1.0); // Avoid division by zero
+      scale_sym = max(smin, smax);
     }
-  }
 
-  // Calculate zero point
-  float zero_point_from_min = float(qmin) - min_val / scale_val;
-  float zero_point_from_max = float(qmax) - max_val / scale_val;
-  float zero_point_from_min_error = abs(float(qmin)) - abs(min_val / scale_val);
-  float zero_point_from_max_error = abs(float(qmax)) - abs(max_val / scale_val);
-  float initial_zero_point = zero_point_from_min_error < zero_point_from_max_error
-      ? zero_point_from_min
-      : zero_point_from_max;
+    // Handle zero or very small scale
+    if (scale_sym == 0.0 || isinf(1.0/scale_sym)) {
+      scale_sym = eps;
+    }
 
-  // Nudge zero point to integer
-  if (initial_zero_point < float(qmin)) {
-    zero_point_val = qmin;
-  } else if (initial_zero_point > float(qmax)) {
-    zero_point_val = qmax;
-  } else {
-    zero_point_val = int(round(initial_zero_point));
+    scale = max(scale_sym, eps);
+    zp = int((qmax + qmin + 1) >> 1); // mid-point – always fits
   }
 }