Update fused quant broadcast logic (#20171)

Summary:

Unifies QuantParamsStruct (sas_compiler's central quant-params abstraction) onto a single affine-quantization representation and drops the axis argument from every fused-quant op interface.

Core change (ops.py): scale/zero_point are now either a singleton (per-tensor, auto-expanded internally) or a full-rank tensor whose shape encodes the affine block layout — block_size[i] = tensor.shape[i] // scale.shape[i]. This one representation covers per-tensor, per-channel, per-group, and blockwise uniformly. quantize/dequantize delegate to torch.ops.torchao.(de)quantize_affine. The axis field is removed from QuantParamsStruct, all ~60 op-schema fields, _make_qp, and the _lib.define strings. is_per_tensor/is_per_channel/is_per_group and a new channel_axis() helper are now derived from scale shape (channel_axis() returns 0 if all dims are unary, the single non-unary dim if there's
exactly one, else None).

Fusion (fusion_pass.py, fusion_passes/utils.py): the qparams flat block is 6→5 tuple; the per-channel branch inserts an aten.view to make 1-D scales full-rank [1, …, C, …, 1] so their shape encodes the block layout.

Lowering boundary (graph_utils.py, lower_to_turing_linear.py, lower_to_turing_conv_no_nlu_params.py): Helios ParameterExtraction wants the compact scale form ([K] / [K, num_groups]), but the fused op now carries full-rank scales. New compact_scale_node() squeezes size-1 dims at the lowering boundary; the inserted view_copy folds via ConstantPropPass before extraction. Lowering asserts channel_axis() is not None.

Broadcast fix (fuse_mul_into_linear.py): channel_scale is an activation-space [K] vector (out-features is the trailing dim of the mul constant), whereas the weight scale is now full-rank [K, 1] (out-features at dim 0). Reshape the channel factor to [-1, 1, …] so it broadcasts along the weight scale's channel axis instead of producing a [K, K] outer product. The 1-D bias multiply is unchanged (bias is already [K]).

Misc consumers: quant_absorption.py per-tensor check is now out_scale.numel() == 1; BUCK adds the torchao dep.

Reviewed By: ethansfng

Differential Revision: D108065588

Author: DrJessop

backends/cadence/fused_quant/op_add.cpp

@@ -43,19 +43,16 @@ Tensor& add_out(
     ScalarType inp_dtype,
     int64_t inp_quant_min,
     int64_t inp_quant_max,
-    optional<int64_t> inp_axis,
     const optional<Tensor>& other_scale,
     const optional<Tensor>& other_zero_point,
     ScalarType other_dtype,
     int64_t other_quant_min,
     int64_t other_quant_max,
-    optional<int64_t> other_axis,
     const optional<Tensor>& out_scale,
     const optional<Tensor>& out_zero_point,
     ScalarType out_dtype,
     int64_t out_quant_min,
     int64_t out_quant_max,
-    optional<int64_t> out_axis,
     double alpha,
     Tensor& out) {
   int64_t numel = inp.numel();
@@ -72,7 +69,7 @@ Tensor& add_out(
     }
     inp_buf.resize(numel);
     QParams qp = extract_qparams(
-        inp_scale, inp_zero_point, inp_quant_min, inp_quant_max, inp_axis, inp);
+        inp_scale, inp_zero_point, inp_quant_min, inp_quant_max, inp);
     FUSED_QUANT_DTYPE_SWITCH(
         inp.scalar_type(),
         scalar_t,
@@ -88,12 +85,7 @@ Tensor& add_out(
     }
     other_buf.resize(numel);
     QParams qp = extract_qparams(
-        other_scale,
-        other_zero_point,
-        other_quant_min,
-        other_quant_max,
-        other_axis,
-        other);
+        other_scale, other_zero_point, other_quant_min, other_quant_max, other);
     FUSED_QUANT_DTYPE_SWITCH(
         other.scalar_type(),
         scalar_t,
@@ -107,7 +99,7 @@ Tensor& add_out(
     add_kernel(inp_float, other_float, result_float.data(), numel, alpha_f);
 
     QParams qp = extract_qparams(
-        out_scale, out_zero_point, out_quant_min, out_quant_max, out_axis, out);
+        out_scale, out_zero_point, out_quant_min, out_quant_max, out);
     FUSED_QUANT_DTYPE_SWITCH(
         out.scalar_type(),
         scalar_t,

backends/cadence/fused_quant/op_add.h

@@ -24,20 +24,17 @@ executorch::aten::Tensor& add_out(
     executorch::aten::ScalarType inp_dtype,
     int64_t inp_quant_min,
     int64_t inp_quant_max,
-    executorch::aten::optional<int64_t> inp_axis,
     const executorch::aten::optional<executorch::aten::Tensor>& other_scale,
     const executorch::aten::optional<executorch::aten::Tensor>&
         other_zero_point,
     executorch::aten::ScalarType other_dtype,
     int64_t other_quant_min,
     int64_t other_quant_max,
-    executorch::aten::optional<int64_t> other_axis,
     const executorch::aten::optional<executorch::aten::Tensor>& out_scale,
     const executorch::aten::optional<executorch::aten::Tensor>& out_zero_point,
     executorch::aten::ScalarType out_dtype,
     int64_t out_quant_min,
     int64_t out_quant_max,
-    executorch::aten::optional<int64_t> out_axis,
     double alpha,
     executorch::aten::Tensor& out);
 

backends/cadence/fused_quant/op_bmm.cpp

@@ -53,19 +53,16 @@ Tensor& bmm_out(
     ScalarType inp_dtype,
     int64_t inp_quant_min,
     int64_t inp_quant_max,
-    optional<int64_t> inp_axis,
     const optional<Tensor>& other_scale,
     const optional<Tensor>& other_zero_point,
     ScalarType other_dtype,
     int64_t other_quant_min,
     int64_t other_quant_max,
-    optional<int64_t> other_axis,
     const optional<Tensor>& out_scale,
     const optional<Tensor>& out_zero_point,
     ScalarType out_dtype,
     int64_t out_quant_min,
     int64_t out_quant_max,
-    optional<int64_t> out_axis,
     Tensor& out) {
   int64_t batch = inp.size(0);
   int64_t M = inp.size(1);
@@ -87,7 +84,7 @@ Tensor& bmm_out(
     }
     inp_buf.resize(inp_numel);
     QParams qp = extract_qparams(
-        inp_scale, inp_zero_point, inp_quant_min, inp_quant_max, inp_axis, inp);
+        inp_scale, inp_zero_point, inp_quant_min, inp_quant_max, inp);
     FUSED_QUANT_DTYPE_SWITCH(
         inp.scalar_type(),
         scalar_t,
@@ -104,12 +101,7 @@ Tensor& bmm_out(
     }
     other_buf.resize(other_numel);
     QParams qp = extract_qparams(
-        other_scale,
-        other_zero_point,
-        other_quant_min,
-        other_quant_max,
-        other_axis,
-        other);
+        other_scale, other_zero_point, other_quant_min, other_quant_max, other);
     FUSED_QUANT_DTYPE_SWITCH(other.scalar_type(),
                              scalar_t,
                              dequantize_buffer(
@@ -126,7 +118,7 @@ Tensor& bmm_out(
     bmm_kernel(inp_float, other_float, result_float.data(), batch, M, K, N);
 
     QParams qp = extract_qparams(
-        out_scale, out_zero_point, out_quant_min, out_quant_max, out_axis, out);
+        out_scale, out_zero_point, out_quant_min, out_quant_max, out);
     FUSED_QUANT_DTYPE_SWITCH(out.scalar_type(),
                              scalar_t,
                              quantize_buffer(

backends/cadence/fused_quant/op_bmm.h

@@ -24,20 +24,17 @@ executorch::aten::Tensor& bmm_out(
     executorch::aten::ScalarType inp_dtype,
     int64_t inp_quant_min,
     int64_t inp_quant_max,
-    executorch::aten::optional<int64_t> inp_axis,
     const executorch::aten::optional<executorch::aten::Tensor>& other_scale,
     const executorch::aten::optional<executorch::aten::Tensor>&
         other_zero_point,
     executorch::aten::ScalarType other_dtype,
     int64_t other_quant_min,
     int64_t other_quant_max,
-    executorch::aten::optional<int64_t> other_axis,
     const executorch::aten::optional<executorch::aten::Tensor>& out_scale,
     const executorch::aten::optional<executorch::aten::Tensor>& out_zero_point,
     executorch::aten::ScalarType out_dtype,
     int64_t out_quant_min,
     int64_t out_quant_max,
-    executorch::aten::optional<int64_t> out_axis,
     executorch::aten::Tensor& out);
 
 } // namespace native

backends/cadence/fused_quant/op_hardswish.cpp

@@ -40,13 +40,11 @@ Tensor& hardswish_out(
     ScalarType inp_dtype,
     int64_t inp_quant_min,
     int64_t inp_quant_max,
-    optional<int64_t> inp_axis,
     const optional<Tensor>& out_scale,
     const optional<Tensor>& out_zero_point,
     ScalarType out_dtype,
     int64_t out_quant_min,
     int64_t out_quant_max,
-    optional<int64_t> out_axis,
     Tensor& out) {
   int64_t numel = inp.numel();
 
@@ -60,7 +58,7 @@ Tensor& hardswish_out(
     }
     inp_buf.resize(numel);
     QParams qp = extract_qparams(
-        inp_scale, inp_zero_point, inp_quant_min, inp_quant_max, inp_axis, inp);
+        inp_scale, inp_zero_point, inp_quant_min, inp_quant_max, inp);
     FUSED_QUANT_DTYPE_SWITCH(
         inp.scalar_type(),
         scalar_t,
@@ -74,7 +72,7 @@ Tensor& hardswish_out(
     hardswish_kernel(inp_float, result_float.data(), numel);
 
     QParams qp = extract_qparams(
-        out_scale, out_zero_point, out_quant_min, out_quant_max, out_axis, out);
+        out_scale, out_zero_point, out_quant_min, out_quant_max, out);
     FUSED_QUANT_DTYPE_SWITCH(
         out.scalar_type(),
         scalar_t,

backends/cadence/fused_quant/op_hardswish.h

@@ -23,13 +23,11 @@ executorch::aten::Tensor& hardswish_out(
     executorch::aten::ScalarType inp_dtype,
     int64_t inp_quant_min,
     int64_t inp_quant_max,
-    executorch::aten::optional<int64_t> inp_axis,
     const executorch::aten::optional<executorch::aten::Tensor>& out_scale,
     const executorch::aten::optional<executorch::aten::Tensor>& out_zero_point,
     executorch::aten::ScalarType out_dtype,
     int64_t out_quant_min,
     int64_t out_quant_max,
-    executorch::aten::optional<int64_t> out_axis,
     executorch::aten::Tensor& out);
 
 } // namespace native

backends/cadence/fused_quant/op_mul.cpp

@@ -42,19 +42,16 @@ Tensor& mul_out(
     ScalarType inp_dtype,
     int64_t inp_quant_min,
     int64_t inp_quant_max,
-    optional<int64_t> inp_axis,
     const optional<Tensor>& other_scale,
     const optional<Tensor>& other_zero_point,
     ScalarType other_dtype,
     int64_t other_quant_min,
     int64_t other_quant_max,
-    optional<int64_t> other_axis,
     const optional<Tensor>& out_scale,
     const optional<Tensor>& out_zero_point,
     ScalarType out_dtype,
     int64_t out_quant_min,
     int64_t out_quant_max,
-    optional<int64_t> out_axis,
     Tensor& out) {
   (void)ctx;
   (void)inp_dtype;
@@ -74,7 +71,7 @@ Tensor& mul_out(
     }
     inp_buf.resize(numel);
     QParams qp = extract_qparams(
-        inp_scale, inp_zero_point, inp_quant_min, inp_quant_max, inp_axis, inp);
+        inp_scale, inp_zero_point, inp_quant_min, inp_quant_max, inp);
     FUSED_QUANT_DTYPE_SWITCH(
         inp.scalar_type(),
         scalar_t,
@@ -90,12 +87,7 @@ Tensor& mul_out(
     }
     other_buf.resize(numel);
     QParams qp = extract_qparams(
-        other_scale,
-        other_zero_point,
-        other_quant_min,
-        other_quant_max,
-        other_axis,
-        other);
+        other_scale, other_zero_point, other_quant_min, other_quant_max, other);
     FUSED_QUANT_DTYPE_SWITCH(
         other.scalar_type(),
         scalar_t,
@@ -109,7 +101,7 @@ Tensor& mul_out(
     mul_kernel(inp_float, other_float, result_float.data(), numel);
 
     QParams qp = extract_qparams(
-        out_scale, out_zero_point, out_quant_min, out_quant_max, out_axis, out);
+        out_scale, out_zero_point, out_quant_min, out_quant_max, out);
     FUSED_QUANT_DTYPE_SWITCH(
         out.scalar_type(),
         scalar_t,

backends/cadence/fused_quant/op_mul.h

@@ -24,20 +24,17 @@ executorch::aten::Tensor& mul_out(
     executorch::aten::ScalarType inp_dtype,
     int64_t inp_quant_min,
     int64_t inp_quant_max,
-    executorch::aten::optional<int64_t> inp_axis,
     const executorch::aten::optional<executorch::aten::Tensor>& other_scale,
     const executorch::aten::optional<executorch::aten::Tensor>&
         other_zero_point,
     executorch::aten::ScalarType other_dtype,
     int64_t other_quant_min,
     int64_t other_quant_max,
-    executorch::aten::optional<int64_t> other_axis,
     const executorch::aten::optional<executorch::aten::Tensor>& out_scale,
     const executorch::aten::optional<executorch::aten::Tensor>& out_zero_point,
     executorch::aten::ScalarType out_dtype,
     int64_t out_quant_min,
     int64_t out_quant_max,
-    executorch::aten::optional<int64_t> out_axis,
     executorch::aten::Tensor& out);
 
 } // namespace native

backends/cadence/fused_quant/op_relu.cpp

@@ -39,13 +39,11 @@ Tensor& relu_out(
     ScalarType inp_dtype,
     int64_t inp_quant_min,
     int64_t inp_quant_max,
-    optional<int64_t> inp_axis,
     const optional<Tensor>& out_scale,
     const optional<Tensor>& out_zero_point,
     ScalarType out_dtype,
     int64_t out_quant_min,
     int64_t out_quant_max,
-    optional<int64_t> out_axis,
     Tensor& out) {
   int64_t numel = inp.numel();
 
@@ -59,7 +57,7 @@ Tensor& relu_out(
     }
     inp_buf.resize(numel);
     QParams qp = extract_qparams(
-        inp_scale, inp_zero_point, inp_quant_min, inp_quant_max, inp_axis, inp);
+        inp_scale, inp_zero_point, inp_quant_min, inp_quant_max, inp);
     FUSED_QUANT_DTYPE_SWITCH(
         inp.scalar_type(),
         scalar_t,
@@ -73,7 +71,7 @@ Tensor& relu_out(
     relu_kernel(inp_float, result_float.data(), numel);
 
     QParams qp = extract_qparams(
-        out_scale, out_zero_point, out_quant_min, out_quant_max, out_axis, out);
+        out_scale, out_zero_point, out_quant_min, out_quant_max, out);
     FUSED_QUANT_DTYPE_SWITCH(
         out.scalar_type(),
         scalar_t,

backends/cadence/fused_quant/op_relu.h

@@ -23,13 +23,11 @@ executorch::aten::Tensor& relu_out(
     executorch::aten::ScalarType inp_dtype,
     int64_t inp_quant_min,
     int64_t inp_quant_max,
-    executorch::aten::optional<int64_t> inp_axis,
     const executorch::aten::optional<executorch::aten::Tensor>& out_scale,
     const executorch::aten::optional<executorch::aten::Tensor>& out_zero_point,
     executorch::aten::ScalarType out_dtype,
     int64_t out_quant_min,
     int64_t out_quant_max,
-    executorch::aten::optional<int64_t> out_axis,
     executorch::aten::Tensor& out);
 
 } // namespace native