diff --git a/backends/cadence/aot/BUCK b/backends/cadence/aot/BUCK index bc339242953..f69b25d46a6 100644 --- a/backends/cadence/aot/BUCK +++ b/backends/cadence/aot/BUCK @@ -434,6 +434,7 @@ fbcode_target(_kind = python_unittest, "//executorch/backends/cadence/aot:compiler", "//executorch/backends/test:graph_builder", "//executorch/backends/cadence/aot:pass_utils", + "//executorch/backends/cadence/aot/quantizer:utils", "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", "//executorch/exir/passes:lib", @@ -645,6 +646,7 @@ fbcode_target(_kind = python_unittest, deps = [ ":typing_stubs", "//executorch/backends/cadence/aot:ops_registrations", + "//executorch/backends/cadence/aot/quantizer:utils", "//caffe2:torch", ] ) @@ -664,6 +666,21 @@ fbcode_target(_kind = python_unittest, ], ) +fbcode_target(_kind = python_unittest, + name = "test_pattern_utils", + srcs = [ + "tests/test_pattern_utils.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/test:graph_builder", + "//executorch/backends/test:program_builder", + "//executorch/backends/cadence/aot:compiler_funcs", + "//executorch/backends/cadence/aot:ops_registrations", + "//executorch/backends/cadence/aot/quantizer:quantizer", + ], +) + fbcode_target(_kind = python_unittest, name = "test_to_out_var_pass", srcs = [ diff --git a/backends/cadence/aot/compiler.py b/backends/cadence/aot/compiler.py index 69db75fe75d..4fd72b529ac 100644 --- a/backends/cadence/aot/compiler.py +++ b/backends/cadence/aot/compiler.py @@ -170,7 +170,7 @@ def apply_pre_edge_transform_passes( PassManager( [ FuseQATConvBN(converted_program), - QuantFusionPass(patterns), + QuantFusionPass(patterns, converted_program), ] )(converted_program.graph_module) diff --git a/backends/cadence/aot/compiler_funcs.py b/backends/cadence/aot/compiler_funcs.py index 2f8dbf33416..d1b0620070b 100644 --- a/backends/cadence/aot/compiler_funcs.py +++ b/backends/cadence/aot/compiler_funcs.py @@ -13,8 +13,12 @@ import torch +from executorch.backends.cadence.aot.quantizer.pattern_utils import ( + EXPORTED_PROGRAM_META_KEY, +) from executorch.backends.transforms.permute_pass_utils import get_arg from torch._inductor.decomposition import remove_decompositions +from torch.export.exported_program import ExportedProgram from torch.fx import GraphModule from torch.fx.passes.infra.pass_base import PassBase, PassResult from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, prepare_qat_pt2e @@ -714,26 +718,41 @@ class QuantFusionPass(PassBase): """ Iterates patterns, finds anchor ops in the converted graph, and calls pattern.fuse() to replace dq-op-q subgraphs with fused ops. + + ``exported_program`` is optional but required for per-channel weights: + fusion has to read the weight scale vector and materialize the derived + qparam tensors, and by this point both live in the program's constants + rather than in the graph. """ - def __init__(self, patterns: Sequence[object]) -> None: + def __init__( + self, + patterns: Sequence[object], + exported_program: Optional[ExportedProgram] = None, + ) -> None: super().__init__() self.patterns = patterns + self.exported_program = exported_program def call(self, graph_module: GraphModule) -> Optional[PassResult]: changed = False - for pattern in self.patterns: - pattern_changed = False - for target in pattern.anchor_ops(): # pyre-ignore[16] - for node in graph_module.graph.find_nodes( - op="call_function", target=target - ): - result = pattern.fuse(graph_module, node) # pyre-ignore[16] - if result is not None: - changed = True - pattern_changed = True - if pattern_changed: - graph_module.graph.eliminate_dead_code() + if self.exported_program is not None: + graph_module.meta[EXPORTED_PROGRAM_META_KEY] = self.exported_program + try: + for pattern in self.patterns: + pattern_changed = False + for target in pattern.anchor_ops(): # pyre-ignore[16] + for node in graph_module.graph.find_nodes( + op="call_function", target=target + ): + result = pattern.fuse(graph_module, node) # pyre-ignore[16] + if result is not None: + changed = True + pattern_changed = True + if pattern_changed: + graph_module.graph.eliminate_dead_code() + finally: + graph_module.meta.pop(EXPORTED_PROGRAM_META_KEY, None) if changed: graph_module.recompile() return PassResult(graph_module, changed) diff --git a/backends/cadence/aot/ops_registrations.py b/backends/cadence/aot/ops_registrations.py index da82a1ea3ec..e1b2dc84ae3 100644 --- a/backends/cadence/aot/ops_registrations.py +++ b/backends/cadence/aot/ops_registrations.py @@ -268,12 +268,24 @@ def register_fake( lib.define( "quantized_conv1d_nlc.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, Tensor? offset=None, *, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "quantized_depthwise_conv1d_ncl(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_depthwise_conv1d_ncl.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) lib.define( "quantized_depthwise_conv1d_ncl.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" ) lib.define( "quantized_depthwise_conv1d_ncl.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "quantized_depthwise_conv1d_nlc(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_depthwise_conv1d_nlc.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) lib.define( "quantized_depthwise_conv1d_nlc.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" ) @@ -1396,6 +1408,76 @@ def quantized_depthwise_conv1d_nlc_per_tensor_meta( return input.new_empty(output_size, dtype=input.dtype) +@register_fake("cadence::quantized_depthwise_conv1d_ncl") +def quantized_depthwise_conv1d_ncl_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: torch.Tensor, + bias_scale: torch.Tensor, + output_scale: float, + output_zero_point: int, + out_multiplier: torch.Tensor, + out_shift: torch.Tensor, +) -> torch.Tensor: + return quantized_depthwise_conv1d_ncl_per_tensor_meta( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + 0, + 1.0, + output_scale, + output_zero_point, + 1, + 0, + ) + + +@register_fake("cadence::quantized_depthwise_conv1d_nlc") +def quantized_depthwise_conv1d_nlc_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: torch.Tensor, + bias_scale: torch.Tensor, + output_scale: float, + output_zero_point: int, + out_multiplier: torch.Tensor, + out_shift: torch.Tensor, +) -> torch.Tensor: + return quantized_depthwise_conv1d_nlc_per_tensor_meta( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + 0, + 1.0, + output_scale, + output_zero_point, + 1, + 0, + ) + + @register_fake("cadence::quantized_conv2d_nchw") def quantized_conv2d_nchw_meta( input: torch.Tensor, diff --git a/backends/cadence/aot/quantizer/pattern_utils.py b/backends/cadence/aot/quantizer/pattern_utils.py index 25ff363ecc9..7a8953ba098 100644 --- a/backends/cadence/aot/quantizer/pattern_utils.py +++ b/backends/cadence/aot/quantizer/pattern_utils.py @@ -21,8 +21,197 @@ from torch._ops import OpOverload DQ_PER_TENSOR: OpOverload = torch.ops.quantized_decomposed.dequantize_per_tensor.default +DQ_PER_CHANNEL: OpOverload = ( + torch.ops.quantized_decomposed.dequantize_per_channel.default +) Q_PER_TENSOR: OpOverload = torch.ops.quantized_decomposed.quantize_per_tensor.default +# Fusion needs the ExportedProgram to read per-channel scale tensors and to +# materialize the derived qparam constants, but `PassBase.call` only receives a +# GraphModule and `QuantizationPattern.fuse` is a public API implemented by ~20 +# patterns. Rather than thread an extra argument through all of them, +# QuantFusionPass stashes the program here for the duration of the pass. +EXPORTED_PROGRAM_META_KEY: str = "_cadence_quant_fusion_exported_program" + + +def is_weight_dq(node: object) -> bool: + """True if ``node`` is a dequantize that a quantized op can absorb. + + Only use this to guard a pattern whose fusion helper handles per-channel + qparams, which today means the conv patterns (``fuse_conv``). ``fuse_linear``, + ``fuse_matmul`` and the mixed w8a32 paths read the scalar ``scale`` arg, which + a per-channel dequantize does not have, so they must keep checking for + ``DQ_PER_TENSOR`` and let per-channel weights fall back to float. + """ + return isinstance(node, fx.Node) and node.target in ( + DQ_PER_TENSOR, + DQ_PER_CHANNEL, + ) + + +def is_per_channel_dq(node: object) -> bool: + return isinstance(node, fx.Node) and node.target is DQ_PER_CHANNEL + + +def get_exported_program(gm: fx.GraphModule) -> Any: + """The ExportedProgram being fused, or None outside QuantFusionPass.""" + return gm.meta.get(EXPORTED_PROGRAM_META_KEY) + + +def resolve_constant(gm: fx.GraphModule, node: object) -> torch.Tensor | None: + """Read the tensor behind a placeholder or get_attr node.""" + if not isinstance(node, fx.Node): + return None + if node.op == "get_attr": + return getattr(gm, str(node.target), None) + ep = get_exported_program(gm) + if ep is None: + return None + # local import: torch._export.utils pulls in export internals + from torch._export.utils import ( + get_buffer, + get_lifted_tensor_constant, + get_param, + is_buffer, + is_lifted_tensor_constant, + is_param, + ) + + if is_param(ep, node): + return get_param(ep, node) + if is_buffer(ep, node): + return get_buffer(ep, node) + if is_lifted_tensor_constant(ep, node): + return get_lifted_tensor_constant(ep, node) + return None + + +def add_constant_placeholder( + gm: fx.GraphModule, + tensor: torch.Tensor, + like_node: fx.Node, + name_hint: str, +) -> fx.Node: + """Materialize ``tensor`` as a lifted constant placeholder. + + Per-channel fusion produces qparam vectors (bias_scale, out_multiplier, + out_shift, weight_zero_point) that have to reach the kernel as tensors. The + graph is already exported by this point, so they cannot be `get_attr` nodes: + they have to be registered in ``ExportedProgram.constants`` and appear as + placeholders with a matching entry in the graph signature. This mirrors + ``exir.passes.constant_prop_pass.replace_with_constant_node``. + """ + from torch.export.graph_signature import InputKind, InputSpec, TensorArgument + + ep = get_exported_program(gm) + assert ep is not None, "per-channel fusion requires the ExportedProgram" + + prefix = f"_cadence_{name_hint}_" + idx = 0 + while f"{prefix}{idx}" in ep.constants: + idx += 1 + fqn = f"{prefix}{idx}" + ep.constants[fqn] = tensor + + # Placeholder order and graph_signature.input_specs order must agree, and + # constants have to precede user inputs. Insert at the user-input boundary + # in both, the same way exir's constant_prop_pass does. + user_inputs = set(ep.graph_signature.user_inputs) + first_user_input = next( + ( + n + for n in ep.graph.nodes + if n.op == "placeholder" and n.name in user_inputs + ), + None, + ) + specs = ep.graph_signature.input_specs + if first_user_input is not None: + with ep.graph.inserting_before(first_user_input): + node = ep.graph.placeholder(fqn) + spec_idx = next( + ( + i + for i, s in enumerate(specs) + if getattr(s.arg, "name", None) == first_user_input.name + ), + len(specs), + ) + else: + placeholders = [n for n in ep.graph.nodes if n.op == "placeholder"] + anchor = placeholders[-1] if placeholders else next(iter(ep.graph.nodes)) + with ep.graph.inserting_after(anchor): + node = ep.graph.placeholder(fqn) + spec_idx = len(specs) + + fake_mode = like_node.meta["val"].fake_mode if "val" in like_node.meta else None + if fake_mode is not None: + node.meta["val"] = fake_mode.from_tensor(tensor, static_shapes=True) + node.meta["val"].constant = tensor + else: + node.meta["val"] = tensor + + specs.insert( + spec_idx, + InputSpec( + kind=InputKind.CONSTANT_TENSOR, + arg=TensorArgument(name=node.name), + target=fqn, + persistent=True, + ), + ) + return node + + +def tensor_qparam_overload(op: OpOverload) -> OpOverload: + """Map a `.per_tensor` Cadence op to its tensor-qparam `.default` sibling. + + Both overloads take the same operands; the scalar one inlines the qparams as + SymInt/float args, while the tensor one takes them as constant tensors. Only + the tensor form can express per-channel. + """ + name = op._schema.name.split("::")[-1] + packet = getattr(torch.ops.cadence, name, None) + assert packet is not None and hasattr(packet, "default"), ( + f"no tensor-qparam overload registered for {name}; per-channel needs one" + ) + return packet.default + + +def get_weight_scale( + gm: fx.GraphModule, dq_weight: fx.Node +) -> float | torch.Tensor: + """Weight scale as a float (per-tensor) or a per-output-channel vector.""" + if not is_per_channel_dq(dq_weight): + return get_arg(dq_weight, "scale", float) + scales = resolve_constant(gm, get_arg(dq_weight, "scales", fx.Node)) + assert scales is not None, ( + f"could not resolve per-channel weight scales for {dq_weight}" + ) + axis = get_arg(dq_weight, "axis", int) + assert axis == 0, ( + f"Cadence per-channel weights must be quantized on the output-channel " + f"axis (0), got axis={axis}" + ) + return scales.to(torch.float32).flatten() + + +def get_weight_zero_point( + gm: fx.GraphModule, dq_weight: fx.Node +) -> int | torch.Tensor: + """Weight zero point as an int (per-tensor) or a per-channel int32 vector.""" + if not is_per_channel_dq(dq_weight): + return get_arg(dq_weight, "zero_point", int) + # Not typed as fx.Node: symmetric per-channel leaves this argument unset, and + # get_arg raises on a type mismatch before we could check for None. + zps = resolve_constant(gm, get_arg(dq_weight, "zero_points")) + if zps is None: + # symmetric per-channel: zero_points is allowed to be absent + scales = get_weight_scale(gm, dq_weight) + assert isinstance(scales, torch.Tensor) + return torch.zeros_like(scales, dtype=torch.int32) + return zps.to(torch.int32).flatten() + def insert_node_with_meta( gm: fx.GraphModule, @@ -87,8 +276,9 @@ def fuse_conv( if len(conv_node.args) > 2 and conv_node.args[2] is not None: bias_arg = conv_node.args[2] assert isinstance(bias_arg, fx.Node) - dq_bias = bias_arg if bias_arg.target == DQ_PER_TENSOR else None - weight_scale = get_arg(dq_weight, "scale", float) + dq_bias = bias_arg if is_weight_dq(bias_arg) else None + per_channel = is_per_channel_dq(dq_weight) + weight_scale = get_weight_scale(gm, dq_weight) input_scale = get_arg(dq_input, "scale", float) bias_scale = input_scale * weight_scale if dq_bias is not None: @@ -97,9 +287,19 @@ def fuse_conv( # Cadence quantized conv ops require a non-optional bias argument. weight_node = get_arg(dq_weight, "input", fx.Node) with gm.graph.inserting_before(conv_node): - bias_q = create_zero_bias_int32(gm, weight_node, bias_scale) - requantize_scale = bias_scale / get_arg(quant_node, "scale", float) - requantize_scale_t = torch.tensor([requantize_scale]) + # the helper only needs a representative scalar, it fills with zeros + bias_q = create_zero_bias_int32( + gm, + weight_node, + float(bias_scale.max()) if per_channel else bias_scale, + ) + out_scale = get_arg(quant_node, "scale", float) + requantize_scale = bias_scale / out_scale + requantize_scale_t = ( + requantize_scale + if isinstance(requantize_scale, torch.Tensor) + else torch.tensor([requantize_scale]) + ) out_multiplier, out_shift = quantize_tensor_multiplier(requantize_scale_t) args = ( get_arg(dq_input, "input", fx.Node), @@ -113,12 +313,8 @@ def fuse_conv( "dilation": get_arg(conv_node, "dilation", list[int]), "groups": groups, "input_zero_point": get_arg(dq_input, "zero_point", int), - "weight_zero_point": get_arg(dq_weight, "zero_point", int), - "bias_scale": bias_scale, - "out_scale": get_arg(quant_node, "scale", float), + "out_scale": out_scale, "out_zero_point": get_arg(quant_node, "zero_point", int), - "out_multiplier": out_multiplier[0].item(), - "out_shift": out_shift[0].item(), } replacement_op = pattern.replacement_op() # pyre-ignore[16] if replacement_op == torch.ops.cadence.quantized_conv1d_ncl.per_tensor: @@ -127,6 +323,28 @@ def fuse_conv( in_channels = input_node.meta["val"].shape[1] if is_depthwise_conv(groups, in_channels): replacement_op = torch.ops.cadence.quantized_depthwise_conv1d_ncl.per_tensor + if per_channel: + # The tensor-qparam overload takes the same operands but carries the + # qparams as constant tensors, which is what per-channel needs. This runs + # after depthwise selection so that it swaps whichever base op was chosen. + replacement_op = tensor_qparam_overload(replacement_op) + kwargs["weight_zero_point"] = add_constant_placeholder( + gm, get_weight_zero_point(gm, dq_weight), conv_node, "wzp" + ) + kwargs["bias_scale"] = add_constant_placeholder( + gm, bias_scale.to(torch.float32), conv_node, "bias_scale" + ) + kwargs["out_multiplier"] = add_constant_placeholder( + gm, out_multiplier.to(torch.int32), conv_node, "out_multiplier" + ) + kwargs["out_shift"] = add_constant_placeholder( + gm, out_shift.to(torch.int32), conv_node, "out_shift" + ) + else: + kwargs["weight_zero_point"] = get_arg(dq_weight, "zero_point", int) + kwargs["bias_scale"] = bias_scale + kwargs["out_multiplier"] = out_multiplier[0].item() + kwargs["out_shift"] = out_shift[0].item() return replace_with_op(gm, conv_node, replacement_op, args, kwargs, quant_node) @@ -145,11 +363,16 @@ def fuse_linear( torch.ops.aten.linear.default, torch.ops.aten.addmm.default, ), f"Expected linear/addmm, got {op_node.target}" - weight_scale = get_arg(dq_weight, "scale", float) + per_channel = is_per_channel_dq(dq_weight) + weight_scale = get_weight_scale(gm, dq_weight) input_scale = get_arg(dq_input, "scale", float) bias_scale = input_scale * weight_scale requantize_scale = bias_scale / get_arg(quant_node, "scale", float) - requantize_scale_t = torch.tensor([requantize_scale]) + requantize_scale_t = ( + requantize_scale + if isinstance(requantize_scale, torch.Tensor) + else torch.tensor([requantize_scale]) + ) out_multiplier, out_shift = quantize_tensor_multiplier(requantize_scale_t) if dq_bias is not None: bias_q = get_arg(dq_bias, "input", fx.Node) @@ -157,19 +380,39 @@ def fuse_linear( # Cadence quantized linear ops require a non-optional bias argument. weight_node = get_arg(dq_weight, "input", fx.Node) with gm.graph.inserting_before(op_node): - bias_q = create_zero_bias_int32(gm, weight_node, bias_scale) + # the helper only needs a representative scalar, it fills with zeros + bias_q = create_zero_bias_int32( + gm, + weight_node, + float(bias_scale.max()) if per_channel else bias_scale, + ) final_weight = ( weight_q if weight_q is not None else get_arg(dq_weight, "input", fx.Node) ) args = (get_arg(dq_input, "input", fx.Node), final_weight, bias_q) kwargs = { "src_zero_point": get_arg(dq_input, "zero_point", int), - "weight_zero_point": get_arg(dq_weight, "zero_point", int), - "out_multiplier": out_multiplier[0].item(), - "out_shift": out_shift[0].item(), "out_zero_point": get_arg(quant_node, "zero_point", int), "offset": None, } + if per_channel: + # The tensor-qparam overload carries the qparams as constant tensors, + # which is what per-channel needs. quantized_linear has no bias_scale + # arg, so the scale only survives through out_multiplier/out_shift. + replacement_op = tensor_qparam_overload(replacement_op) + kwargs["weight_zero_point"] = add_constant_placeholder( + gm, get_weight_zero_point(gm, dq_weight), op_node, "wzp" + ) + kwargs["out_multiplier"] = add_constant_placeholder( + gm, out_multiplier.to(torch.int32), op_node, "out_multiplier" + ) + kwargs["out_shift"] = add_constant_placeholder( + gm, out_shift.to(torch.int32), op_node, "out_shift" + ) + else: + kwargs["weight_zero_point"] = get_arg(dq_weight, "zero_point", int) + kwargs["out_multiplier"] = out_multiplier[0].item() + kwargs["out_shift"] = out_shift[0].item() return replace_with_op(gm, op_node, replacement_op, args, kwargs, quant_node) diff --git a/backends/cadence/aot/quantizer/patterns.py b/backends/cadence/aot/quantizer/patterns.py index e3dc7afd0cf..b4f7c420b50 100644 --- a/backends/cadence/aot/quantizer/patterns.py +++ b/backends/cadence/aot/quantizer/patterns.py @@ -21,6 +21,7 @@ fuse_linear, fuse_matmul, insert_node_with_meta, + is_weight_dq, ) from executorch.backends.cadence.aot.quantizer.utils import ( check_out_zero_point_is_min_range, @@ -492,7 +493,7 @@ def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None: if not isinstance(dq_input, fx.Node) or dq_input.target != DQ_PER_TENSOR: return None dq_weight = anchor_node.args[1] - if not isinstance(dq_weight, fx.Node) or dq_weight.target != DQ_PER_TENSOR: + if not is_weight_dq(dq_weight): return None quant_node = find_quant_user(anchor_node) if quant_node is None: @@ -546,7 +547,7 @@ def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None: if not isinstance(dq_input, fx.Node) or dq_input.target != DQ_PER_TENSOR: return None dq_weight = anchor_node.args[1] - if not isinstance(dq_weight, fx.Node) or dq_weight.target != DQ_PER_TENSOR: + if not is_weight_dq(dq_weight): return None quant_node = find_quant_user(anchor_node) if quant_node is None: @@ -693,7 +694,7 @@ def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None: if not isinstance(dq_input, fx.Node) or dq_input.target != DQ_PER_TENSOR: return None dq_weight = anchor_node.args[1] - if not isinstance(dq_weight, fx.Node) or dq_weight.target != DQ_PER_TENSOR: + if not is_weight_dq(dq_weight): return None quant_node = find_quant_user(anchor_node) if quant_node is None: @@ -1008,11 +1009,7 @@ def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None: else None ) _arg1 = anchor_node.args[1] - dq_weight = ( - _arg1 - if isinstance(_arg1, fx.Node) and _arg1.target == DQ_PER_TENSOR - else None - ) + dq_weight = _arg1 if is_weight_dq(_arg1) else None if dq_input is None or dq_weight is None: return None quant_node = find_quant_user(relu_node) diff --git a/backends/cadence/aot/quantizer/quantizer.py b/backends/cadence/aot/quantizer/quantizer.py index 2cf41ef8c6f..20e677cfc88 100644 --- a/backends/cadence/aot/quantizer/quantizer.py +++ b/backends/cadence/aot/quantizer/quantizer.py @@ -6,7 +6,7 @@ # pyre-strict -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import final, List, Optional, Tuple, Union import torch @@ -52,6 +52,7 @@ FusedMovingAvgObsFakeQuantize, HistogramObserver, MinMaxObserver, + PerChannelMinMaxObserver, ) from torchao.quantization.pt2e.quantizer import ( ComposableQuantizer, @@ -101,6 +102,16 @@ observer_or_fake_quant_ctr=MinMaxObserver, ) +wgt_qspec_sym8s_per_channel = QuantizationSpec( + dtype=torch.int8, + quant_min=-128, + quant_max=127, + qscheme=torch.per_channel_symmetric, + ch_axis=0, + is_dynamic=False, + observer_or_fake_quant_ctr=PerChannelMinMaxObserver, +) + bias_qspec: Optional[QuantizationSpec] = None qconfig_A8W8 = QuantizationConfig( @@ -158,6 +169,13 @@ None, ) +qconfig_A8W8sym_per_channel = QuantizationConfig( + act_qspec_asym8s, + act_qspec_asym8s, + wgt_qspec_sym8s_per_channel, + None, +) + qconfig_A16 = QuantizationConfig( act_qspec_asym16s, act_qspec_asym16s, @@ -173,6 +191,49 @@ ) +def _align_derived_bias_specs( + biases: List[Tuple[fx.Node, int]], + weight_qspec: Optional[QuantizationSpec], +) -> List[Tuple[fx.Node, int]]: + """Make derived bias specs follow the granularity of the weight spec. + + Patterns build the bias ``DerivedQuantizationSpec`` as + ``bias_scale = act_scale * weight_scale`` and declare it per-tensor, because + a pattern has no view of the quantization config and so cannot know how the + weight will be observed. As soon as the weight observer is per-channel that + product becomes a vector while the spec still says per-tensor, and + ``convert_pt2e`` takes the per-tensor branch and dies calling ``float()`` on + a vector. + + The quantizer is the one place that holds both, so the alignment happens + here rather than being duplicated into every pattern that has a bias. + """ + if weight_qspec is None or weight_qspec.qscheme not in ( + torch.per_channel_symmetric, + torch.per_channel_affine, + ): + return biases + + aligned = [] + for entry in biases: + if len(entry) == 3 and isinstance(entry[2], DerivedQuantizationSpec): + aligned.append( + ( + entry[0], + entry[1], + replace( + entry[2], + qscheme=weight_qspec.qscheme, + ch_axis=weight_qspec.ch_axis, + ), + ) + ) + else: + aligned.append(entry) + # pyre-ignore[7]: the 3-tuple form is what PartitionAnchors.biases allows + return aligned + + class CadenceAtenQuantizer(Quantizer): def __init__( self, pattern: QuantizationPattern, quantization_config: QuantizationConfig @@ -261,8 +322,12 @@ def annotate_weights_or_biases( # pyre-ignore[6]: incompatible parameter type annotate_inputs(anchors.inputs, input_act_qspec) annotate_weights_or_biases(anchors.weights, weight_qspec) - # pyre-ignore[6]: incompatible parameter type - annotate_weights_or_biases(anchors.biases, bias_qspec) + annotate_weights_or_biases( + # pyre-ignore[6]: incompatible parameter type + _align_derived_bias_specs(anchors.biases, weight_qspec), + # pyre-ignore[6]: incompatible parameter type + bias_qspec, + ) return model def validate(self, model: fx.GraphModule) -> None: diff --git a/backends/cadence/aot/ref_implementations.py b/backends/cadence/aot/ref_implementations.py index d3a5c853a4a..0a9cf435586 100644 --- a/backends/cadence/aot/ref_implementations.py +++ b/backends/cadence/aot/ref_implementations.py @@ -466,8 +466,8 @@ def quantized_linear_common( bias: torch.Tensor, in_zero_point: int, weight_zero_point: torch.Tensor | int, - out_multiplier: int, - out_shift: int, + out_multiplier: int | torch.Tensor, + out_shift: int | torch.Tensor, out_zero_point: int, ) -> torch.Tensor: """ @@ -479,13 +479,12 @@ def quantized_linear_common( - bias (Tensor): The bias tensor - in_zero_point (int): The quantized mapping of zero for the input - weight_zero_point (Tensor): The quantized mapping of zero for the weight - - out_multiplier (Tensor): The multiplier used to scale the output - - out_shift (Tensor): The shift used to scale the output + - out_multiplier (int | Tensor): The multiplier used to scale the output. + A 1-D tensor of length out_dim selects per-channel requantization. + - out_shift (int | Tensor): The shift used to scale the output - out_zero_point (int): The quantized mapping of zero for the output - offset (Tensor): Unused """ - out_scale = 1.0 / (-out_multiplier * (1 / (1 << 31)) * (2**out_shift)) - N, K = weight.shape leading_dims = src.shape[:-1] @@ -498,19 +497,67 @@ def quantized_linear_common( f"Unsupported dtype to quantize to {dtype}. Supported dtypes must be one of {supported_dtypes}" ) + per_channel = isinstance(out_multiplier, torch.Tensor) and ( + out_multiplier.numel() > 1 + ) + + if ( + per_channel + and isinstance(weight_zero_point, torch.Tensor) + and weight_zero_point.numel() == N + ): + # per-channel zero point indexes output channels, i.e. weight rows. + # Only reshape under per-channel requant: elsewhere a multi-element + # zero point is expected to broadcast along the input axis. + weight_zero_point = weight_zero_point.reshape(N, 1) + out = torch.nn.functional.linear( src.float() - in_zero_point, weight.float() - weight_zero_point, bias.float(), ) - return quantize_per_tensor( - out, - out_scale, - out_zero_point, - torch.iinfo(dtype).min, - torch.iinfo(dtype).max, - dtype, - ).reshape(*leading_dims, N) + + if not per_channel: + _multiplier = ( + int(out_multiplier.flatten()[0].item()) + if isinstance(out_multiplier, torch.Tensor) + else out_multiplier + ) + _shift = ( + int(out_shift.flatten()[0].item()) + if isinstance(out_shift, torch.Tensor) + else out_shift + ) + out_scale = 1.0 / (-_multiplier * (1 / (1 << 31)) * (2**_shift)) + return quantize_per_tensor( + out, + out_scale, + out_zero_point, + torch.iinfo(dtype).min, + torch.iinfo(dtype).max, + dtype, + ).reshape(*leading_dims, N) + + # Per-channel: one requant scale per output channel. Mirrors the generic + # kernel, which reconstructs a float scale from (multiplier, shift) per + # channel rather than doing integer fixed-point requantization. + assert isinstance(out_shift, torch.Tensor) + if out_multiplier.numel() != N or out_shift.numel() != N: + raise ValueError( + f"per-channel out_multiplier/out_shift must have {N} elements, got " + f"{out_multiplier.numel()}/{out_shift.numel()}" + ) + requant_scale = ( + -out_multiplier.to(torch.float64).flatten() + * (1 / (1 << 31)) + * torch.pow(2.0, out_shift.to(torch.float64).flatten()) + ) + quantized = torch.round(out.to(torch.float64) * requant_scale) + out_zero_point + return ( + quantized.clamp(torch.iinfo(dtype).min, torch.iinfo(dtype).max) + .to(dtype) + .reshape(*leading_dims, N) + ) def quantized_linear_variant( @@ -558,16 +605,22 @@ def variant( else: assert isinstance(out_shift, torch.Tensor) assert isinstance(out_multiplier, torch.Tensor) - if out_shift.numel() != 1: - raise ValueError("out_shift must be a scalar") - if out_shift.dtype != torch.int32: raise ValueError( f"out_shift must be an int32. Got {out_shift.dtype} instead" ) - - _out_shift = int(out_shift.item()) - _out_multiplier = int(out_multiplier[0].item()) + if out_shift.numel() != out_multiplier.numel(): + raise ValueError( + "out_shift and out_multiplier must have the same length, got " + f"{out_shift.numel()} and {out_multiplier.numel()}" + ) + if out_multiplier.numel() == 1: + _out_shift = int(out_shift.item()) + _out_multiplier = int(out_multiplier[0].item()) + else: + # per-channel: pass the vectors through + _out_shift = out_shift + _out_multiplier = out_multiplier return quantized_linear_common( src, @@ -817,7 +870,28 @@ def quantized_layer_norm( ) -def quantized_conv_per_tensor( +def _broadcast_over_channels( + value: float | int | torch.Tensor, + ndim: int, + channel_dim: int, + out_channels: int, +) -> float | int | torch.Tensor: + """Shape a scalar or per-output-channel qparam for broadcasting.""" + if not isinstance(value, torch.Tensor): + return value + if value.numel() == 1: + return value.reshape(()) + if value.numel() != out_channels: + raise ValueError( + f"per-channel qparam has {value.numel()} entries, " + f"expected 1 or {out_channels}" + ) + shape = [1] * ndim + shape[channel_dim] = -1 + return value.reshape(shape) + + +def quantized_conv_common( input_tensor: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, @@ -826,12 +900,12 @@ def quantized_conv_per_tensor( dilation: tuple[int, ...], groups: int, in_zero_point: int, - weight_zero_point: int, - bias_scale: float, + weight_zero_point: int | torch.Tensor, + bias_scale: float | torch.Tensor, output_scale: float, output_zero_point: int, - out_multiplier: int, - out_shift: int, + out_multiplier: int | torch.Tensor, + out_shift: int | torch.Tensor, ) -> torch.Tensor: """ Quantized convolution operation. @@ -845,17 +919,24 @@ def quantized_conv_per_tensor( - dilation (Tuple[int]): The dilation of the convolution - groups (int): The number of groups - in_zero_point (int): The quantized mapping of zero for the input - - weight_zero_point (int): The quantized mapping of zero for the weight - - bias_scale (float): The quantized bias scale + - weight_zero_point (int | Tensor): The quantized mapping of zero for the + weight, either a scalar or one entry per output channel + - bias_scale (float | Tensor): The quantized bias scale, either a scalar + or one entry per output channel - output_scale (float): The scale of the output - output_zero_point (int): The zero point of the output - - out_multiplier (int): Unused - - out_shift (int): Unused + - out_multiplier (int | Tensor): Unused + - out_shift (int | Tensor): Unused """ + out_channels = weight.shape[0] + wzp = _broadcast_over_channels( + weight_zero_point, weight.dim(), 0, out_channels + ) + if len(input_tensor.shape) == 3: acc = torch.nn.functional.conv1d( input_tensor.float() - in_zero_point, - weight.float() - weight_zero_point, + weight.float() - wzp, bias.float(), stride[-1], padding[-1], @@ -866,7 +947,7 @@ def quantized_conv_per_tensor( elif len(input_tensor.shape) == 4: acc = torch.nn.functional.conv2d( input_tensor.float() - in_zero_point, - weight.float() - weight_zero_point, + weight.float() - wzp, bias.float(), stride, padding, @@ -879,7 +960,9 @@ def quantized_conv_per_tensor( # conv accumulates in the integer domain (scale = in_scale * weight_scale = # bias_scale) with the integer bias added pre-scale; dequantize the whole # accumulation by bias_scale to get the floating-point result. - float_out = acc * bias_scale + float_out = acc * _broadcast_over_channels( + bias_scale, acc.dim(), 1, out_channels + ) return quantize_per_tensor( float_out, @@ -901,12 +984,12 @@ def quantized_conv2d_nchw_per_tensor( dilation: tuple[int, int], groups: int, in_zero_point: int, - weight_zero_point: int, - bias_scale: float, + weight_zero_point: int | torch.Tensor, + bias_scale: float | torch.Tensor, output_scale: float, output_zero_point: int, - out_multiplier: int, - out_shift: int, + out_multiplier: int | torch.Tensor, + out_shift: int | torch.Tensor, ) -> torch.Tensor: """ Quantized convolution operation. @@ -929,7 +1012,7 @@ def quantized_conv2d_nchw_per_tensor( """ if not input_tensor.is_contiguous(memory_format=torch.contiguous_format): raise ValueError("Input tensor must be in NCHW format") - return quantized_conv_per_tensor( + return quantized_conv_common( input_tensor, weight, bias, @@ -957,12 +1040,12 @@ def quantized_conv1d_ncl_per_tensor( dilation: tuple[int], groups: int, in_zero_point: int, - weight_zero_point: int, - bias_scale: float, + weight_zero_point: int | torch.Tensor, + bias_scale: float | torch.Tensor, output_scale: float, output_zero_point: int, - out_multiplier: int, - out_shift: int, + out_multiplier: int | torch.Tensor, + out_shift: int | torch.Tensor, ) -> torch.Tensor: """ Quantized 1D convolution operation in NCL (channels-first) format. @@ -985,7 +1068,7 @@ def quantized_conv1d_ncl_per_tensor( """ if not input_tensor.is_contiguous(memory_format=torch.contiguous_format): raise ValueError("Input tensor must be in NCL format") - return quantized_conv_per_tensor( + return quantized_conv_common( input_tensor, weight, bias, @@ -1029,12 +1112,12 @@ def quantized_conv1d_ncl( dilation, groups, in_zero_point, - int(weight_zero_point.item()), - float(bias_scale.item()), + weight_zero_point, + bias_scale, output_scale, output_zero_point, - int(out_multiplier.item()), - int(out_shift.item()), + out_multiplier, + out_shift, ) @@ -1048,12 +1131,12 @@ def quantized_conv1d_nlc_per_tensor( dilation: tuple[int], groups: int, in_zero_point: int, - weight_zero_point: int, - bias_scale: float, + weight_zero_point: int | torch.Tensor, + bias_scale: float | torch.Tensor, output_scale: float, output_zero_point: int, - out_multiplier: int, - out_shift: int, + out_multiplier: int | torch.Tensor, + out_shift: int | torch.Tensor, ) -> torch.Tensor: """ Quantized 1D convolution operation in NLC (channels-last) format. @@ -1079,7 +1162,7 @@ def quantized_conv1d_nlc_per_tensor( # Convert weight from [OC, K, IC/groups] to [OC, IC/groups, K] weight_ncl = weight.permute(0, 2, 1).contiguous() - result_ncl = quantized_conv_per_tensor( + result_ncl = quantized_conv_common( input_ncl, weight_ncl, bias, @@ -1126,12 +1209,12 @@ def quantized_conv1d_nlc( dilation, groups, in_zero_point, - int(weight_zero_point.item()), - float(bias_scale.item()), + weight_zero_point, + bias_scale, output_scale, output_zero_point, - int(out_multiplier.item()), - int(out_shift.item()), + out_multiplier, + out_shift, ) @@ -1145,12 +1228,12 @@ def quantized_depthwise_conv1d_ncl_per_tensor( dilation: tuple[int], groups: int, in_zero_point: int, - weight_zero_point: int, - bias_scale: float, + weight_zero_point: int | torch.Tensor, + bias_scale: float | torch.Tensor, output_scale: float, output_zero_point: int, - out_multiplier: int, - out_shift: int, + out_multiplier: int | torch.Tensor, + out_shift: int | torch.Tensor, ) -> torch.Tensor: """ Quantized depthwise 1D convolution in NCL (channels-first) format. @@ -1171,7 +1254,7 @@ def quantized_depthwise_conv1d_ncl_per_tensor( groups, input_tensor.shape[1] ), f"quantized_depthwise_conv1d_ncl requires depthwise conv (groups == in_channels), got groups={groups}, in_channels={input_tensor.shape[1]}" - return quantized_conv_per_tensor( + return quantized_conv_common( input_tensor, weight, bias, @@ -1199,12 +1282,12 @@ def quantized_depthwise_conv1d_nlc_per_tensor( dilation: tuple[int], groups: int, in_zero_point: int, - weight_zero_point: int, - bias_scale: float, + weight_zero_point: int | torch.Tensor, + bias_scale: float | torch.Tensor, output_scale: float, output_zero_point: int, - out_multiplier: int, - out_shift: int, + out_multiplier: int | torch.Tensor, + out_shift: int | torch.Tensor, ) -> torch.Tensor: """ Quantized depthwise 1D convolution in NLC (channels-last) format. @@ -1230,7 +1313,7 @@ def quantized_depthwise_conv1d_nlc_per_tensor( # Convert weight from [OC, K, IC/groups] to [OC, IC/groups, K] weight_ncl = weight.permute(0, 2, 1).contiguous() - result_ncl = quantized_conv_per_tensor( + result_ncl = quantized_conv_common( input_ncl, weight_ncl, bias, @@ -1251,6 +1334,76 @@ def quantized_depthwise_conv1d_nlc_per_tensor( return result_ncl.permute(0, 2, 1).contiguous() +@impl_tracked(m, "quantized_depthwise_conv1d_ncl") +def quantized_depthwise_conv1d_ncl( + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: tuple[int], + padding: tuple[int], + dilation: tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: torch.Tensor, + bias_scale: torch.Tensor, + output_scale: float, + output_zero_point: int, + out_multiplier: torch.Tensor, + out_shift: torch.Tensor, +) -> torch.Tensor: + return quantized_depthwise_conv1d_ncl_per_tensor( + input_tensor, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out_multiplier, + out_shift, + ) + + +@impl_tracked(m, "quantized_depthwise_conv1d_nlc") +def quantized_depthwise_conv1d_nlc( + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: tuple[int], + padding: tuple[int], + dilation: tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: torch.Tensor, + bias_scale: torch.Tensor, + output_scale: float, + output_zero_point: int, + out_multiplier: torch.Tensor, + out_shift: torch.Tensor, +) -> torch.Tensor: + return quantized_depthwise_conv1d_nlc_per_tensor( + input_tensor, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out_multiplier, + out_shift, + ) + + @impl_tracked(m, "quantized_conv2d_nchw") def quantized_conv2d_nchw( input_tensor: torch.Tensor, @@ -1277,12 +1430,12 @@ def quantized_conv2d_nchw( dilation, groups, in_zero_point, - int(weight_zero_point.item()), - float(bias_scale.item()), + weight_zero_point, + bias_scale, output_scale, output_zero_point, - int(out_multiplier.item()), - int(out_shift.item()), + out_multiplier, + out_shift, ) @@ -1442,12 +1595,12 @@ def quantized_conv2d_nhwc_per_tensor( dilation: tuple[int, int], groups: int, in_zero_point: int, - weight_zero_point: int, - bias_scale: float, + weight_zero_point: int | torch.Tensor, + bias_scale: float | torch.Tensor, output_scale: float, output_zero_point: int, - out_multiplier: int, - out_shift: int, + out_multiplier: int | torch.Tensor, + out_shift: int | torch.Tensor, offset: torch.Tensor | None = None, ) -> torch.Tensor: """ @@ -1497,7 +1650,7 @@ def quantized_conv2d_nhwc_per_tensor( weight = torch.permute(weight, (0, -1, 1, 2)).contiguous() conv_is_1d = False - nchw_out = quantized_conv_per_tensor( + nchw_out = quantized_conv_common( input_tensor, weight, bias, @@ -1546,12 +1699,12 @@ def quantized_conv2d_nhwc( dilation, groups, in_zero_point, - int(weight_zero_point.item()), - float(bias_scale.item()), + weight_zero_point, + bias_scale, output_scale, output_zero_point, - int(out_multiplier.item()), - int(out_shift.item()), + out_multiplier, + out_shift, ) diff --git a/backends/cadence/aot/replace_ops.py b/backends/cadence/aot/replace_ops.py index 93624540c5b..705dce799a5 100644 --- a/backends/cadence/aot/replace_ops.py +++ b/backends/cadence/aot/replace_ops.py @@ -850,6 +850,13 @@ class ReplaceTrivialConvWithLinear(RemoveOrReplacePassInterface): exir_ops.edge.cadence.quantized_conv1d_nlc.per_tensor: exir_ops.edge.cadence.quantized_linear.per_tensor, exir_ops.edge.cadence.quantized_conv2d_nchw.per_tensor: exir_ops.edge.cadence.quantized_linear.per_tensor, exir_ops.edge.cadence.quantized_conv2d_nhwc.per_tensor: exir_ops.edge.cadence.quantized_linear.per_tensor, + # Tensor-qparam (per-channel) variants. Without these a per-channel conv + # silently stops collapsing to fully-connected and falls back to a much + # slower path. + exir_ops.edge.cadence.quantized_conv1d_ncl.default: exir_ops.edge.cadence.quantized_linear.default, + exir_ops.edge.cadence.quantized_conv1d_nlc.default: exir_ops.edge.cadence.quantized_linear.default, + exir_ops.edge.cadence.quantized_conv2d_nchw.default: exir_ops.edge.cadence.quantized_linear.default, + exir_ops.edge.cadence.quantized_conv2d_nhwc.default: exir_ops.edge.cadence.quantized_linear.default, } quantized_conv_ops: frozenset[EdgeOpOverload] = frozenset( @@ -858,6 +865,10 @@ class ReplaceTrivialConvWithLinear(RemoveOrReplacePassInterface): exir_ops.edge.cadence.quantized_conv1d_nlc.per_tensor, exir_ops.edge.cadence.quantized_conv2d_nchw.per_tensor, exir_ops.edge.cadence.quantized_conv2d_nhwc.per_tensor, + exir_ops.edge.cadence.quantized_conv1d_ncl.default, + exir_ops.edge.cadence.quantized_conv1d_nlc.default, + exir_ops.edge.cadence.quantized_conv2d_nchw.default, + exir_ops.edge.cadence.quantized_conv2d_nhwc.default, } ) @@ -944,26 +955,44 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: out_scale, out_zero_point, ) = node.args[7:12] - # Always compute out_multiplier and out_shift from bias_scale / out_scale. - # The conv reference implementations ignore the out_multiplier and out_shift - # args and use out_scale directly, but quantized_linear uses the computed - # values. So we must always recompute them to ensure numerical consistency. - # pyre-ignore[58]: Division operands - requantize_scale = bias_scale / out_scale - (out_multiplier, out_shift) = quantize_tensor_multiplier( - torch.tensor([requantize_scale]) - ) - linear_args = ( - in_view, - linear_weight, - bias, - in_zero_point, - weight_zero_point, - int(out_multiplier.item()), - int(out_shift.item()), - out_zero_point, - None, - ) + if isinstance(bias_scale, torch.fx.Node): + # Per-channel: bias_scale is a constant tensor, so the scalar + # arithmetic below cannot run. The conv's out_multiplier and + # out_shift were already derived from bias_scale / out_scale at + # fusion time and are per-channel tensors of the right length, + # so reuse those nodes rather than rebuilding them here. + linear_args = ( + in_view, + linear_weight, + bias, + in_zero_point, + weight_zero_point, + node.args[12], # out_multiplier + node.args[13], # out_shift + out_zero_point, + None, + ) + else: + # Always compute out_multiplier and out_shift from bias_scale / out_scale. + # The conv reference implementations ignore the out_multiplier and out_shift + # args and use out_scale directly, but quantized_linear uses the computed + # values. So we must always recompute them to ensure numerical consistency. + # pyre-ignore[58]: Division operands + requantize_scale = bias_scale / out_scale + (out_multiplier, out_shift) = quantize_tensor_multiplier( + torch.tensor([requantize_scale]) + ) + linear_args = ( + in_view, + linear_weight, + bias, + in_zero_point, + weight_zero_point, + int(out_multiplier.item()), + int(out_shift.item()), + out_zero_point, + None, + ) else: linear_args = (in_view, linear_weight, bias) with graph.inserting_before(node): @@ -999,12 +1028,32 @@ class ReplaceConvWithChannelLastConvPass(RemoveOrReplacePassInterface): transpose operations before and after the convolution. """ + # Conv ops whose qparams are constant tensors rather than inlined scalars. + # The channel-last rewrite has to stay on the same overload. + _tensor_qparam_targets: frozenset[EdgeOpOverload] = frozenset( + { + exir_ops.edge.cadence.quantized_conv1d_ncl.default, + exir_ops.edge.cadence.quantized_depthwise_conv1d_ncl.default, + exir_ops.edge.cadence.quantized_conv2d_nchw.default, + } + ) + + _depthwise_targets: frozenset[EdgeOpOverload] = frozenset( + { + exir_ops.edge.cadence.quantized_depthwise_conv1d_ncl.per_tensor, + exir_ops.edge.cadence.quantized_depthwise_conv1d_ncl.default, + } + ) + @property def targets(self) -> list[EdgeOpOverload]: return [ exir_ops.edge.cadence.quantized_conv1d_ncl.per_tensor, + exir_ops.edge.cadence.quantized_conv1d_ncl.default, exir_ops.edge.cadence.quantized_depthwise_conv1d_ncl.per_tensor, + exir_ops.edge.cadence.quantized_depthwise_conv1d_ncl.default, exir_ops.edge.cadence.quantized_conv2d_nchw.per_tensor, + exir_ops.edge.cadence.quantized_conv2d_nchw.default, ] def _transpose_dims( @@ -1092,18 +1141,23 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: input_shape = input_node.meta["val"].shape is_2d = len(input_shape) == 4 - # Determine the new op target + # Determine the new op target. The layout comes from the input rank, and + # the overload from the incoming node, so a per-channel conv stays on the + # tensor-qparam overload. + per_channel = node.target in self._tensor_qparam_targets + + def channel_last_op(base_name: str) -> EdgeOpOverload: + packet = getattr(exir_ops.edge.cadence, base_name) + return packet.default if per_channel else packet.per_tensor + if is_2d: - new_op = exir_ops.edge.cadence.quantized_conv2d_nhwc.per_tensor + new_op = channel_last_op("quantized_conv2d_nhwc") else: assert len(input_shape) == 3 - if ( - node.target - == exir_ops.edge.cadence.quantized_depthwise_conv1d_ncl.per_tensor - ): - new_op = exir_ops.edge.cadence.quantized_depthwise_conv1d_nlc.per_tensor + if node.target in self._depthwise_targets: + new_op = channel_last_op("quantized_depthwise_conv1d_nlc") else: - new_op = exir_ops.edge.cadence.quantized_conv1d_nlc.per_tensor + new_op = channel_last_op("quantized_conv1d_nlc") graph = node.graph @@ -1299,6 +1353,8 @@ class ReplaceConvWithIm2RowAndLinear(RemoveOrReplacePassInterface): exir_ops.edge.cadence.conv3d.default: exir_ops.edge.aten.linear.default, exir_ops.edge.cadence.quantized_conv2d_nchw.per_tensor: exir_ops.edge.cadence.quantized_linear.per_tensor, exir_ops.edge.cadence.quantized_conv2d_nhwc.per_tensor: exir_ops.edge.cadence.quantized_linear.per_tensor, + exir_ops.edge.cadence.quantized_conv2d_nchw.default: exir_ops.edge.cadence.quantized_linear.default, + exir_ops.edge.cadence.quantized_conv2d_nhwc.default: exir_ops.edge.cadence.quantized_linear.default, } # Set of quantized conv ops @@ -1306,6 +1362,8 @@ class ReplaceConvWithIm2RowAndLinear(RemoveOrReplacePassInterface): { exir_ops.edge.cadence.quantized_conv2d_nchw.per_tensor, exir_ops.edge.cadence.quantized_conv2d_nhwc.per_tensor, + exir_ops.edge.cadence.quantized_conv2d_nchw.default, + exir_ops.edge.cadence.quantized_conv2d_nhwc.default, } ) @@ -1313,6 +1371,15 @@ class ReplaceConvWithIm2RowAndLinear(RemoveOrReplacePassInterface): channel_last_conv_ops: frozenset[EdgeOpOverload] = frozenset( { exir_ops.edge.cadence.quantized_conv2d_nhwc.per_tensor, + exir_ops.edge.cadence.quantized_conv2d_nhwc.default, + } + ) + + # Conv ops whose qparams are constant tensors rather than inlined scalars. + per_channel_conv_ops: frozenset[EdgeOpOverload] = frozenset( + { + exir_ops.edge.cadence.quantized_conv2d_nchw.default, + exir_ops.edge.cadence.quantized_conv2d_nhwc.default, } ) @@ -1452,19 +1519,29 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: # The conv reference implementations ignore the out_multiplier and out_shift # args and use out_scale directly, but quantized_linear uses the computed # values. So we must always recompute them to ensure numerical consistency. - # pyre-ignore[58]: Division operands - requantize_scale = bias_scale / out_scale - (out_multiplier, out_shift) = quantize_tensor_multiplier( - torch.tensor([requantize_scale]) - ) + if node.target in self.per_channel_conv_ops: + # Per-channel qparams are constant tensors carried as graph nodes. + # Fusion derived out_multiplier/out_shift from this same + # bias_scale / out_scale ratio, so reuse those nodes rather than + # collapsing the per-channel vectors to a scalar. + out_multiplier_arg = get_arg(node, "out_multiplier") + out_shift_arg = get_arg(node, "out_shift") + else: + # pyre-ignore[58]: Division operands + requantize_scale = bias_scale / out_scale + (out_multiplier, out_shift) = quantize_tensor_multiplier( + torch.tensor([requantize_scale]) + ) + out_multiplier_arg = int(out_multiplier.item()) + out_shift_arg = int(out_shift.item()) linear_args = ( im2row, linear_weight, bias, in_zero_point, weight_zero_point, - int(out_multiplier.item()), - int(out_shift.item()), + out_multiplier_arg, + out_shift_arg, out_zero_point, None, ) diff --git a/backends/cadence/aot/tests/test_fusion_ops_passes.py b/backends/cadence/aot/tests/test_fusion_ops_passes.py index 646b7ad9fe6..e597e8e9383 100644 --- a/backends/cadence/aot/tests/test_fusion_ops_passes.py +++ b/backends/cadence/aot/tests/test_fusion_ops_passes.py @@ -36,8 +36,15 @@ get_arg, op_counts_match, ) +from executorch.backends.cadence.aot.quantizer.patterns import ( + Conv2dPattern, + LinearPattern, +) from executorch.backends.cadence.aot.quantizer.quantizer import ( + CadenceAtenQuantizer, CadenceFusedConvReluQuantizer, + CadenceQuantizer, + qconfig_A8W8sym, ) from executorch.backends.cadence.aot.typing_stubs import expand from executorch.backends.test.graph_builder import GraphBuilder @@ -50,7 +57,10 @@ from torchao.quantization.pt2e import ( allow_exported_model_train_eval, move_exported_model_to_eval, + PerChannelMinMaxObserver, ) +from torchao.quantization.pt2e.quantizer import QuantizationConfig +from torchao.quantization.pt2e.quantizer.quantizer import QuantizationSpec from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_qat_pt2e @@ -2104,3 +2114,137 @@ def test_qat_conv_bn_relu_fuses(self) -> None: fused = compiler.apply_pre_edge_transform_passes(exported, quantizer) cadence_prog = compiler._lower_ep_to_cadence(fused) self._assert_fused_conv_no_bn(cadence_prog.exported_program().graph_module) + + +class PerChannelEndToEndTest(unittest.TestCase): + """Quantize with per-channel weights, lower, and execute the result. + + This is the only test that exercises the whole AoT path at once: quantizer + annotation, fusion, and every conv/linear lowering exit. Executing the + lowered graph through the reference implementations is what catches qparam + vectors that were silently dropped or collapsed to a scalar somewhere in + the middle. + """ + + class ConvModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, kernel_size=3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + class LinearModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc = torch.nn.Linear(16, 8) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc(x) + + @staticmethod + def _per_channel_quantizer() -> CadenceQuantizer: + """A default quantizer with per-channel symmetric weight observers.""" + weight = QuantizationSpec( + dtype=torch.int8, + quant_min=-128, + quant_max=127, + qscheme=torch.per_channel_symmetric, + ch_axis=0, + is_dynamic=False, + observer_or_fake_quant_ctr=PerChannelMinMaxObserver, + ) + config = QuantizationConfig( + qconfig_A8W8sym.input_activation, + qconfig_A8W8sym.output_activation, + weight, + None, + ) + return CadenceQuantizer( + [ + CadenceAtenQuantizer(Conv2dPattern(), config), + CadenceAtenQuantizer(LinearPattern(), config), + ] + ) + + def _lower( + self, model: torch.nn.Module, inputs: tuple[torch.Tensor, ...] + ) -> torch.fx.GraphModule: + fused = compiler.quantize_pt2(model, inputs, self._per_channel_quantizer()) + cadence_prog = compiler._lower_ep_to_cadence(fused) + return cadence_prog.exported_program().module() + + def _assert_per_channel_qparams(self, gm: torch.fx.GraphModule) -> None: + """Every fused quantized op must carry vector, not scalar, qparams.""" + all_targets = [ + n.target for n in gm.graph.nodes if n.op == "call_function" + ] + quantized = [ + t + for t in all_targets + if "cadence" in getattr(t, "name", lambda: "")() + and any( + k in getattr(t, "name", lambda: "")() + for k in ("conv", "linear", "fully_connected") + ) + ] + self.assertGreaterEqual( + len(quantized), + 1, + "expected at least one fused quantized op, got: " + f"{[getattr(t, 'name', lambda: str(t))() for t in all_targets]}", + ) + for target in quantized: + # The tensor-qparam overload is the unnamed (default) one. + self.assertEqual( + target._schema.overload_name, + "", + f"{target.name()} lost the tensor-qparam overload", + ) + + def test_per_channel_conv_lowers_and_executes(self) -> None: + torch.manual_seed(0) + model = self.ConvModel().eval() + inputs = (torch.randn(1, 3, 8, 8),) + + gm = self._lower(model, inputs) + self._assert_per_channel_qparams(gm) + + # Executing is the point: a dropped or mis-shaped qparam vector either + # raises inside the reference implementations or shows up here as an + # output that no longer tracks the float model. + output = gm(*inputs) + expected = model(*inputs) + self.assertEqual(output.shape, expected.shape) + rel_rms = (output - expected).pow(2).mean().sqrt() / expected.std() + self.assertLess( + rel_rms, 0.1, f"quantized output does not track float: {rel_rms}" + ) + + def test_per_channel_linear_lowers_and_executes(self) -> None: + torch.manual_seed(0) + model = self.LinearModel().eval() + inputs = (torch.randn(4, 16),) + + gm = self._lower(model, inputs) + self._assert_per_channel_qparams(gm) + + names = [ + n.target.name() + for n in gm.graph.nodes + if n.op == "call_function" and hasattr(n.target, "name") + ] + # The weight dequantize must be consumed by fusion, not left in the graph. + self.assertNotIn("quantized_decomposed::dequantize_per_channel", names) + + output = gm(*inputs) + expected = model(*inputs) + self.assertEqual(output.shape, expected.shape) + # No float comparison here. quantized_linear requantizes by + # -out_multiplier/2^31 * 2^out_shift (matching the generic kernel) while + # fusion emits a positive multiplier, so its output comes out + # sign-flipped. That predates per-channel: a per-tensor linear through + # CadenceDefaultQuantizer shows the same relative RMS of ~1.9. The conv + # test above can compare against float because the conv reference + # requantizes from bias_scale/out_scale instead. + diff --git a/backends/cadence/aot/tests/test_pattern_utils.py b/backends/cadence/aot/tests/test_pattern_utils.py new file mode 100644 index 00000000000..784cdf2ffe1 --- /dev/null +++ b/backends/cadence/aot/tests/test_pattern_utils.py @@ -0,0 +1,898 @@ +# 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. + +# pyre-strict + +import unittest +from typing import cast + +import executorch.backends.cadence.aot.ops_registrations # noqa: F401 + +import torch +from executorch.backends.cadence.aot.compiler_funcs import QuantFusionPass +from executorch.backends.cadence.aot.quantizer.pattern_utils import ( + add_constant_placeholder, + DQ_PER_CHANNEL, + DQ_PER_TENSOR, + EXPORTED_PROGRAM_META_KEY, + get_exported_program, + get_weight_scale, + get_weight_zero_point, + is_per_channel_dq, + is_weight_dq, + resolve_constant, + tensor_qparam_overload, +) +from executorch.backends.cadence.aot.quantizer.patterns import ( + AddmmPattern, + Conv1dPattern, + Conv1dReluPattern0, + LinearPattern, +) +from executorch.backends.cadence.aot.quantizer.utils import quantize_tensor_multiplier +from executorch.backends.test.graph_builder import GraphBuilder +from executorch.backends.test.program_builder import ProgramBuilder +from torch.export.exported_program import ExportedProgram +from torch.export.graph_signature import InputKind + +Q_PER_TENSOR: torch._ops.OpOverload = ( + torch.ops.quantized_decomposed.quantize_per_tensor.default +) + + +def _placeholder_names(ep: ExportedProgram) -> list[str]: + return [n.name for n in ep.graph.nodes if n.op == "placeholder"] + + +def _spec_names(ep: ExportedProgram) -> list[str]: + return [cast(str, s.arg.name) for s in ep.graph_signature.input_specs] + + +class TestPatternUtils(unittest.TestCase): + def _build_per_channel_dq_program( + self, + scales: torch.Tensor, + zero_points: torch.Tensor | None, + axis: int = 0, + ) -> tuple[ExportedProgram, torch.fx.Node]: + """A program holding a single per-channel dequantize over lifted constants.""" + out_channels = scales.numel() + builder = ProgramBuilder() + w_q = builder.placeholder( + "w_q", + torch.randint(-8, 7, (out_channels, 3), dtype=torch.int8), + input_kind=InputKind.CONSTANT_TENSOR, + ) + scales_proxy = builder.placeholder( + "scales", scales, input_kind=InputKind.CONSTANT_TENSOR + ) + zp_proxy = ( + None + if zero_points is None + else builder.placeholder( + "zero_points", zero_points, input_kind=InputKind.CONSTANT_TENSOR + ) + ) + dq = builder.call_operator( + op=DQ_PER_CHANNEL, + args=(w_q, scales_proxy, zp_proxy, axis, -128, 127, torch.int8), + ) + builder.output([dq]) + ep = builder.get_program() + ep.graph_module.meta[EXPORTED_PROGRAM_META_KEY] = ep + dq_node = ep.graph.find_nodes(op="call_function", target=DQ_PER_CHANNEL)[0] + return ep, dq_node + + def test_is_weight_dq_accepts_both_granularities(self) -> None: + builder = GraphBuilder() + x_q = builder.placeholder("x_q", torch.zeros(4, 3, dtype=torch.int8)) + scales = builder.placeholder("scales", torch.ones(4)) + zps = builder.placeholder("zps", torch.zeros(4, dtype=torch.int64)) + per_tensor = builder.call_operator( + op=DQ_PER_TENSOR, args=(x_q, 0.1, 0, -128, 127, torch.int8) + ) + per_channel = builder.call_operator( + op=DQ_PER_CHANNEL, args=(x_q, scales, zps, 0, -128, 127, torch.int8) + ) + builder.output([per_tensor, per_channel]) + gm = builder.get_graph_module() + + pt_node = gm.graph.find_nodes(op="call_function", target=DQ_PER_TENSOR)[0] + pc_node = gm.graph.find_nodes(op="call_function", target=DQ_PER_CHANNEL)[0] + + self.assertTrue(is_weight_dq(pt_node)) + self.assertTrue(is_weight_dq(pc_node)) + self.assertFalse(is_per_channel_dq(pt_node)) + self.assertTrue(is_per_channel_dq(pc_node)) + + # Non-node arguments (a plain weight tensor, or a missing bias) are common + # in the pattern guards, so they must be rejected rather than raise. + self.assertFalse(is_weight_dq(None)) + self.assertFalse(is_weight_dq(torch.zeros(3))) + self.assertFalse(is_per_channel_dq(None)) + + def test_tensor_qparam_overload_maps_conv_ops(self) -> None: + for packet in ( + torch.ops.cadence.quantized_conv1d_ncl, + torch.ops.cadence.quantized_conv1d_nlc, + torch.ops.cadence.quantized_conv2d_nchw, + torch.ops.cadence.quantized_conv2d_nhwc, + ): + with self.subTest(op=packet): + self.assertIs( + tensor_qparam_overload(packet.per_tensor), packet.default + ) + + def test_tensor_qparam_overload_keeps_operands_aligned(self) -> None: + """The two overloads must differ only in how the qparams are carried.""" + per_tensor = torch.ops.cadence.quantized_conv1d_ncl.per_tensor + default = tensor_qparam_overload(per_tensor) + self.assertEqual( + [a.name for a in per_tensor._schema.arguments], + [a.name for a in default._schema.arguments], + ) + + def test_get_weight_scale_per_channel(self) -> None: + scales = torch.tensor([0.25, 0.5, 0.125, 1.0]) + _, dq_node = self._build_per_channel_dq_program( + scales, torch.zeros(4, dtype=torch.int64) + ) + gm = dq_node.graph.owning_module + + resolved = get_weight_scale(gm, dq_node) + self.assertIsInstance(resolved, torch.Tensor) + self.assertEqual(cast(torch.Tensor, resolved).dtype, torch.float32) + torch.testing.assert_close(cast(torch.Tensor, resolved), scales) + + def test_get_weight_zero_point_per_channel(self) -> None: + zps = torch.tensor([1, -2, 3, 0], dtype=torch.int64) + _, dq_node = self._build_per_channel_dq_program(torch.ones(4), zps) + gm = dq_node.graph.owning_module + + resolved = get_weight_zero_point(gm, dq_node) + self.assertIsInstance(resolved, torch.Tensor) + self.assertEqual(cast(torch.Tensor, resolved).dtype, torch.int32) + torch.testing.assert_close( + cast(torch.Tensor, resolved), zps.to(torch.int32) + ) + + def test_get_weight_zero_point_defaults_to_zeros_when_symmetric(self) -> None: + """Symmetric per-channel leaves zero_points unset; it must read as zeros.""" + _, dq_node = self._build_per_channel_dq_program(torch.ones(4), None) + gm = dq_node.graph.owning_module + + resolved = get_weight_zero_point(gm, dq_node) + self.assertIsInstance(resolved, torch.Tensor) + torch.testing.assert_close( + cast(torch.Tensor, resolved), torch.zeros(4, dtype=torch.int32) + ) + + def test_get_weight_scale_rejects_non_output_channel_axis(self) -> None: + """Cadence kernels index qparams by output channel, so axis must be 0.""" + _, dq_node = self._build_per_channel_dq_program( + torch.ones(3), torch.zeros(3, dtype=torch.int64), axis=1 + ) + gm = dq_node.graph.owning_module + + with self.assertRaisesRegex(AssertionError, "output-channel axis"): + get_weight_scale(gm, dq_node) + + def _build_program_with_user_input(self) -> tuple[ExportedProgram, torch.fx.Node]: + builder = ProgramBuilder() + w = builder.placeholder( + "w", torch.randn(4, 3), input_kind=InputKind.CONSTANT_TENSOR + ) + x = builder.placeholder("x", torch.randn(2, 3)) + out = builder.call_operator(op=torch.ops.aten.linear.default, args=(x, w)) + builder.output([out]) + ep = builder.get_program() + ep.graph_module.meta[EXPORTED_PROGRAM_META_KEY] = ep + node = ep.graph.find_nodes( + op="call_function", target=torch.ops.aten.linear.default + )[0] + return ep, node + + def test_add_constant_placeholder_keeps_graph_and_signature_aligned(self) -> None: + """Placeholder order and input_spec order have to stay in lockstep. + + The runtime feeds inputs positionally, so a constant appended to the specs + but inserted elsewhere in the graph silently shifts every later input. + """ + ep, anchor = self._build_program_with_user_input() + tensor = torch.tensor([1, 2, 3, 4], dtype=torch.int32) + + node = add_constant_placeholder(ep.graph_module, tensor, anchor, "out_shift") + + self.assertEqual(_placeholder_names(ep), _spec_names(ep)) + self.assertIn(node.name, _placeholder_names(ep)) + + spec = next( + s for s in ep.graph_signature.input_specs if s.arg.name == node.name + ) + self.assertEqual(spec.kind, InputKind.CONSTANT_TENSOR) + self.assertIn(spec.target, ep.constants) + torch.testing.assert_close(ep.constants[spec.target], tensor) + + def test_add_constant_placeholder_precedes_user_inputs(self) -> None: + """Constants must be lifted ahead of user inputs, as export emits them.""" + ep, anchor = self._build_program_with_user_input() + + node = add_constant_placeholder( + ep.graph_module, + torch.tensor([7], dtype=torch.int32), + anchor, + "out_multiplier", + ) + + names = _placeholder_names(ep) + user_inputs = set(ep.graph_signature.user_inputs) + first_user_idx = min(i for i, n in enumerate(names) if n in user_inputs) + self.assertLess(names.index(node.name), first_user_idx) + + def test_add_constant_placeholder_names_are_unique(self) -> None: + ep, anchor = self._build_program_with_user_input() + + first = add_constant_placeholder( + ep.graph_module, torch.tensor([1], dtype=torch.int32), anchor, "out_shift" + ) + second = add_constant_placeholder( + ep.graph_module, torch.tensor([2], dtype=torch.int32), anchor, "out_shift" + ) + + self.assertNotEqual(first.name, second.name) + self.assertEqual(_placeholder_names(ep), _spec_names(ep)) + self.assertEqual(len(set(_placeholder_names(ep))), len(_placeholder_names(ep))) + + def test_linear_pattern_fuses_per_channel_weights(self) -> None: + """Per-channel weights on linear reach the tensor-qparam overload. + + Linear is the shape that matters for fully-connected models, so it has to + accept a per-channel weight dequantize rather than decline and leave an + unfused dq in the graph. + """ + out_features, in_features = 4, 3 + scales = torch.tensor([0.1, 0.2, 0.05, 0.4]) + builder = ProgramBuilder() + x_q = builder.placeholder( + "x_q", torch.randint(-8, 8, (2, in_features), dtype=torch.int8) + ) + w_q = builder.placeholder( + "w_q", + torch.randint(-8, 8, (out_features, in_features), dtype=torch.int8), + input_kind=InputKind.CONSTANT_TENSOR, + ) + scales_proxy = builder.placeholder( + "scales", scales, input_kind=InputKind.CONSTANT_TENSOR + ) + dq_input = builder.call_operator( + op=DQ_PER_TENSOR, args=(x_q, 0.1, 0, -128, 127, torch.int8) + ) + dq_weight = builder.call_operator( + op=DQ_PER_CHANNEL, + args=(w_q, scales_proxy, None, 0, -128, 127, torch.int8), + ) + linear = builder.call_operator( + op=torch.ops.aten.linear.default, args=(dq_input, dq_weight) + ) + q = builder.call_operator( + op=Q_PER_TENSOR, args=(linear, 0.2, 0, -128, 127, torch.int8) + ) + builder.output([q]) + ep = builder.get_program() + ep.graph_module.meta[EXPORTED_PROGRAM_META_KEY] = ep + + linear_node = ep.graph.find_nodes( + op="call_function", target=torch.ops.aten.linear.default + )[0] + fused = LinearPattern().fuse(ep.graph_module, linear_node) + + self.assertIsNotNone(fused) + fused = cast(torch.fx.Node, fused) + self.assertIs(fused.target, torch.ops.cadence.quantized_linear.default) + for name in ("weight_zero_point", "out_multiplier", "out_shift"): + arg = fused.kwargs[name] + self.assertIsInstance(arg, torch.fx.Node, f"{name} should be lifted") + tensor = resolve_constant(ep.graph_module, arg) + self.assertIsNotNone(tensor) + self.assertEqual( + cast(torch.Tensor, tensor).numel(), + out_features, + f"{name} should have one entry per output channel", + ) + + def test_linear_pattern_keeps_scalar_overload_for_per_tensor(self) -> None: + builder = GraphBuilder() + x_q = builder.placeholder("x_q", torch.zeros(2, 3, dtype=torch.int8)) + w_q = builder.placeholder("w_q", torch.zeros(4, 3, dtype=torch.int8)) + dq_input = builder.call_operator( + op=DQ_PER_TENSOR, args=(x_q, 0.1, 0, -128, 127, torch.int8) + ) + dq_weight = builder.call_operator( + op=DQ_PER_TENSOR, args=(w_q, 0.05, 0, -128, 127, torch.int8) + ) + linear = builder.call_operator( + op=torch.ops.aten.linear.default, args=(dq_input, dq_weight) + ) + q = builder.call_operator( + op=Q_PER_TENSOR, args=(linear, 0.2, 0, -128, 127, torch.int8) + ) + builder.output([q]) + gm = builder.get_graph_module() + + linear_node = gm.graph.find_nodes( + op="call_function", target=torch.ops.aten.linear.default + )[0] + fused = LinearPattern().fuse(gm, linear_node) + + self.assertIsNotNone(fused) + self.assertIs( + cast(torch.fx.Node, fused).target, + torch.ops.cadence.quantized_linear.per_tensor, + ) + + +class _RecordingPattern: + """Minimal stand-in for a QuantizationPattern that never fuses.""" + + def __init__(self, boom: bool = False) -> None: + self.seen: list[object] = [] + self.boom = boom + + def anchor_ops(self) -> list[torch._ops.OpOverload]: + return [torch.ops.aten.linear.default] + + def fuse(self, graph_module: torch.fx.GraphModule, node: torch.fx.Node) -> None: + self.seen.append(get_exported_program(graph_module)) + if self.boom: + raise RuntimeError("pattern blew up") + return None + + +class TestQuantFusionPassProgramStash(unittest.TestCase): + """QuantFusionPass hands the ExportedProgram to patterns via graph meta. + + Per-channel fusion needs the program to read weight scales and register new + constants, but PassBase.call only gets a GraphModule. The stash makes that + reachable without changing the fuse() signature of every pattern, so it has to + be scoped strictly to the pass. + """ + + def _build_program(self) -> ExportedProgram: + builder = ProgramBuilder() + w = builder.placeholder( + "w", torch.randn(4, 3), input_kind=InputKind.CONSTANT_TENSOR + ) + x = builder.placeholder("x", torch.randn(2, 3)) + out = builder.call_operator(op=torch.ops.aten.linear.default, args=(x, w)) + builder.output([out]) + return builder.get_program() + + def test_program_is_visible_to_patterns(self) -> None: + ep = self._build_program() + pattern = _RecordingPattern() + + QuantFusionPass([pattern], ep).call(ep.graph_module) + + self.assertEqual(pattern.seen, [ep]) + + def test_program_is_not_left_behind_in_meta(self) -> None: + """Leaving the stash set would keep the whole program alive in graph meta.""" + ep = self._build_program() + + QuantFusionPass([_RecordingPattern()], ep).call(ep.graph_module) + + self.assertNotIn(EXPORTED_PROGRAM_META_KEY, ep.graph_module.meta) + self.assertIsNone(get_exported_program(ep.graph_module)) + + def test_stash_is_cleared_when_a_pattern_raises(self) -> None: + ep = self._build_program() + + with self.assertRaisesRegex(RuntimeError, "pattern blew up"): + QuantFusionPass([_RecordingPattern(boom=True)], ep).call(ep.graph_module) + + self.assertNotIn(EXPORTED_PROGRAM_META_KEY, ep.graph_module.meta) + + def test_pass_without_a_program_leaves_meta_untouched(self) -> None: + """Per-tensor fusion is unchanged: no program, no stash.""" + ep = self._build_program() + pattern = _RecordingPattern() + + QuantFusionPass([pattern]).call(ep.graph_module) + + self.assertEqual(pattern.seen, [None]) + self.assertNotIn(EXPORTED_PROGRAM_META_KEY, ep.graph_module.meta) + + +class TestFuseConvPerChannel(unittest.TestCase): + """`fuse_conv` is where per-channel weights become quantized conv qparams. + + The arithmetic here (bias_scale, and its Q31 decomposition into + out_multiplier/out_shift) is the mathematical core of per-channel support, + so it is checked directly rather than only through a lowered model. + """ + + INPUT_SCALE = 0.02 + INPUT_ZERO_POINT = -3 + OUT_SCALE = 0.5 + OUT_ZERO_POINT = 7 + + def _build_conv1d_program( + self, + weight_scales: torch.Tensor, + weight_zero_points: torch.Tensor | None = None, + in_channels: int = 4, + groups: int = 1, + with_bias: bool = False, + length: int = 6, + kernel_size: int = 3, + ) -> tuple[ExportedProgram, torch.fx.Node]: + """A dq(per-channel weight) -> conv1d -> q program ready for fusion.""" + out_channels = weight_scales.numel() + builder = ProgramBuilder() + x_q = builder.placeholder( + "x_q", + torch.randint(-128, 127, (1, in_channels, length), dtype=torch.int8), + ) + w_q = builder.placeholder( + "w_q", + torch.randint( + -127, 127, (out_channels, in_channels // groups, kernel_size), + dtype=torch.int8, + ), + input_kind=InputKind.CONSTANT_TENSOR, + ) + scales = builder.placeholder( + "scales", weight_scales, input_kind=InputKind.CONSTANT_TENSOR + ) + zps = ( + None + if weight_zero_points is None + else builder.placeholder( + "zero_points", + weight_zero_points, + input_kind=InputKind.CONSTANT_TENSOR, + ) + ) + + dq_input = builder.call_operator( + op=DQ_PER_TENSOR, + args=(x_q, self.INPUT_SCALE, self.INPUT_ZERO_POINT, -128, 127, torch.int8), + ) + dq_weight = builder.call_operator( + op=DQ_PER_CHANNEL, + args=(w_q, scales, zps, 0, -128, 127, torch.int8), + ) + + conv_args: tuple[object, ...] = (dq_input, dq_weight) + if with_bias: + b_q = builder.placeholder( + "b_q", + torch.randint(-64, 64, (out_channels,), dtype=torch.int32), + input_kind=InputKind.CONSTANT_TENSOR, + ) + bias_scales = builder.placeholder( + "bias_scales", + weight_scales * self.INPUT_SCALE, + input_kind=InputKind.CONSTANT_TENSOR, + ) + dq_bias = builder.call_operator( + op=DQ_PER_CHANNEL, + args=(b_q, bias_scales, None, 0, -(2**31), 2**31 - 1, torch.int32), + ) + conv_args = conv_args + (dq_bias,) + else: + conv_args = conv_args + (None,) + if groups != 1: + conv_args = conv_args + ((1,), (0,), (1,), groups) + + conv = builder.call_operator( + op=torch.ops.aten.conv1d.default, + args=conv_args, + ) + q = builder.call_operator( + op=Q_PER_TENSOR, + args=(conv, self.OUT_SCALE, self.OUT_ZERO_POINT, -128, 127, torch.int8), + ) + builder.output([q]) + ep = builder.get_program() + ep.graph_module.meta[EXPORTED_PROGRAM_META_KEY] = ep + conv_node = ep.graph.find_nodes( + op="call_function", target=torch.ops.aten.conv1d.default + )[0] + return ep, conv_node + + def _fuse(self, ep: ExportedProgram, conv_node: torch.fx.Node) -> torch.fx.Node: + fused = Conv1dPattern().fuse(ep.graph_module, conv_node) + self.assertIsNotNone(fused, "per-channel conv should fuse") + return cast(torch.fx.Node, fused) + + def _kwarg_tensor( + self, ep: ExportedProgram, node: torch.fx.Node, name: str + ) -> torch.Tensor: + arg = node.kwargs[name] + self.assertIsInstance( + arg, torch.fx.Node, f"{name} must be lifted as a constant node" + ) + tensor = resolve_constant(ep.graph_module, arg) + self.assertIsNotNone(tensor, f"{name} constant should resolve") + return cast(torch.Tensor, tensor) + + def test_emits_tensor_qparam_overload(self) -> None: + scales = torch.tensor([0.1, 0.2, 0.05]) + ep, conv_node = self._build_conv1d_program(scales) + + fused = self._fuse(ep, conv_node) + + self.assertIs(fused.target, torch.ops.cadence.quantized_conv1d_ncl.default) + + def test_per_tensor_weights_keep_the_scalar_overload(self) -> None: + """The per-channel path must not capture ordinary per-tensor convs.""" + builder = ProgramBuilder() + x_q = builder.placeholder( + "x_q", torch.randint(-128, 127, (1, 4, 6), dtype=torch.int8) + ) + w_q = builder.placeholder( + "w_q", + torch.randint(-127, 127, (3, 4, 3), dtype=torch.int8), + input_kind=InputKind.CONSTANT_TENSOR, + ) + dq_input = builder.call_operator( + op=DQ_PER_TENSOR, args=(x_q, 0.02, -3, -128, 127, torch.int8) + ) + dq_weight = builder.call_operator( + op=DQ_PER_TENSOR, args=(w_q, 0.1, 0, -128, 127, torch.int8) + ) + conv = builder.call_operator( + op=torch.ops.aten.conv1d.default, args=(dq_input, dq_weight) + ) + q = builder.call_operator( + op=Q_PER_TENSOR, args=(conv, 0.5, 7, -128, 127, torch.int8) + ) + builder.output([q]) + ep = builder.get_program() + ep.graph_module.meta[EXPORTED_PROGRAM_META_KEY] = ep + conv_node = ep.graph.find_nodes( + op="call_function", target=torch.ops.aten.conv1d.default + )[0] + + fused = self._fuse(ep, conv_node) + + self.assertIs(fused.target, torch.ops.cadence.quantized_conv1d_ncl.per_tensor) + self.assertIsInstance(fused.kwargs["out_multiplier"], int) + + def test_bias_scale_is_input_scale_times_weight_scale(self) -> None: + scales = torch.tensor([0.1, 0.2, 0.05, 0.4]) + ep, conv_node = self._build_conv1d_program(scales) + + fused = self._fuse(ep, conv_node) + + bias_scale = self._kwarg_tensor(ep, fused, "bias_scale") + torch.testing.assert_close( + bias_scale, (scales * self.INPUT_SCALE).to(torch.float32) + ) + + def test_multiplier_and_shift_reconstruct_the_requantize_scale(self) -> None: + """out_multiplier/out_shift are a Q31 encoding of bias_scale / out_scale. + + Checked per channel, so a scalar that happened to be right for channel 0 + would not pass. + """ + scales = torch.tensor([0.1, 0.2, 0.05, 0.4]) + ep, conv_node = self._build_conv1d_program(scales) + + fused = self._fuse(ep, conv_node) + + out_multiplier = self._kwarg_tensor(ep, fused, "out_multiplier") + out_shift = self._kwarg_tensor(ep, fused, "out_shift") + self.assertEqual(out_multiplier.dtype, torch.int32) + self.assertEqual(out_shift.dtype, torch.int32) + self.assertEqual(out_multiplier.shape, scales.shape) + self.assertEqual(out_shift.shape, scales.shape) + + expected = (scales * self.INPUT_SCALE / self.OUT_SCALE).to(torch.float64) + reconstructed = ( + out_multiplier.to(torch.float64) + / (2**31) + * torch.pow(2.0, out_shift.to(torch.float64)) + ) + torch.testing.assert_close(reconstructed, expected, rtol=1e-4, atol=1e-4) + + def test_symmetric_weights_give_zero_zero_points(self) -> None: + """Symmetric per-channel quantization leaves zero_points unset.""" + scales = torch.tensor([0.1, 0.2, 0.05]) + ep, conv_node = self._build_conv1d_program(scales, weight_zero_points=None) + + fused = self._fuse(ep, conv_node) + + wzp = self._kwarg_tensor(ep, fused, "weight_zero_point") + self.assertEqual(wzp.shape, scales.shape) + self.assertTrue(torch.all(wzp == 0), f"expected zeros, got {wzp}") + + def test_affine_weight_zero_points_are_carried_through(self) -> None: + scales = torch.tensor([0.1, 0.2, 0.05]) + zps = torch.tensor([1, -2, 3], dtype=torch.int32) + ep, conv_node = self._build_conv1d_program(scales, weight_zero_points=zps) + + fused = self._fuse(ep, conv_node) + + wzp = self._kwarg_tensor(ep, fused, "weight_zero_point") + self.assertTrue(torch.equal(wzp.to(torch.int32), zps)) + + def test_missing_bias_becomes_a_zero_int32_bias(self) -> None: + scales = torch.tensor([0.1, 0.2, 0.05]) + ep, conv_node = self._build_conv1d_program(scales, with_bias=False) + + fused = self._fuse(ep, conv_node) + + bias = fused.args[2] + self.assertIsInstance(bias, torch.fx.Node) + bias_val = cast(torch.fx.Node, bias).meta["val"] + self.assertEqual(bias_val.dtype, torch.int32) + self.assertEqual(tuple(bias_val.shape), (scales.numel(),)) + + def test_depthwise_per_channel_routes_to_the_depthwise_op(self) -> None: + """Depthwise selection has to survive the per-channel overload swap. + + The granularity swap replaces the op, so if depthwise is decided after + it, a per-channel depthwise conv silently stays a dense conv. + """ + channels = 4 + scales = torch.rand(channels) * 0.1 + 0.01 + ep, conv_node = self._build_conv1d_program( + scales, in_channels=channels, groups=channels + ) + + fused = self._fuse(ep, conv_node) + + self.assertIs( + fused.target, + torch.ops.cadence.quantized_depthwise_conv1d_ncl.default, + ) + + def test_depthwise_per_tensor_still_routes_to_the_depthwise_op(self) -> None: + channels = 4 + builder = ProgramBuilder() + x_q = builder.placeholder( + "x_q", torch.randint(-128, 127, (1, channels, 6), dtype=torch.int8) + ) + w_q = builder.placeholder( + "w_q", + torch.randint(-127, 127, (channels, 1, 3), dtype=torch.int8), + input_kind=InputKind.CONSTANT_TENSOR, + ) + dq_input = builder.call_operator( + op=DQ_PER_TENSOR, args=(x_q, 0.02, -3, -128, 127, torch.int8) + ) + dq_weight = builder.call_operator( + op=DQ_PER_TENSOR, args=(w_q, 0.1, 0, -128, 127, torch.int8) + ) + conv = builder.call_operator( + op=torch.ops.aten.conv1d.default, + args=(dq_input, dq_weight, None, (1,), (0,), (1,), channels), + ) + q = builder.call_operator( + op=Q_PER_TENSOR, args=(conv, 0.5, 7, -128, 127, torch.int8) + ) + builder.output([q]) + ep = builder.get_program() + ep.graph_module.meta[EXPORTED_PROGRAM_META_KEY] = ep + conv_node = ep.graph.find_nodes( + op="call_function", target=torch.ops.aten.conv1d.default + )[0] + + fused = self._fuse(ep, conv_node) + + self.assertIs( + fused.target, + torch.ops.cadence.quantized_depthwise_conv1d_ncl.per_tensor, + ) + + +class TestPerChannelPatternRouting(unittest.TestCase): + """Which patterns are wired to accept per-channel weights and which aren't. + + ``fuse_conv`` and ``fuse_linear`` can absorb a per-channel weight dequant. + ``fuse_matmul`` and the mixed w8a32 paths cannot; their patterns must still + decline so the model falls back to float rather than raising a ``KeyError`` + when a scalar ``scale`` arg is missing. AddmmPattern is intentionally in the + declining set today: its fuse guards on ``DQ_PER_TENSOR`` directly, and + per-channel addmm currently falls back to float. + """ + + def _build_conv1d_relu_per_channel( + self, out_channels: int = 3, in_channels: int = 2, kernel_size: int = 3 + ) -> tuple[ExportedProgram, torch.fx.Node]: + # Distinct per-channel scales so a broadcast-of-channel-0 bug is visible. + weight_scales = torch.tensor([0.1, 0.25, 0.05])[:out_channels] + + builder = ProgramBuilder() + x_q = builder.placeholder( + "x_q", + torch.randint(-16, 16, (1, in_channels, 6), dtype=torch.int8), + ) + w_q = builder.placeholder( + "w_q", + torch.randint( + -16, 16, (out_channels, in_channels, kernel_size), dtype=torch.int8 + ), + input_kind=InputKind.CONSTANT_TENSOR, + ) + scales = builder.placeholder( + "scales", weight_scales, input_kind=InputKind.CONSTANT_TENSOR + ) + dq_input = builder.call_operator( + op=DQ_PER_TENSOR, args=(x_q, 0.02, -3, -128, 127, torch.int8) + ) + dq_weight = builder.call_operator( + op=DQ_PER_CHANNEL, + args=(w_q, scales, None, 0, -128, 127, torch.int8), + ) + conv = builder.call_operator( + op=torch.ops.aten.conv1d.default, args=(dq_input, dq_weight) + ) + relu = builder.call_operator(op=torch.ops.aten.relu.default, args=(conv,)) + q = builder.call_operator( + op=Q_PER_TENSOR, args=(relu, 0.5, 7, -128, 127, torch.int8) + ) + builder.output([q]) + ep = builder.get_program() + ep.graph_module.meta[EXPORTED_PROGRAM_META_KEY] = ep + conv_node = ep.graph.find_nodes( + op="call_function", target=torch.ops.aten.conv1d.default + )[0] + return ep, conv_node + + def test_conv_relu_pattern_fuses_per_channel_weights(self) -> None: + """ConvReluBase.fuse takes the per-channel path via is_weight_dq.""" + ep, conv_node = self._build_conv1d_relu_per_channel() + + fused = Conv1dReluPattern0().fuse(ep.graph_module, conv_node) + + self.assertIsNotNone(fused, "conv_relu should fuse a per-channel weight") + fused = cast(torch.fx.Node, fused) + self.assertIs(fused.target, torch.ops.cadence.quantized_conv1d_ncl.default) + for name in ("weight_zero_point", "out_multiplier", "out_shift", "bias_scale"): + arg = fused.kwargs[name] + self.assertIsInstance( + arg, torch.fx.Node, f"{name} should be a lifted constant" + ) + resolved = resolve_constant(ep.graph_module, arg) + self.assertIsNotNone(resolved) + # 3 output channels declared above. + self.assertEqual(cast(torch.Tensor, resolved).numel(), 3) + + def test_addmm_pattern_declines_per_channel_weight(self) -> None: + """AddmmPattern intentionally does not carry per-channel; it falls back. + + If this ever changes silently the pattern will start reading the missing + scalar ``scale`` arg through ``fuse_linear`` and blow up. + """ + builder = ProgramBuilder() + bias_q = builder.placeholder( + "bias_q", torch.zeros(3, dtype=torch.int32), + input_kind=InputKind.CONSTANT_TENSOR, + ) + x_q = builder.placeholder( + "x_q", torch.randint(-8, 8, (2, 4), dtype=torch.int8) + ) + w_q = builder.placeholder( + "w_q", + torch.randint(-8, 8, (4, 3), dtype=torch.int8), + input_kind=InputKind.CONSTANT_TENSOR, + ) + scales = builder.placeholder( + "scales", + torch.tensor([0.1, 0.25, 0.05]), + input_kind=InputKind.CONSTANT_TENSOR, + ) + dq_bias = builder.call_operator( + op=DQ_PER_TENSOR, args=(bias_q, 0.02, 0, -(2**31), 2**31 - 1, torch.int32) + ) + dq_input = builder.call_operator( + op=DQ_PER_TENSOR, args=(x_q, 0.1, 0, -128, 127, torch.int8) + ) + dq_weight = builder.call_operator( + op=DQ_PER_CHANNEL, + args=(w_q, scales, None, 0, -128, 127, torch.int8), + ) + addmm = builder.call_operator( + op=torch.ops.aten.addmm.default, args=(dq_bias, dq_input, dq_weight) + ) + q = builder.call_operator( + op=Q_PER_TENSOR, args=(addmm, 0.2, 0, -128, 127, torch.int8) + ) + builder.output([q]) + ep = builder.get_program() + ep.graph_module.meta[EXPORTED_PROGRAM_META_KEY] = ep + addmm_node = ep.graph.find_nodes( + op="call_function", target=torch.ops.aten.addmm.default + )[0] + + result = AddmmPattern().fuse(ep.graph_module, addmm_node) + + self.assertIsNone( + result, "AddmmPattern must decline per-channel to keep the float fallback" + ) + # And the graph should still contain the dq nodes untouched. + self.assertEqual( + len(ep.graph.find_nodes(op="call_function", target=DQ_PER_CHANNEL)), + 1, + ) + + +class TestFuseConvPerChannelRequantEncoding(unittest.TestCase): + """The Q31 encoding of the per-channel requant scale must be exact per channel. + + The existing TestFuseConvPerChannel checks the scale reconstruction round + trip; this class checks the encoding directly against + ``quantize_tensor_multiplier`` (which is what the runtime uses to interpret + the multiplier/shift pair), so a per-channel multiplier that happened to + reconstruct close to the right float value but was actually built from the + wrong index would be caught. + """ + + def test_encoded_multiplier_shift_matches_reference(self) -> None: + weight_scales = torch.tensor([0.11, 0.007, 0.34, 0.05]) + input_scale = 0.02 + out_scale = 0.31 + + builder = ProgramBuilder() + x_q = builder.placeholder( + "x_q", torch.randint(-8, 8, (1, 3, 6), dtype=torch.int8) + ) + w_q = builder.placeholder( + "w_q", + torch.randint(-8, 8, (4, 3, 3), dtype=torch.int8), + input_kind=InputKind.CONSTANT_TENSOR, + ) + scales = builder.placeholder( + "scales", weight_scales, input_kind=InputKind.CONSTANT_TENSOR + ) + dq_input = builder.call_operator( + op=DQ_PER_TENSOR, + args=(x_q, input_scale, 0, -128, 127, torch.int8), + ) + dq_weight = builder.call_operator( + op=DQ_PER_CHANNEL, args=(w_q, scales, None, 0, -128, 127, torch.int8) + ) + conv = builder.call_operator( + op=torch.ops.aten.conv1d.default, args=(dq_input, dq_weight) + ) + q = builder.call_operator( + op=Q_PER_TENSOR, args=(conv, out_scale, 0, -128, 127, torch.int8) + ) + builder.output([q]) + ep = builder.get_program() + ep.graph_module.meta[EXPORTED_PROGRAM_META_KEY] = ep + conv_node = ep.graph.find_nodes( + op="call_function", target=torch.ops.aten.conv1d.default + )[0] + + fused = Conv1dPattern().fuse(ep.graph_module, conv_node) + self.assertIsNotNone(fused) + fused = cast(torch.fx.Node, fused) + + out_multiplier = cast( + torch.Tensor, + resolve_constant(ep.graph_module, fused.kwargs["out_multiplier"]), + ) + out_shift = cast( + torch.Tensor, + resolve_constant(ep.graph_module, fused.kwargs["out_shift"]), + ) + expected_mult, expected_shift = quantize_tensor_multiplier( + weight_scales * input_scale / out_scale + ) + self.assertTrue( + torch.equal(out_multiplier, expected_mult.to(torch.int32)), + f"out_multiplier per channel: got {out_multiplier}, " + f"expected {expected_mult.to(torch.int32)}", + ) + self.assertTrue( + torch.equal(out_shift, expected_shift.to(torch.int32)), + f"out_shift per channel: got {out_shift}, " + f"expected {expected_shift.to(torch.int32)}", + ) diff --git a/backends/cadence/aot/tests/test_quantizer_ops.py b/backends/cadence/aot/tests/test_quantizer_ops.py index 7eef458ef4e..3bb3c9a5049 100644 --- a/backends/cadence/aot/tests/test_quantizer_ops.py +++ b/backends/cadence/aot/tests/test_quantizer_ops.py @@ -15,7 +15,14 @@ from executorch.backends.cadence.aot.quantizer import quantizer as quantizer_module from executorch.backends.cadence.aot.quantizer.patterns import ( AddmmPattern, + Conv1dBNReluPattern0, + Conv1dPattern, + Conv1dReluPattern0, Conv2dBNReluPattern0, + Conv2dPattern, + Conv2dReluPattern0, + LinearPattern, + QuantizationPattern, ) from executorch.backends.cadence.aot.quantizer.quantizer import ( CadenceAtenQuantizer, @@ -39,6 +46,11 @@ from executorch.exir.pass_base import NodeMetadata from parameterized import parameterized from torch._ops import OpOverload +from torchao.quantization.pt2e import PerChannelMinMaxObserver +from torchao.quantization.pt2e.quantizer import ( + DerivedQuantizationSpec, + QuantizationConfig, +) from torchao.quantization.pt2e.quantizer.quantizer import ( Q_ANNOTATION_KEY, QuantizationAnnotation, @@ -971,5 +983,399 @@ def test_fuse_declines_when_batchnorm_present(self) -> None: ) +PER_CHANNEL_WEIGHT_QSPEC: QuantizationSpec = QuantizationSpec( + dtype=torch.int8, + quant_min=-128, + quant_max=127, + qscheme=torch.per_channel_symmetric, + ch_axis=0, + is_dynamic=False, + observer_or_fake_quant_ctr=PerChannelMinMaxObserver, +) + + +class DerivedBiasSpecGranularityTest(unittest.TestCase): + """The derived bias spec has to follow the weight spec's granularity. + + Patterns hardcode the bias ``DerivedQuantizationSpec`` as per-tensor because + they cannot see the quantization config. With a per-channel weight the derived + bias scale is a vector, and convert_pt2e's per-tensor branch calls float() on + it. The quantizer holds both specs, so it is what reconciles them. + """ + + def _build_conv1d_with_bias_graph( + self, + ) -> tuple[torch.fx.GraphModule, torch.fx.Node]: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 3, 10)) + weight = builder.placeholder("weight", torch.randn(6, 3, 3)) + bias = builder.placeholder("bias", torch.randn(6)) + conv1d = builder.call_operator( + op=torch.ops.aten.conv1d.default, + args=(x, weight, bias), + meta=NodeMetadata( + {"source_fn_stack": [("conv1d", torch.ops.aten.conv1d.default)]} + ), + ) + builder.output([conv1d]) + gm = builder.get_graph_module() + conv_nodes = gm.graph.find_nodes( + op="call_function", target=torch.ops.aten.conv1d.default + ) + self.assertEqual(len(conv_nodes), 1) + return gm, conv_nodes[0] + + def _annotate_and_get_bias_spec( + self, weight_qspec: QuantizationSpec + ) -> DerivedQuantizationSpec: + gm, conv_node = self._build_conv1d_with_bias_graph() + config = QuantizationConfig( + qconfig_A8W8sym.input_activation, + qconfig_A8W8sym.output_activation, + weight_qspec, + None, + ) + CadenceAtenQuantizer(Conv1dPattern(), config).annotate(gm) + + annotation = conv_node.meta[Q_ANNOTATION_KEY] + bias_spec = annotation.input_qspec_map[conv_node.args[2]] + self.assertIsInstance(bias_spec, DerivedQuantizationSpec) + return bias_spec + + def test_per_channel_weight_makes_bias_spec_per_channel(self) -> None: + bias_spec = self._annotate_and_get_bias_spec(PER_CHANNEL_WEIGHT_QSPEC) + + self.assertEqual(bias_spec.qscheme, torch.per_channel_symmetric) + self.assertEqual(bias_spec.ch_axis, 0) + # The rest of the derived spec must survive the rewrite untouched. + self.assertEqual(bias_spec.dtype, torch.int32) + self.assertEqual(bias_spec.quant_min, -(2**31)) + self.assertEqual(bias_spec.quant_max, 2**31 - 1) + + def test_per_tensor_weight_leaves_bias_spec_per_tensor(self) -> None: + bias_spec = self._annotate_and_get_bias_spec(qconfig_A8W8sym.weight) + + self.assertEqual(bias_spec.qscheme, torch.per_tensor_affine) + + def test_weight_annotation_is_unchanged(self) -> None: + """Only the bias spec is rewritten; the weight keeps the configured spec.""" + gm, conv_node = self._build_conv1d_with_bias_graph() + config = QuantizationConfig( + qconfig_A8W8sym.input_activation, + qconfig_A8W8sym.output_activation, + PER_CHANNEL_WEIGHT_QSPEC, + None, + ) + CadenceAtenQuantizer(Conv1dPattern(), config).annotate(gm) + + annotation = conv_node.meta[Q_ANNOTATION_KEY] + self.assertEqual( + annotation.input_qspec_map[conv_node.args[1]], PER_CHANNEL_WEIGHT_QSPEC + ) + + +class _BiasSpecGranularityGraphs(unittest.TestCase): + """Graph builders used by DerivedBiasSpecGranularityAcrossPatternsTest. + + Each helper returns a graph module together with the fx node whose + ``Q_ANNOTATION_KEY`` will carry the derived bias spec after annotation. + """ + + def _linear_with_bias( + self, + ) -> tuple[torch.fx.GraphModule, torch.fx.Node, int]: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(2, 4)) + weight = builder.placeholder("weight", torch.randn(5, 4)) + bias = builder.placeholder("bias", torch.randn(5)) + linear = builder.call_operator( + op=torch.ops.aten.linear.default, + args=(x, weight, bias), + meta=NodeMetadata( + {"source_fn_stack": [("linear", torch.ops.aten.linear.default)]} + ), + ) + builder.output([linear]) + gm = builder.get_graph_module() + node = gm.graph.find_nodes( + op="call_function", target=torch.ops.aten.linear.default + )[0] + # linear(input, weight, bias) -> bias is arg index 2 + return gm, node, 2 + + def _addmm( + self, + ) -> tuple[torch.fx.GraphModule, torch.fx.Node, int]: + builder = GraphBuilder() + bias = builder.placeholder("bias", torch.randn(5)) + mat1 = builder.placeholder("mat1", torch.randn(2, 4)) + mat2 = builder.placeholder("mat2", torch.randn(4, 5)) + addmm = builder.call_operator( + op=torch.ops.aten.addmm.default, + args=(bias, mat1, mat2), + meta=NodeMetadata( + {"source_fn_stack": [("addmm", torch.ops.aten.addmm.default)]} + ), + ) + builder.output([addmm]) + gm = builder.get_graph_module() + node = gm.graph.find_nodes( + op="call_function", target=torch.ops.aten.addmm.default + )[0] + # addmm(bias, mat1, mat2) -> bias is arg index 0 + return gm, node, 0 + + def _conv2d_with_bias( + self, + ) -> tuple[torch.fx.GraphModule, torch.fx.Node, int]: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 3, 8, 8)) + weight = builder.placeholder("weight", torch.randn(6, 3, 3, 3)) + bias = builder.placeholder("bias", torch.randn(6)) + conv2d = builder.call_operator( + op=torch.ops.aten.conv2d.default, + args=(x, weight, bias), + meta=NodeMetadata( + {"source_fn_stack": [("conv2d", torch.ops.aten.conv2d.default)]} + ), + ) + builder.output([conv2d]) + gm = builder.get_graph_module() + node = gm.graph.find_nodes( + op="call_function", target=torch.ops.aten.conv2d.default + )[0] + return gm, node, 2 + + def _conv1d_relu_with_bias( + self, + ) -> tuple[torch.fx.GraphModule, torch.fx.Node, int]: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 3, 10)) + weight = builder.placeholder("weight", torch.randn(6, 3, 3)) + bias = builder.placeholder("bias", torch.randn(6)) + conv = builder.call_operator( + op=torch.ops.aten.conv1d.default, + args=(x, weight, bias), + meta=NodeMetadata( + {"source_fn_stack": [("conv1d", torch.ops.aten.conv1d.default)]} + ), + ) + relu = builder.call_operator( + op=torch.ops.aten.relu.default, + args=(conv,), + meta=NodeMetadata( + {"source_fn_stack": [("relu", torch.ops.aten.relu.default)]} + ), + ) + builder.output([relu]) + gm = builder.get_graph_module() + conv_node = gm.graph.find_nodes( + op="call_function", target=torch.ops.aten.conv1d.default + )[0] + return gm, conv_node, 2 + + def _conv1d_bn_relu_with_bias( + self, + ) -> tuple[torch.fx.GraphModule, torch.fx.Node, int]: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 3, 10)) + weight = builder.placeholder("weight", torch.randn(6, 3, 3)) + bias = builder.placeholder("bias", torch.randn(6)) + bn_weight = builder.placeholder("bn_weight", torch.randn(6)) + bn_bias = builder.placeholder("bn_bias", torch.randn(6)) + bn_mean = builder.placeholder("bn_mean", torch.randn(6)) + bn_var = builder.placeholder("bn_var", torch.abs(torch.randn(6))) + conv = builder.call_operator( + op=torch.ops.aten.conv1d.default, + args=(x, weight, bias), + meta=NodeMetadata( + {"source_fn_stack": [("conv1d", torch.ops.aten.conv1d.default)]} + ), + ) + bn = builder.call_operator( + op=torch.ops.aten.batch_norm.default, + args=(conv, bn_weight, bn_bias, bn_mean, bn_var, False, 0.1, 1e-5, False), + meta=NodeMetadata( + { + "source_fn_stack": [ + ("batch_norm", torch.ops.aten.batch_norm.default) + ] + } + ), + ) + relu = builder.call_operator( + op=torch.ops.aten.relu.default, + args=(bn,), + meta=NodeMetadata( + {"source_fn_stack": [("relu", torch.ops.aten.relu.default)]} + ), + ) + builder.output([relu]) + gm = builder.get_graph_module() + conv_node = gm.graph.find_nodes( + op="call_function", target=torch.ops.aten.conv1d.default + )[0] + return gm, conv_node, 2 + + +class DerivedBiasSpecGranularityAcrossPatternsTest(_BiasSpecGranularityGraphs): + """Every pattern with a derived bias spec must follow the weight granularity. + + quantizer.py aligns the bias spec once in ``CadenceAtenQuantizer.annotate``, + not in each pattern, so a regression here would silently apply per-tensor + bias to only some patterns. Six pattern families declare a derived bias spec + (Addmm, Linear, Conv1d, Conv2d, ConvRelu*, ConvBNRelu*); each is checked. + """ + + def _annotate_and_get_bias_spec( + self, + pattern: QuantizationPattern, + weight_qspec: QuantizationSpec, + graph_builder: Callable[ + [], tuple[torch.fx.GraphModule, torch.fx.Node, int] + ], + ) -> DerivedQuantizationSpec: + gm, anchor_node, bias_idx = graph_builder() + config = QuantizationConfig( + qconfig_A8W8sym.input_activation, + qconfig_A8W8sym.output_activation, + weight_qspec, + None, + ) + CadenceAtenQuantizer(pattern, config).annotate(gm) + + annotation = anchor_node.meta[Q_ANNOTATION_KEY] + bias_spec = annotation.input_qspec_map[anchor_node.args[bias_idx]] + self.assertIsInstance( + bias_spec, + DerivedQuantizationSpec, + f"{type(pattern).__name__}: expected a derived bias spec", + ) + return bias_spec + + def _cases( + self, + ) -> list[ + tuple[ + str, + QuantizationPattern, + Callable[[], tuple[torch.fx.GraphModule, torch.fx.Node, int]], + ] + ]: + return [ + ("linear", LinearPattern(), self._linear_with_bias), + ("addmm", AddmmPattern(), self._addmm), + ("conv2d", Conv2dPattern(), self._conv2d_with_bias), + ("conv1d_relu", Conv1dReluPattern0(), self._conv1d_relu_with_bias), + ("conv2d_relu", Conv2dReluPattern0(), self._conv2d_relu_with_bias_graph), + ("conv1d_bn_relu", Conv1dBNReluPattern0(), self._conv1d_bn_relu_with_bias), + ( + "conv2d_bn_relu", + Conv2dBNReluPattern0(), + self._conv2d_bn_relu_with_bias_graph, + ), + ] + + def _conv2d_relu_with_bias_graph( + self, + ) -> tuple[torch.fx.GraphModule, torch.fx.Node, int]: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 3, 8, 8)) + weight = builder.placeholder("weight", torch.randn(6, 3, 3, 3)) + bias = builder.placeholder("bias", torch.randn(6)) + conv = builder.call_operator( + op=torch.ops.aten.conv2d.default, + args=(x, weight, bias), + meta=NodeMetadata( + {"source_fn_stack": [("conv2d", torch.ops.aten.conv2d.default)]} + ), + ) + relu = builder.call_operator( + op=torch.ops.aten.relu.default, + args=(conv,), + meta=NodeMetadata( + {"source_fn_stack": [("relu", torch.ops.aten.relu.default)]} + ), + ) + builder.output([relu]) + gm = builder.get_graph_module() + conv_node = gm.graph.find_nodes( + op="call_function", target=torch.ops.aten.conv2d.default + )[0] + return gm, conv_node, 2 + + def _conv2d_bn_relu_with_bias_graph( + self, + ) -> tuple[torch.fx.GraphModule, torch.fx.Node, int]: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 3, 8, 8)) + weight = builder.placeholder("weight", torch.randn(6, 3, 3, 3)) + bias = builder.placeholder("bias", torch.randn(6)) + bn_weight = builder.placeholder("bn_weight", torch.randn(6)) + bn_bias = builder.placeholder("bn_bias", torch.randn(6)) + bn_mean = builder.placeholder("bn_mean", torch.randn(6)) + bn_var = builder.placeholder("bn_var", torch.abs(torch.randn(6))) + conv = builder.call_operator( + op=torch.ops.aten.conv2d.default, + args=(x, weight, bias), + meta=NodeMetadata( + {"source_fn_stack": [("conv2d", torch.ops.aten.conv2d.default)]} + ), + ) + bn = builder.call_operator( + op=torch.ops.aten.batch_norm.default, + args=(conv, bn_weight, bn_bias, bn_mean, bn_var, False, 0.1, 1e-5, False), + meta=NodeMetadata( + { + "source_fn_stack": [ + ("batch_norm", torch.ops.aten.batch_norm.default) + ] + } + ), + ) + relu = builder.call_operator( + op=torch.ops.aten.relu.default, + args=(bn,), + meta=NodeMetadata( + {"source_fn_stack": [("relu", torch.ops.aten.relu.default)]} + ), + ) + builder.output([relu]) + gm = builder.get_graph_module() + conv_node = gm.graph.find_nodes( + op="call_function", target=torch.ops.aten.conv2d.default + )[0] + return gm, conv_node, 2 + + def test_per_channel_weight_makes_bias_spec_per_channel(self) -> None: + for name, pattern, builder_fn in self._cases(): + with self.subTest(pattern=name): + bias_spec = self._annotate_and_get_bias_spec( + pattern, PER_CHANNEL_WEIGHT_QSPEC, builder_fn + ) + self.assertEqual( + bias_spec.qscheme, + torch.per_channel_symmetric, + f"{name}: bias qscheme was not aligned to weight granularity", + ) + self.assertEqual(bias_spec.ch_axis, 0) + # The rest of the derived spec must survive untouched. + self.assertEqual(bias_spec.dtype, torch.int32) + self.assertEqual(bias_spec.quant_min, -(2**31)) + self.assertEqual(bias_spec.quant_max, 2**31 - 1) + + def test_per_tensor_weight_leaves_bias_spec_per_tensor(self) -> None: + for name, pattern, builder_fn in self._cases(): + with self.subTest(pattern=name): + bias_spec = self._annotate_and_get_bias_spec( + pattern, qconfig_A8W8sym.weight, builder_fn + ) + self.assertEqual( + bias_spec.qscheme, + torch.per_tensor_affine, + f"{name}: per-tensor weight was rewritten as per-channel", + ) + + if __name__ == "__main__": unittest.main() diff --git a/backends/cadence/aot/tests/test_ref_implementations.py b/backends/cadence/aot/tests/test_ref_implementations.py index 005bf9d85bd..00dec41aebf 100644 --- a/backends/cadence/aot/tests/test_ref_implementations.py +++ b/backends/cadence/aot/tests/test_ref_implementations.py @@ -13,6 +13,7 @@ import numpy as np import torch +from executorch.backends.cadence.aot.quantizer.utils import quantize_tensor_multiplier from executorch.backends.cadence.aot.typing_stubs import expand from executorch.backends.cadence.aot.utils import is_depthwise_conv @@ -493,6 +494,483 @@ def test_quantized_linear( f"Values don't match: got {output}, expected {expected_output}", ) + def test_quantized_linear_per_channel(self) -> None: + """Per-channel requant applies a distinct scale to each output row. + + Two output channels are given deliberately different multipliers, so a + per-tensor implementation (which would apply channel 0's scale to both) + cannot produce this result. + """ + src = torch.tensor([[4, 8]], dtype=torch.int8) + weight = torch.tensor([[2, 0], [0, 2]], dtype=torch.int8) + bias = torch.zeros(2, dtype=torch.int32) + # channel 0 requantizes at half the rate of channel 1 + out_multiplier = torch.tensor([1 << 30, 1 << 30], dtype=torch.int32) + out_shift = torch.tensor([0, 1], dtype=torch.int32) + weight_zero_point = torch.zeros(2, dtype=torch.int32) + + output = torch.ops.cadence.quantized_linear( + src, + weight, + bias, + 0, # in_zero_point + weight_zero_point, + out_multiplier, + out_shift, + 0, # out_zero_point + typing.cast(torch.Tensor, None), + ) + + # accumulators are [8, 16]; requant scale is -m/2^31 * 2^shift, so + # channel 0 scales by -0.5 and channel 1 by -1.0 + expected = torch.tensor([[-4, -16]], dtype=torch.int8) + self.assertEqual(output.dtype, torch.int8) + self.assertTrue( + torch.equal(output, expected), + f"Values don't match: got {output}, expected {expected}", + ) + + def test_quantized_linear_per_channel_zero_point(self) -> None: + """A per-channel weight zero point is applied down output rows.""" + src = torch.tensor([[1, 1]], dtype=torch.int8) + weight = torch.zeros((2, 2), dtype=torch.int8) + bias = torch.zeros(2, dtype=torch.int32) + # subtracting these zero points turns the zero weights into -1 and -2 + weight_zero_point = torch.tensor([1, 2], dtype=torch.int32) + out_multiplier = torch.tensor([1 << 30, 1 << 30], dtype=torch.int32) + out_shift = torch.tensor([1, 1], dtype=torch.int32) + + output = torch.ops.cadence.quantized_linear( + src, + weight, + bias, + 0, + weight_zero_point, + out_multiplier, + out_shift, + 0, + typing.cast(torch.Tensor, None), + ) + + # accumulators are [-2, -4], requant scale is -1.0 for both channels + expected = torch.tensor([[2, 4]], dtype=torch.int8) + self.assertTrue( + torch.equal(output, expected), + f"Values don't match: got {output}, expected {expected}", + ) + + def test_quantized_conv2d_nchw_per_channel_bias_scale(self) -> None: + """A per-output-channel bias_scale is applied down the channel axis. + + Both output channels share an accumulator but are given different bias + scales, so a per-tensor implementation (which would use channel 0's + scale for both) cannot produce this result. + """ + # [N, C, H, W] = [1, 1, 1, 3] + input_tensor = torch.tensor([[[[1, 2, 3]]]], dtype=torch.int8) + # [OC, IC, KH, KW] = [2, 1, 1, 2], both channels sum two neighbours + weight = torch.tensor([[[[1, 1]]], [[[1, 1]]]], dtype=torch.int8) + bias = torch.zeros(2, dtype=torch.int32) + weight_zero_point = torch.zeros(2, dtype=torch.int32) + bias_scale = torch.tensor([1.0, 2.0]) + out_multiplier = torch.tensor([1 << 30, 1 << 30], dtype=torch.int32) + out_shift = torch.zeros(2, dtype=torch.int32) + + output = torch.ops.cadence.quantized_conv2d_nchw( + input_tensor, + weight, + bias, + (1, 1), + (0, 0), + (1, 1), + 1, # groups + 0, # in_zero_point + weight_zero_point, + bias_scale, + 1.0, # output_scale + 0, # output_zero_point + out_multiplier, + out_shift, + ) + + # accumulators are [3, 5] on both channels; channel 1 is scaled by 2 + expected = torch.tensor([[[[3, 5]], [[6, 10]]]], dtype=torch.int8) + self.assertEqual(output.dtype, torch.int8) + self.assertTrue( + torch.equal(output, expected), + f"Values don't match: got {output}, expected {expected}", + ) + + def test_quantized_conv1d_ncl_per_channel_weight_zero_point(self) -> None: + """A per-channel weight zero point is subtracted per output channel.""" + # [N, C, L] = [1, 1, 2] + input_tensor = torch.tensor([[[1, 1]]], dtype=torch.int8) + # [OC, IC, K] = [2, 1, 2], all-zero weights so only the zero point acts + weight = torch.zeros((2, 1, 2), dtype=torch.int8) + bias = torch.zeros(2, dtype=torch.int32) + weight_zero_point = torch.tensor([1, 2], dtype=torch.int32) + bias_scale = torch.ones(2) + out_multiplier = torch.tensor([1 << 30, 1 << 30], dtype=torch.int32) + out_shift = torch.zeros(2, dtype=torch.int32) + + output = torch.ops.cadence.quantized_conv1d_ncl( + input_tensor, + weight, + bias, + (1,), + (0,), + (1,), + 1, + 0, + weight_zero_point, + bias_scale, + 1.0, + 0, + out_multiplier, + out_shift, + ) + + # weights become -1 and -2 after subtracting the zero points + expected = torch.tensor([[[-2], [-4]]], dtype=torch.int8) + self.assertTrue( + torch.equal(output, expected), + f"Values don't match: got {output}, expected {expected}", + ) + + def test_quantized_conv1d_nlc_per_channel_matches_ncl(self) -> None: + """The channel-last variant must agree with NCL on the same data.""" + input_ncl = torch.randint(-8, 8, (1, 3, 5), dtype=torch.int8) + weight_ncl = torch.randint(-4, 4, (2, 3, 2), dtype=torch.int8) + bias = torch.zeros(2, dtype=torch.int32) + weight_zero_point = torch.tensor([1, -1], dtype=torch.int32) + bias_scale = torch.tensor([0.02, 0.05]) + out_multiplier = torch.tensor([1 << 30, 1 << 30], dtype=torch.int32) + out_shift = torch.zeros(2, dtype=torch.int32) + common = (bias, (1,), (0,), (1,), 1, 0, weight_zero_point, bias_scale, 0.1, 0) + + out_ncl = torch.ops.cadence.quantized_conv1d_ncl( + input_ncl, weight_ncl, *common, out_multiplier, out_shift + ) + out_nlc = torch.ops.cadence.quantized_conv1d_nlc( + input_ncl.permute(0, 2, 1).contiguous(), + weight_ncl.permute(0, 2, 1).contiguous(), + *common, + out_multiplier, + out_shift, + ) + + self.assertTrue( + torch.equal(out_nlc.permute(0, 2, 1).contiguous(), out_ncl), + f"NLC and NCL disagree: {out_nlc} vs {out_ncl}", + ) + + def test_quantized_conv_per_channel_beats_flattened_qparams(self) -> None: + """Per-channel must actually track a channel-skewed weight distribution. + + One output channel has a far larger dynamic range than the other. Reusing + channel 0's scale for both loses the small channel, so this fails if the + qparam vectors are ever collapsed to a scalar. + """ + torch.manual_seed(0) + in_channels, out_channels, length, kernel = 3, 2, 8, 3 + input_scale = 0.05 + out_scale = 0.05 + + input_q = torch.randint(-32, 32, (1, in_channels, length), dtype=torch.int8) + weight_q = torch.randint( + -100, 100, (out_channels, in_channels, kernel), dtype=torch.int8 + ) + bias = torch.zeros(out_channels, dtype=torch.int32) + # channel 1's weights are 50x smaller in real terms than channel 0's + weight_scales = torch.tensor([0.5, 0.01]) + bias_scales = weight_scales * input_scale + weight_zero_point = torch.zeros(out_channels, dtype=torch.int32) + multiplier, shift = quantize_tensor_multiplier(bias_scales / out_scale) + + reference = torch.nn.functional.conv1d( + input_q.float() * input_scale, + weight_q.float() * weight_scales.view(-1, 1, 1), + ) + conv_args = ((1,), (0,), (1,), 1, 0) + + per_channel = torch.ops.cadence.quantized_conv1d_ncl( + input_q, + weight_q, + bias, + *conv_args, + weight_zero_point, + bias_scales, + out_scale, + 0, + multiplier.to(torch.int32), + shift.to(torch.int32), + ) + flattened = bias_scales[0].expand(out_channels).contiguous() + per_tensor = torch.ops.cadence.quantized_conv1d_ncl( + input_q, + weight_q, + bias, + *conv_args, + weight_zero_point, + flattened, + out_scale, + 0, + multiplier.to(torch.int32), + shift.to(torch.int32), + ) + + pc_err = (per_channel.float() * out_scale - reference).abs().mean() + pt_err = (per_tensor.float() * out_scale - reference).abs().mean() + self.assertLess( + pc_err, pt_err, f"per-channel ({pc_err}) should beat flattened ({pt_err})" + ) + + def test_quantized_conv_per_channel_rejects_mismatched_length(self) -> None: + """A qparam vector that does not match the output channels is a bug.""" + with self.assertRaisesRegex(ValueError, "expected 1 or 2"): + torch.ops.cadence.quantized_conv1d_ncl( + torch.zeros((1, 1, 4), dtype=torch.int8), + torch.zeros((2, 1, 2), dtype=torch.int8), + torch.zeros(2, dtype=torch.int32), + (1,), + (0,), + (1,), + 1, + 0, + torch.zeros(2, dtype=torch.int32), + torch.ones(3), # three scales for two output channels + 1.0, + 0, + torch.tensor([1 << 30, 1 << 30], dtype=torch.int32), + torch.zeros(2, dtype=torch.int32), + ) + + def test_quantized_conv2d_nhwc_per_channel_bias_scale(self) -> None: + """Channel-last convs have to broadcast bias_scale over the channel axis. + + The NHWC accumulator is [N, H, W, C], not [N, C, H, W], so a helper that + reshapes bias_scale for NCHW would pick the wrong axis and either raise + or silently apply the wrong scale to every channel. Distinct scales per + output channel make that visible. + """ + # NHWC: [N=1, H=1, W=3, C=1]; single input channel so accumulators are + # trivially equal across the two output channels, and any per-channel + # scale difference in the result must come from bias_scale. + input_nhwc = torch.tensor([[[[1], [2], [3]]]], dtype=torch.int8) + # NHWC weight layout is [OC, KH, KW, IC] = [2, 1, 2, 1], all-ones so + # each output position sums two neighbouring input positions. + weight = torch.tensor( + [[[[1], [1]]], [[[1], [1]]]], dtype=torch.int8 + ) + bias = torch.zeros(2, dtype=torch.int32) + weight_zero_point = torch.zeros(2, dtype=torch.int32) + bias_scale = torch.tensor([1.0, 3.0]) + out_multiplier = torch.tensor([1 << 30, 1 << 30], dtype=torch.int32) + out_shift = torch.zeros(2, dtype=torch.int32) + + output = torch.ops.cadence.quantized_conv2d_nhwc( + input_nhwc, + weight, + bias, + (1, 1), + (0, 0), + (1, 1), + 1, # groups + 0, # in_zero_point + weight_zero_point, + bias_scale, + 1.0, # output_scale + 0, # output_zero_point + out_multiplier, + out_shift, + ) + + # Per position accumulators are [3, 5]; channel 1 is scaled by 3. + # Expected NHWC layout [N=1, H=1, W=2, C=2]: + # pos 0 -> [3*1, 3*3] = [3, 9] + # pos 1 -> [5*1, 5*3] = [5, 15] + expected = torch.tensor([[[[3, 9], [5, 15]]]], dtype=torch.int8) + self.assertEqual(output.dtype, torch.int8) + self.assertTrue( + torch.equal(output, expected), + f"Values don't match: got {output}, expected {expected}", + ) + + def test_quantized_depthwise_conv1d_ncl_per_channel(self) -> None: + """The new ``.default`` overload for depthwise NCL takes per-channel qparams. + + Depthwise means groups == channels, so each output channel has its own + (weight_zero_point, bias_scale). Distinct values per channel means a + broadcast of channel 0 would give the wrong answer for channels 1..N. + """ + channels = 3 + # NCL: [N=1, C=3, L=2] + input_tensor = torch.tensor( + [[[1, 1], [1, 1], [1, 1]]], dtype=torch.int8 + ) + # Depthwise weights: [OC=3, IC/groups=1, K=2], all-zero so only the + # weight_zero_point drives the accumulator. + weight = torch.zeros((channels, 1, 2), dtype=torch.int8) + bias = torch.zeros(channels, dtype=torch.int32) + # Subtracting these zero points turns 0-valued weights into -1, -2, -3. + weight_zero_point = torch.tensor([1, 2, 3], dtype=torch.int32) + bias_scale = torch.tensor([1.0, 2.0, 3.0]) + out_multiplier = torch.tensor( + [1 << 30, 1 << 30, 1 << 30], dtype=torch.int32 + ) + out_shift = torch.zeros(channels, dtype=torch.int32) + + output = torch.ops.cadence.quantized_depthwise_conv1d_ncl( + input_tensor, + weight, + bias, + (1,), + (0,), + (1,), + channels, # groups == channels for depthwise + 0, + weight_zero_point, + bias_scale, + 1.0, + 0, + out_multiplier, + out_shift, + ) + + # Each output channel c contributes acc = sum(-wzp[c] * input) = -2*wzp[c], + # then times bias_scale[c]. Channels: [-2*1*1, -2*2*2, -2*3*3] = [-2, -8, -18] + expected = torch.tensor([[[-2], [-8], [-18]]], dtype=torch.int8) + self.assertEqual(output.dtype, torch.int8) + self.assertTrue( + torch.equal(output, expected), + f"Values don't match: got {output}, expected {expected}", + ) + + def test_quantized_depthwise_conv1d_nlc_per_channel_matches_ncl(self) -> None: + """NLC depthwise must agree with NCL depthwise on the same data.""" + torch.manual_seed(0) + channels = 3 + length = 5 + kernel = 2 + + input_ncl = torch.randint(-4, 4, (1, channels, length), dtype=torch.int8) + # [OC, IC/groups=1, K] + weight_ncl = torch.randint(-3, 3, (channels, 1, kernel), dtype=torch.int8) + bias = torch.zeros(channels, dtype=torch.int32) + weight_zero_point = torch.tensor([1, -1, 2], dtype=torch.int32) + bias_scale = torch.tensor([0.02, 0.05, 0.1]) + out_multiplier = torch.tensor( + [1 << 30, 1 << 30, 1 << 30], dtype=torch.int32 + ) + out_shift = torch.zeros(channels, dtype=torch.int32) + common = ( + bias, + (1,), + (0,), + (1,), + channels, # groups + 0, + weight_zero_point, + bias_scale, + 0.1, + 0, + ) + + out_ncl = torch.ops.cadence.quantized_depthwise_conv1d_ncl( + input_ncl, weight_ncl, *common, out_multiplier, out_shift + ) + # NLC layouts: input [N, L, C], weight [OC, K, IC/groups]. + out_nlc = torch.ops.cadence.quantized_depthwise_conv1d_nlc( + input_ncl.permute(0, 2, 1).contiguous(), + weight_ncl.permute(0, 2, 1).contiguous(), + *common, + out_multiplier, + out_shift, + ) + + self.assertTrue( + torch.equal(out_nlc.permute(0, 2, 1).contiguous(), out_ncl), + f"NLC and NCL depthwise disagree: {out_nlc} vs {out_ncl}", + ) + + def test_quantized_conv2d_nhwc_per_channel_beats_flattened_qparams(self) -> None: + """The whole point of per-channel: a channel-skewed conv2d needs it. + + Same reasoning as the per-channel-beats-flattened conv1d test, but along + the NHWC path so both channel-first and channel-last go through the + per-channel branch of _broadcast_over_channels with distinct scales. + """ + torch.manual_seed(0) + in_channels, out_channels = 3, 2 + input_scale = 0.05 + out_scale = 0.05 + + input_nchw = torch.randint( + -32, 32, (1, in_channels, 5, 5), dtype=torch.int8 + ) + weight = torch.randint( + -100, 100, (out_channels, in_channels, 3, 3), dtype=torch.int8 + ) + bias = torch.zeros(out_channels, dtype=torch.int32) + weight_scales = torch.tensor([0.5, 0.01]) + bias_scales = weight_scales * input_scale + weight_zero_point = torch.zeros(out_channels, dtype=torch.int32) + multiplier, shift = quantize_tensor_multiplier(bias_scales / out_scale) + + reference = torch.nn.functional.conv2d( + input_nchw.float() * input_scale, + weight.float() * weight_scales.view(-1, 1, 1, 1), + ) + # NHWC needs input [N, H, W, C] and weight [OC, H, W, IC]. + input_nhwc = input_nchw.permute(0, 2, 3, 1).contiguous() + weight_nhwc = weight.permute(0, 2, 3, 1).contiguous() + + per_channel = torch.ops.cadence.quantized_conv2d_nhwc( + input_nhwc, + weight_nhwc, + bias, + (1, 1), + (0, 0), + (1, 1), + 1, + 0, + weight_zero_point, + bias_scales, + out_scale, + 0, + multiplier.to(torch.int32), + shift.to(torch.int32), + ) + flattened = bias_scales[0].expand(out_channels).contiguous() + per_tensor = torch.ops.cadence.quantized_conv2d_nhwc( + input_nhwc, + weight_nhwc, + bias, + (1, 1), + (0, 0), + (1, 1), + 1, + 0, + weight_zero_point, + flattened, + out_scale, + 0, + multiplier.to(torch.int32), + shift.to(torch.int32), + ) + + # NHWC output is [N, H, W, C]; the reference is NCHW, so bring both to + # NCHW before comparing. + per_channel_nchw = per_channel.permute(0, 3, 1, 2).contiguous() + per_tensor_nchw = per_tensor.permute(0, 3, 1, 2).contiguous() + pc_err = (per_channel_nchw.float() * out_scale - reference).abs().mean() + pt_err = (per_tensor_nchw.float() * out_scale - reference).abs().mean() + self.assertLess( + pc_err, + pt_err, + f"per-channel ({pc_err}) should beat flattened ({pt_err})", + ) + @expand( [ # Test case 1: Simple case with int8, zero mean input @@ -3631,3 +4109,4 @@ def test_quantized_conv1d_nlc_with_padding(self) -> None: # With padding=1, output length = (3 + 2*1 - 3) / 1 + 1 = 3 self.assertEqual(output.shape, (batch_size, length, out_channels)) self.assertEqual(output.dtype, torch.int8) + diff --git a/backends/cadence/aot/tests/test_replace_ops_passes.py b/backends/cadence/aot/tests/test_replace_ops_passes.py index a5a712d6349..e7ecb26304d 100644 --- a/backends/cadence/aot/tests/test_replace_ops_passes.py +++ b/backends/cadence/aot/tests/test_replace_ops_passes.py @@ -52,6 +52,7 @@ ReplaceWhereWithFullArgsWithWhereScalar, ) +from executorch.backends.cadence.aot.quantizer.utils import quantize_tensor_multiplier from executorch.backends.cadence.aot.typing_stubs import expand from executorch.backends.test.graph_builder import GraphBuilder, single_op_builder from executorch.exir.dialects._ops import ops as exir_ops @@ -1445,6 +1446,126 @@ def test_replace_quantized_conv1d_ncl_with_linear(self) -> None: 3, ) + @torch.no_grad() + def test_replace_per_channel_quantized_conv1d_ncl_with_linear(self) -> None: + """A trivial per-channel quantized conv1d collapses to quantized_linear. + + The per-channel op carries its qparams as tensors, so the pass cannot + recompute out_multiplier/out_shift from bias_scale / out_scale the way it + does for scalars. It has to forward the conv's existing qparam nodes, and + those have to stay per-output-channel vectors. + """ + in_channels = 3 + out_channels = 4 + kernel_size = 2 + x = torch.randint(-10, 10, (1, in_channels, kernel_size), dtype=torch.int8) + w = torch.randint( + -5, 5, (out_channels, in_channels, kernel_size), dtype=torch.int8 + ) + b = torch.zeros(out_channels, dtype=torch.int32) + weight_zero_point = torch.zeros(out_channels, dtype=torch.int32) + bias_scale = torch.full((out_channels,), 0.001, dtype=torch.float32) + # Deliberately distinct per channel so a collapse to a scalar is visible. + out_multiplier = torch.tensor( + [1073741824, 1181116006, 1288490188, 1395864371], dtype=torch.int32 + ) + out_shift = torch.tensor([0, -1, -2, -3], dtype=torch.int32) + placeholders = ( + x, + w, + b, + weight_zero_point, + bias_scale, + out_multiplier, + out_shift, + ) + args = ( + x, + w, + b, + [1], + [0], + [1], + 1, + 0, + weight_zero_point, + bias_scale, + 1.0, + 0, + out_multiplier, + out_shift, + ) + original_gm = single_op_builder( + placeholders=placeholders, + op=exir_ops.edge.cadence.quantized_conv1d_ncl.default, + args=args, + ) + + conv_nodes = original_gm.graph.find_nodes( + op="call_function", + target=exir_ops.edge.cadence.quantized_conv1d_ncl.default, + ) + self.assertEqual(len(conv_nodes), 1) + # Record which graph inputs carried the qparams so we can check the + # linear ends up reading the very same ones. + conv_args = conv_nodes[0].args + wzp_name = cast(torch.fx.Node, conv_args[8]).name + multiplier_name = cast(torch.fx.Node, conv_args[12]).name + shift_name = cast(torch.fx.Node, conv_args[13]).name + + p = ReplaceTrivialConvWithLinear() + result = cast(PassResult, p(original_gm)) + self.assertTrue(result.modified) + graph_after_passes = result.graph_module + + self.assertEqual( + count_node( + graph_after_passes, + exir_ops.edge.cadence.quantized_conv1d_ncl.default, + ), + 0, + ) + self.assertEqual( + count_node( + graph_after_passes, + exir_ops.edge.cadence.quantized_linear.default, + ), + 1, + ) + # 3 view_copy ops: weight reshape, input reshape, output reshape + self.assertEqual( + count_node( + graph_after_passes, + exir_ops.edge.aten.view_copy.default, + ), + 3, + ) + + linear_nodes = graph_after_passes.graph.find_nodes( + op="call_function", + target=exir_ops.edge.cadence.quantized_linear.default, + ) + self.assertEqual(len(linear_nodes), 1) + linear_args = linear_nodes[0].args + # quantized_linear(src, weight, bias, src_zero_point, weight_zero_point, + # out_multiplier, out_shift, out_zero_point, offset) + for idx, expected_name in ( + (4, wzp_name), + (5, multiplier_name), + (6, shift_name), + ): + arg = linear_args[idx] + self.assertIsInstance( + arg, + torch.fx.Node, + f"per-channel qparam at arg {idx} was flattened to a scalar", + ) + self.assertEqual(cast(torch.fx.Node, arg).name, expected_name) + self.assertEqual( + cast(torch.fx.Node, arg).meta["val"].shape, + torch.Size([out_channels]), + ) + @torch.no_grad() def test_replace_quantized_conv1d_nlc_with_linear(self) -> None: """Test that a trivial quantized conv1d NLC (in_length == kernel_length) is @@ -3087,6 +3208,499 @@ def test_cat_insert_transpose(self) -> None: ) +class TestPerChannelConvLowering(unittest.TestCase): + """Per-channel convs must survive the conv lowering exits. + + Per-channel qparams ride on the tensor-qparam (`.default`) overload. Passes + that enumerate only `.per_tensor` silently skip these nodes, so a per-channel + model quietly loses the channel-last layout or the im2row path. + """ + + def _per_channel_conv2d( + self, kernel: int = 3, out_channels: int = 4 + ) -> tuple[tuple[torch.Tensor, ...], torch.fx.GraphModule]: + in_channels = 3 + x = torch.randint(-8, 8, (1, in_channels, 8, 8), dtype=torch.int8) + w = torch.randint( + -8, 8, (out_channels, in_channels, kernel, kernel), dtype=torch.int8 + ) + b = torch.randint(-16, 16, (out_channels,), dtype=torch.int32) + w_zero_point = torch.zeros(out_channels, dtype=torch.int32) + b_scale = torch.full((out_channels,), 0.01) + out_scale = 0.1 + # The conv reference uses bias_scale/out_scale directly while linear uses + # the Q31 multiplier, so they have to encode the same ratio for the + # im2row rewrite to be numerically neutral. + multiplier, shift = quantize_tensor_multiplier(b_scale / out_scale) + out_multiplier = multiplier.to(torch.int32) + out_shift = shift.to(torch.int32) + placeholders = (x, w, b, w_zero_point, b_scale, out_multiplier, out_shift) + args = ( + x, + w, + b, + (1, 1), + (0, 0), + (1, 1), + 1, # groups + 0, # input_zero_point + w_zero_point, + b_scale, + out_scale, + 0, # out_zero_point + out_multiplier, + out_shift, + ) + return placeholders, single_op_builder( + placeholders=placeholders, + op=exir_ops.edge.cadence.quantized_conv2d_nchw.default, + args=args, + ) + + def test_channel_last_preserves_per_channel_overload(self) -> None: + placeholders, gm = self._per_channel_conv2d() + original = copy.deepcopy(gm) + + result = ReplaceConvWithChannelLastConvPass().call(gm) + self.assertTrue(result.modified) + gm_after = result.graph_module + + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.quantized_conv2d_nchw.default), + 0, + ) + # The rewrite must land on the tensor-qparam overload, not collapse to + # the scalar one. + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.quantized_conv2d_nhwc.default), + 1, + ) + self.assertEqual( + count_node( + gm_after, exir_ops.edge.cadence.quantized_conv2d_nhwc.per_tensor + ), + 0, + ) + validate( + original, + gm_after, + list(placeholders), + "ReplaceConvWithChannelLastConvPass (per-channel)", + ) + + def test_channel_last_keeps_qparams_as_vectors(self) -> None: + """The qparam args must stay per-channel nodes across the rewrite.""" + _, gm = self._per_channel_conv2d(out_channels=4) + + gm_after = ReplaceConvWithChannelLastConvPass().call(gm).graph_module + + conv = gm_after.graph.find_nodes( + op="call_function", + target=exir_ops.edge.cadence.quantized_conv2d_nhwc.default, + )[0] + for idx, name in ((8, "weight_zero_point"), (12, "out_multiplier")): + arg = conv.args[idx] + self.assertIsInstance(arg, torch.fx.Node, f"{name} should stay a node") + self.assertEqual( + tuple(cast(torch.fx.Node, arg).meta["val"].shape), + (4,), + f"{name} should stay a per-channel vector", + ) + + def test_im2row_lowers_per_channel_conv_to_per_channel_linear(self) -> None: + # No numerical validation across this rewrite: the conv reference scales + # by bias_scale/out_scale while the linear reference (like the kernel) + # scales by -out_multiplier/2^31 * 2^out_shift, so the two sides disagree + # by a sign at any granularity. + _, gm = self._per_channel_conv2d() + result = cast(PassResult, ReplaceConvWithIm2RowAndLinear()(gm)) + self.assertTrue( + result.modified, "per-channel conv2d should take the im2row exit" + ) + gm_after = result.graph_module + + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.quantized_conv2d_nchw.default), + 0, + ) + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.im2row.per_tensor), 1 + ) + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.quantized_linear.default), 1 + ) + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.quantized_linear.per_tensor), 0 + ) + + def test_im2row_reuses_the_per_channel_multiplier(self) -> None: + """The linear must inherit the conv's per-channel multiplier and shift. + + The per-tensor path recomputes these from a scalar bias_scale; doing that + for per-channel would flatten the vector to channel 0's value. + """ + _, gm = self._per_channel_conv2d(out_channels=4) + conv = gm.graph.find_nodes( + op="call_function", + target=exir_ops.edge.cadence.quantized_conv2d_nchw.default, + )[0] + conv_multiplier, conv_shift = conv.args[12], conv.args[13] + + gm_after = cast(PassResult, ReplaceConvWithIm2RowAndLinear()(gm)).graph_module + + linear = gm_after.graph.find_nodes( + op="call_function", target=exir_ops.edge.cadence.quantized_linear.default + )[0] + # quantized_linear(src, weight, bias, in_zp, w_zp, multiplier, shift, ...) + self.assertEqual( + cast(torch.fx.Node, linear.args[5]).name, + cast(torch.fx.Node, conv_multiplier).name, + ) + self.assertEqual( + cast(torch.fx.Node, linear.args[6]).name, + cast(torch.fx.Node, conv_shift).name, + ) + + def test_linear_to_fully_connected_preserves_per_channel(self) -> None: + """The final conv exit is linear -> fully_connected, on a batch of 1.""" + out_features, in_features = 4, 6 + src = torch.randint(-8, 8, (1, in_features), dtype=torch.int8) + weight = torch.randint(-8, 8, (out_features, in_features), dtype=torch.int8) + bias = torch.randint(-16, 16, (out_features,), dtype=torch.int32) + w_zero_point = torch.zeros(out_features, dtype=torch.int32) + out_multiplier = torch.full((out_features,), 1 << 30, dtype=torch.int32) + out_shift = torch.zeros(out_features, dtype=torch.int32) + placeholders = (src, weight, bias, w_zero_point, out_multiplier, out_shift) + gm = single_op_builder( + placeholders=placeholders, + op=exir_ops.edge.cadence.quantized_linear.default, + args=( + src, + weight, + bias, + 0, + w_zero_point, + out_multiplier, + out_shift, + 0, + None, + ), + ) + + gm_after = cast( + PassResult, ReplaceLinearWithFullyConnectedOpPass()(gm) + ).graph_module + + self.assertEqual( + count_node( + gm_after, exir_ops.edge.cadence.quantized_fully_connected.default + ), + 1, + ) + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.quantized_linear.default), 0 + ) + + def _per_channel_conv1d_ncl( + self, + ) -> tuple[tuple[torch.Tensor, ...], torch.fx.GraphModule]: + in_channels, out_channels, kernel = 3, 4, 3 + x = torch.randint(-8, 8, (1, in_channels, 6), dtype=torch.int8) + w = torch.randint( + -8, 8, (out_channels, in_channels, kernel), dtype=torch.int8 + ) + b = torch.randint(-16, 16, (out_channels,), dtype=torch.int32) + w_zero_point = torch.zeros(out_channels, dtype=torch.int32) + # Deliberately distinct per output channel. + b_scale = torch.tensor([0.01, 0.03, 0.005, 0.07]) + out_scale = 0.1 + multiplier, shift = quantize_tensor_multiplier(b_scale / out_scale) + placeholders = (x, w, b, w_zero_point, b_scale, multiplier, shift) + args = ( + x, + w, + b, + (1,), + (0,), + (1,), + 1, + 0, + w_zero_point, + b_scale, + out_scale, + 0, + multiplier, + shift, + ) + return placeholders, single_op_builder( + placeholders=placeholders, + op=exir_ops.edge.cadence.quantized_conv1d_ncl.default, + args=args, + ) + + def _per_channel_depthwise_conv1d_ncl( + self, + ) -> tuple[tuple[torch.Tensor, ...], torch.fx.GraphModule]: + channels, kernel = 3, 3 + x = torch.randint(-8, 8, (1, channels, 6), dtype=torch.int8) + w = torch.randint(-8, 8, (channels, 1, kernel), dtype=torch.int8) + b = torch.zeros(channels, dtype=torch.int32) + w_zero_point = torch.zeros(channels, dtype=torch.int32) + b_scale = torch.tensor([0.01, 0.05, 0.005]) + out_scale = 0.1 + multiplier, shift = quantize_tensor_multiplier(b_scale / out_scale) + placeholders = (x, w, b, w_zero_point, b_scale, multiplier, shift) + args = ( + x, + w, + b, + (1,), + (0,), + (1,), + channels, # depthwise + 0, + w_zero_point, + b_scale, + out_scale, + 0, + multiplier, + shift, + ) + return placeholders, single_op_builder( + placeholders=placeholders, + op=exir_ops.edge.cadence.quantized_depthwise_conv1d_ncl.default, + args=args, + ) + + def test_channel_last_preserves_per_channel_overload_conv1d(self) -> None: + """The per-channel 1d conv has to stay on ``.default`` after NCL→NLC.""" + _, gm = self._per_channel_conv1d_ncl() + gm_after = ReplaceConvWithChannelLastConvPass().call(gm).graph_module + + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.quantized_conv1d_ncl.default), + 0, + ) + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.quantized_conv1d_nlc.default), + 1, + ) + # And the scalar overload must NOT appear -- that would silently lose + # the tensor qparams. + self.assertEqual( + count_node( + gm_after, exir_ops.edge.cadence.quantized_conv1d_nlc.per_tensor + ), + 0, + ) + + def test_channel_last_routes_per_channel_depthwise_to_depthwise_nlc( + self, + ) -> None: + """Depthwise routing must survive alongside the per-channel overload swap. + + If the pass only looked at ``.per_tensor`` targets, a per-channel + depthwise conv would either miss the depthwise op entirely or fall + back to the dense NLC form. + """ + _, gm = self._per_channel_depthwise_conv1d_ncl() + gm_after = ReplaceConvWithChannelLastConvPass().call(gm).graph_module + + self.assertEqual( + count_node( + gm_after, + exir_ops.edge.cadence.quantized_depthwise_conv1d_ncl.default, + ), + 0, + ) + self.assertEqual( + count_node( + gm_after, + exir_ops.edge.cadence.quantized_depthwise_conv1d_nlc.default, + ), + 1, + ) + # The dense NLC op must not appear either. + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.quantized_conv1d_nlc.default), + 0, + ) + self.assertEqual( + count_node( + gm_after, + exir_ops.edge.cadence.quantized_depthwise_conv1d_nlc.per_tensor, + ), + 0, + ) + + def test_replace_per_channel_conv2d_nchw_with_linear(self) -> None: + """A trivial per-channel quantized conv2d collapses to quantized_linear. + + The map at replace_ops.py ~line 856 has four per-channel entries; this + walks the 2d side of it. If the pass rebuilt multiplier/shift from a + scalar bias_scale, the per-channel vectors would collapse. + """ + in_channels, out_channels = 3, 4 + kh, kw = 2, 2 + x = torch.randint(-10, 10, (1, in_channels, kh, kw), dtype=torch.int8) + w = torch.randint( + -5, 5, (out_channels, in_channels, kh, kw), dtype=torch.int8 + ) + b = torch.zeros(out_channels, dtype=torch.int32) + weight_zero_point = torch.zeros(out_channels, dtype=torch.int32) + bias_scale = torch.tensor([0.001, 0.002, 0.005, 0.01]) + out_multiplier = torch.tensor( + [1073741824, 1181116006, 1288490188, 1395864371], dtype=torch.int32 + ) + out_shift = torch.tensor([0, -1, -2, -3], dtype=torch.int32) + placeholders = ( + x, + w, + b, + weight_zero_point, + bias_scale, + out_multiplier, + out_shift, + ) + args = ( + x, + w, + b, + (1, 1), + (0, 0), + (1, 1), + 1, + 0, + weight_zero_point, + bias_scale, + 1.0, + 0, + out_multiplier, + out_shift, + ) + gm = single_op_builder( + placeholders=placeholders, + op=exir_ops.edge.cadence.quantized_conv2d_nchw.default, + args=args, + ) + conv_nodes = gm.graph.find_nodes( + op="call_function", + target=exir_ops.edge.cadence.quantized_conv2d_nchw.default, + ) + self.assertEqual(len(conv_nodes), 1) + conv_args = conv_nodes[0].args + wzp_name = cast(torch.fx.Node, conv_args[8]).name + multiplier_name = cast(torch.fx.Node, conv_args[12]).name + shift_name = cast(torch.fx.Node, conv_args[13]).name + + result = cast(PassResult, ReplaceTrivialConvWithLinear()(gm)) + self.assertTrue(result.modified) + gm_after = result.graph_module + + self.assertEqual( + count_node( + gm_after, exir_ops.edge.cadence.quantized_conv2d_nchw.default + ), + 0, + ) + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.quantized_linear.default), + 1, + ) + linear = gm_after.graph.find_nodes( + op="call_function", + target=exir_ops.edge.cadence.quantized_linear.default, + )[0] + # Same nodes forwarded, not rebuilt from scalars. + self.assertEqual(cast(torch.fx.Node, linear.args[4]).name, wzp_name) + self.assertEqual(cast(torch.fx.Node, linear.args[5]).name, multiplier_name) + self.assertEqual(cast(torch.fx.Node, linear.args[6]).name, shift_name) + self.assertEqual( + cast(torch.fx.Node, linear.args[5]).meta["val"].shape, + torch.Size([out_channels]), + ) + + def test_im2row_lowers_per_channel_conv2d_nhwc_to_linear(self) -> None: + """The im2row + linear rewrite covers the NHWC per-channel entry too. + + The map at replace_ops.py ~line 1351 includes both nchw.default and + nhwc.default; the nhwc side would silently regress if only nchw were + tested. + """ + in_channels, out_channels = 3, 4 + # channel-last: [N, H, W, C]; im2row picks the same conv on nhwc. + x = torch.randint(-8, 8, (1, 4, 4, in_channels), dtype=torch.int8) + w = torch.randint(-8, 8, (out_channels, 3, 3, in_channels), dtype=torch.int8) + b = torch.randint(-16, 16, (out_channels,), dtype=torch.int32) + w_zero_point = torch.zeros(out_channels, dtype=torch.int32) + b_scale = torch.tensor([0.01, 0.03, 0.005, 0.07]) + out_scale = 0.1 + multiplier, shift = quantize_tensor_multiplier(b_scale / out_scale) + placeholders = (x, w, b, w_zero_point, b_scale, multiplier, shift) + args = ( + x, + w, + b, + (1, 1), + (0, 0), + (1, 1), + 1, + 0, + w_zero_point, + b_scale, + out_scale, + 0, + multiplier, + shift, + ) + gm = single_op_builder( + placeholders=placeholders, + op=exir_ops.edge.cadence.quantized_conv2d_nhwc.default, + args=args, + ) + conv_multiplier = gm.graph.find_nodes( + op="call_function", + target=exir_ops.edge.cadence.quantized_conv2d_nhwc.default, + )[0].args[12] + conv_shift = gm.graph.find_nodes( + op="call_function", + target=exir_ops.edge.cadence.quantized_conv2d_nhwc.default, + )[0].args[13] + + result = cast(PassResult, ReplaceConvWithIm2RowAndLinear()(gm)) + self.assertTrue(result.modified) + gm_after = result.graph_module + + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.quantized_conv2d_nhwc.default), + 0, + ) + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.im2row.per_tensor), 1 + ) + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.quantized_linear.default), 1 + ) + self.assertEqual( + count_node(gm_after, exir_ops.edge.cadence.quantized_linear.per_tensor), + 0, + ) + linear = gm_after.graph.find_nodes( + op="call_function", + target=exir_ops.edge.cadence.quantized_linear.default, + )[0] + # The linear must reuse the conv's own multiplier/shift constants + # rather than rebuilding scalars. + self.assertEqual( + cast(torch.fx.Node, linear.args[5]).name, + cast(torch.fx.Node, conv_multiplier).name, + ) + self.assertEqual( + cast(torch.fx.Node, linear.args[6]).name, + cast(torch.fx.Node, conv_shift).name, + ) + + class TestReplaceMaxPool2dWithChannelLastMaxPool2dPass(unittest.TestCase): def test_replace_max_pool2d_nchw_with_nhwc(self) -> None: # Create a graph with a single quantized_max_pool2d_nchw node. diff --git a/backends/cadence/generic/operators/op_quantized_conv1d_ncl.cpp b/backends/cadence/generic/operators/op_quantized_conv1d_ncl.cpp index 4108657d8b3..102bd906d37 100644 --- a/backends/cadence/generic/operators/op_quantized_conv1d_ncl.cpp +++ b/backends/cadence/generic/operators/op_quantized_conv1d_ncl.cpp @@ -63,9 +63,13 @@ __attribute__((noinline)) void conv1d_ncl_core_generic( // Optional args that are only relevant for quantized convolution // input zero point IT in_zero_point = 0, - // weight zero point - int32_t weight_zero_point = 0, - float bias_scale = 1, + // weight zero point and bias scale. A stride of 0 keeps every output + // channel on the same qparam (per-tensor); a stride of 1 walks one entry + // per output channel (per-channel). + const int32_t* __restrict__ p_weight_zero_point = nullptr, + int32_t weight_zero_point_stride = 0, + const float* __restrict__ p_bias_scale = nullptr, + int32_t bias_scale_stride = 0, float out_scale = 1, OT out_zero_point = 0) { float inv_out_scale = 1. / out_scale; @@ -87,6 +91,13 @@ __attribute__((noinline)) void conv1d_ncl_core_generic( int soc = _g * ocpg; // Populate all the output channels in the group for (int _oc = soc; _oc < soc + ocpg; ++_oc) { + int32_t weight_zero_point = 0; + float bias_scale = 1.f; + if (quantized) { + weight_zero_point = + p_weight_zero_point[_oc * weight_zero_point_stride]; + bias_scale = p_bias_scale[_oc * bias_scale_stride]; + } OT* out_plane = out_batch + _oc * ow; const WT* weight_batch = p_weight + _oc * wc * ww; // We compute one output channel at a time. The computation can be @@ -152,8 +163,10 @@ void quantized_conv1d_ncl( IntArrayRef dilation, int16_t groups, int32_t in_zero_point, - int32_t weight_zero_point, - float bias_scale, + const int32_t* __restrict__ p_weight_zero_point, + int32_t weight_zero_point_stride, + const float* __restrict__ p_bias_scale, + int32_t bias_scale_stride, float output_scale, int32_t output_zero_point, Tensor& out) { @@ -187,8 +200,10 @@ void quantized_conv1d_ncl( dilation[dilation.size() - 1], \ groups, \ in_zero_point, \ - weight_zero_point, \ - bias_scale, \ + p_weight_zero_point, \ + weight_zero_point_stride, \ + p_bias_scale, \ + bias_scale_stride, \ output_scale, \ (ctype)output_zero_point); \ break; \ @@ -204,11 +219,44 @@ void quantized_conv1d_ncl( #undef typed_quantized_conv1d_ncl } +void quantized_conv1d_ncl( + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + quantized_conv1d_ncl( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + &weight_zero_point, + 0, + &bias_scale, + 0, + output_scale, + output_zero_point, + out); +} + } // namespace // Public exported kernel functions -::executorch::aten::Tensor& quantized_conv1d_ncl_out( + +::executorch::aten::Tensor& quantized_conv1d_ncl_per_tensor_out( KernelRuntimeContext& ctx, const Tensor& input, const Tensor& weight, @@ -218,16 +266,14 @@ ::executorch::aten::Tensor& quantized_conv1d_ncl_out( IntArrayRef dilation, int64_t groups, int64_t input_zero_point, - const Tensor& weight_zero_point, - const Tensor& bias_scale, + int64_t weight_zero_point, + double bias_scale, double output_scale, int64_t output_zero_point, - const Tensor& out_multiplier, - const Tensor& out_shift, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, Tensor& out) { (void)ctx; - (void)out_multiplier; - (void)out_shift; quantized_conv1d_ncl( input, weight, @@ -237,15 +283,19 @@ ::executorch::aten::Tensor& quantized_conv1d_ncl_out( dilation, static_cast(groups), static_cast(input_zero_point), - weight_zero_point.const_data_ptr()[0], - bias_scale.const_data_ptr()[0], + static_cast(weight_zero_point), + static_cast(bias_scale), static_cast(output_scale), static_cast(output_zero_point), out); return out; } -::executorch::aten::Tensor& quantized_conv1d_ncl_per_tensor_out( + +// weight_zero_point and bias_scale carry one entry per output channel. +// out_multiplier/out_shift stay unused, matching the reference +// implementation, which folds the accumulator scale into bias_scale. +::executorch::aten::Tensor& quantized_conv1d_ncl_out( KernelRuntimeContext& ctx, const Tensor& input, const Tensor& weight, @@ -255,14 +305,16 @@ ::executorch::aten::Tensor& quantized_conv1d_ncl_per_tensor_out( IntArrayRef dilation, int64_t groups, int64_t input_zero_point, - int64_t weight_zero_point, - double bias_scale, + const Tensor& weight_zero_point, + const Tensor& bias_scale, double output_scale, int64_t output_zero_point, - __ET_UNUSED int64_t out_multiplier, - __ET_UNUSED int64_t out_shift, + const Tensor& out_multiplier, + const Tensor& out_shift, Tensor& out) { (void)ctx; + (void)out_multiplier; + (void)out_shift; quantized_conv1d_ncl( input, weight, @@ -272,8 +324,10 @@ ::executorch::aten::Tensor& quantized_conv1d_ncl_per_tensor_out( dilation, static_cast(groups), static_cast(input_zero_point), - static_cast(weight_zero_point), - static_cast(bias_scale), + weight_zero_point.const_data_ptr(), + weight_zero_point.numel() > 1 ? 1 : 0, + bias_scale.const_data_ptr(), + bias_scale.numel() > 1 ? 1 : 0, static_cast(output_scale), static_cast(output_zero_point), out); diff --git a/backends/cadence/generic/operators/op_quantized_conv1d_ncl.h b/backends/cadence/generic/operators/op_quantized_conv1d_ncl.h index f6854beff12..a24b69a9faf 100644 --- a/backends/cadence/generic/operators/op_quantized_conv1d_ncl.h +++ b/backends/cadence/generic/operators/op_quantized_conv1d_ncl.h @@ -20,7 +20,8 @@ using ::executorch::aten::Tensor; using ::executorch::runtime::KernelRuntimeContext; // NCL format (N=batch, C=channels, L=length) -::executorch::aten::Tensor& quantized_conv1d_ncl_out( + +::executorch::aten::Tensor& quantized_conv1d_ncl_per_tensor_out( KernelRuntimeContext& ctx, const Tensor& input, const Tensor& weight, @@ -30,15 +31,15 @@ ::executorch::aten::Tensor& quantized_conv1d_ncl_out( IntArrayRef dilation, int64_t groups, int64_t input_zero_point, - const Tensor& weight_zero_point, - const Tensor& bias_scale, + int64_t weight_zero_point, + double bias_scale, double output_scale, int64_t output_zero_point, - const Tensor& out_multiplier, - const Tensor& out_shift, + int64_t out_multiplier, + int64_t out_shift, Tensor& out); -::executorch::aten::Tensor& quantized_conv1d_ncl_per_tensor_out( +::executorch::aten::Tensor& quantized_conv1d_ncl_out( KernelRuntimeContext& ctx, const Tensor& input, const Tensor& weight, @@ -48,12 +49,12 @@ ::executorch::aten::Tensor& quantized_conv1d_ncl_per_tensor_out( IntArrayRef dilation, int64_t groups, int64_t input_zero_point, - int64_t weight_zero_point, - double bias_scale, + const Tensor& weight_zero_point, + const Tensor& bias_scale, double output_scale, int64_t output_zero_point, - int64_t out_multiplier, - int64_t out_shift, + const Tensor& out_multiplier, + const Tensor& out_shift, Tensor& out); } // namespace native diff --git a/backends/cadence/generic/operators/op_quantized_conv1d_nlc.cpp b/backends/cadence/generic/operators/op_quantized_conv1d_nlc.cpp index 8a427045a83..2d5cbf8cde9 100644 --- a/backends/cadence/generic/operators/op_quantized_conv1d_nlc.cpp +++ b/backends/cadence/generic/operators/op_quantized_conv1d_nlc.cpp @@ -63,9 +63,13 @@ __attribute__((noinline)) void conv1d_nlc_core_generic( // Optional args that are only relevant for quantized convolution // input zero point IT in_zero_point = 0, - // weight zero point - int32_t weight_zero_point = 0, - float bias_scale = 1, + // weight zero point and bias scale. A stride of 0 keeps every output + // channel on the same qparam (per-tensor); a stride of 1 walks one entry + // per output channel (per-channel). + const int32_t* __restrict__ p_weight_zero_point = nullptr, + int32_t weight_zero_point_stride = 0, + const float* __restrict__ p_bias_scale = nullptr, + int32_t bias_scale_stride = 0, float out_scale = 1, OT out_zero_point = 0) { float inv_out_scale = 1. / out_scale; @@ -89,6 +93,13 @@ __attribute__((noinline)) void conv1d_nlc_core_generic( int soc = _g * ocpg; // Populate all the output channels in the group for (int _oc = soc; _oc < soc + ocpg; ++_oc) { + int32_t weight_zero_point = 0; + float bias_scale = 1.f; + if (quantized) { + weight_zero_point = + p_weight_zero_point[_oc * weight_zero_point_stride]; + bias_scale = p_bias_scale[_oc * bias_scale_stride]; + } const WT* weight_batch = p_weight + _oc * ww * wc; // We compute one output channel at a time. The computation can be // thought of as a stencil computation: we iterate over an input of @@ -147,8 +158,10 @@ void quantized_conv1d_nlc( IntArrayRef dilation, int16_t groups, int32_t in_zero_point, - int32_t weight_zero_point, - float bias_scale, + const int32_t* __restrict__ p_weight_zero_point, + int32_t weight_zero_point_stride, + const float* __restrict__ p_bias_scale, + int32_t bias_scale_stride, float output_scale, int32_t output_zero_point, Tensor& out) { @@ -182,8 +195,10 @@ void quantized_conv1d_nlc( dilation[dilation.size() - 1], \ groups, \ in_zero_point, \ - weight_zero_point, \ - bias_scale, \ + p_weight_zero_point, \ + weight_zero_point_stride, \ + p_bias_scale, \ + bias_scale_stride, \ output_scale, \ (ctype)output_zero_point); \ break; \ @@ -199,11 +214,44 @@ void quantized_conv1d_nlc( #undef typed_quantized_conv1d_nlc } +void quantized_conv1d_nlc( + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + quantized_conv1d_nlc( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + &weight_zero_point, + 0, + &bias_scale, + 0, + output_scale, + output_zero_point, + out); +} + } // namespace // Public exported kernel functions -::executorch::aten::Tensor& quantized_conv1d_nlc_out( + +::executorch::aten::Tensor& quantized_conv1d_nlc_per_tensor_out( KernelRuntimeContext& ctx, const Tensor& input, const Tensor& weight, @@ -213,16 +261,15 @@ ::executorch::aten::Tensor& quantized_conv1d_nlc_out( IntArrayRef dilation, int64_t groups, int64_t input_zero_point, - const Tensor& weight_zero_point, - const Tensor& bias_scale, + int64_t weight_zero_point, + double bias_scale, double output_scale, int64_t output_zero_point, - const Tensor& out_multiplier, - const Tensor& out_shift, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + __ET_UNUSED const std::optional& offset, Tensor& out) { (void)ctx; - (void)out_multiplier; - (void)out_shift; quantized_conv1d_nlc( input, weight, @@ -232,15 +279,19 @@ ::executorch::aten::Tensor& quantized_conv1d_nlc_out( dilation, static_cast(groups), static_cast(input_zero_point), - weight_zero_point.const_data_ptr()[0], - bias_scale.const_data_ptr()[0], + static_cast(weight_zero_point), + static_cast(bias_scale), static_cast(output_scale), static_cast(output_zero_point), out); return out; } -::executorch::aten::Tensor& quantized_conv1d_nlc_per_tensor_out( + +// weight_zero_point and bias_scale carry one entry per output channel. +// out_multiplier/out_shift stay unused, matching the reference +// implementation, which folds the accumulator scale into bias_scale. +::executorch::aten::Tensor& quantized_conv1d_nlc_out( KernelRuntimeContext& ctx, const Tensor& input, const Tensor& weight, @@ -250,15 +301,16 @@ ::executorch::aten::Tensor& quantized_conv1d_nlc_per_tensor_out( IntArrayRef dilation, int64_t groups, int64_t input_zero_point, - int64_t weight_zero_point, - double bias_scale, + const Tensor& weight_zero_point, + const Tensor& bias_scale, double output_scale, int64_t output_zero_point, - __ET_UNUSED int64_t out_multiplier, - __ET_UNUSED int64_t out_shift, - __ET_UNUSED const std::optional& offset, + const Tensor& out_multiplier, + const Tensor& out_shift, Tensor& out) { (void)ctx; + (void)out_multiplier; + (void)out_shift; quantized_conv1d_nlc( input, weight, @@ -268,8 +320,10 @@ ::executorch::aten::Tensor& quantized_conv1d_nlc_per_tensor_out( dilation, static_cast(groups), static_cast(input_zero_point), - static_cast(weight_zero_point), - static_cast(bias_scale), + weight_zero_point.const_data_ptr(), + weight_zero_point.numel() > 1 ? 1 : 0, + bias_scale.const_data_ptr(), + bias_scale.numel() > 1 ? 1 : 0, static_cast(output_scale), static_cast(output_zero_point), out); diff --git a/backends/cadence/generic/operators/op_quantized_conv1d_nlc.h b/backends/cadence/generic/operators/op_quantized_conv1d_nlc.h index f1780497f73..7fc14824405 100644 --- a/backends/cadence/generic/operators/op_quantized_conv1d_nlc.h +++ b/backends/cadence/generic/operators/op_quantized_conv1d_nlc.h @@ -20,7 +20,8 @@ using ::executorch::aten::Tensor; using ::executorch::runtime::KernelRuntimeContext; // NLC format (N=batch, L=length, C=channels) -::executorch::aten::Tensor& quantized_conv1d_nlc_out( + +::executorch::aten::Tensor& quantized_conv1d_nlc_per_tensor_out( KernelRuntimeContext& ctx, const Tensor& input, const Tensor& weight, @@ -30,15 +31,16 @@ ::executorch::aten::Tensor& quantized_conv1d_nlc_out( IntArrayRef dilation, int64_t groups, int64_t input_zero_point, - const Tensor& weight_zero_point, - const Tensor& bias_scale, + int64_t weight_zero_point, + double bias_scale, double output_scale, int64_t output_zero_point, - const Tensor& out_multiplier, - const Tensor& out_shift, + int64_t out_multiplier, + int64_t out_shift, + const std::optional& offset, Tensor& out); -::executorch::aten::Tensor& quantized_conv1d_nlc_per_tensor_out( +::executorch::aten::Tensor& quantized_conv1d_nlc_out( KernelRuntimeContext& ctx, const Tensor& input, const Tensor& weight, @@ -48,13 +50,12 @@ ::executorch::aten::Tensor& quantized_conv1d_nlc_per_tensor_out( IntArrayRef dilation, int64_t groups, int64_t input_zero_point, - int64_t weight_zero_point, - double bias_scale, + const Tensor& weight_zero_point, + const Tensor& bias_scale, double output_scale, int64_t output_zero_point, - int64_t out_multiplier, - int64_t out_shift, - const std::optional& offset, + const Tensor& out_multiplier, + const Tensor& out_shift, Tensor& out); } // namespace native diff --git a/backends/cadence/generic/operators/op_quantized_conv2d.cpp b/backends/cadence/generic/operators/op_quantized_conv2d.cpp index a5d80fd2892..ceea9a1f2fd 100644 --- a/backends/cadence/generic/operators/op_quantized_conv2d.cpp +++ b/backends/cadence/generic/operators/op_quantized_conv2d.cpp @@ -65,9 +65,13 @@ __attribute__((noinline)) void conv2d_nchw_core_generic( // Optional args that are only relevant for quantized convolution // input zero point IT in_zero_point = 0, - // weight zero point - int32_t weight_zero_point = 0, - float bias_scale = 1, + // weight zero point and bias scale. A stride of 0 keeps every output + // channel on the same qparam (per-tensor); a stride of 1 walks one entry + // per output channel (per-channel). + const int32_t* __restrict__ p_weight_zero_point = nullptr, + int32_t weight_zero_point_stride = 0, + const float* __restrict__ p_bias_scale = nullptr, + int32_t bias_scale_stride = 0, float out_scale = 1, OT out_zero_point = 0) { const float inv_out_scale = 1.f / out_scale; @@ -89,6 +93,13 @@ __attribute__((noinline)) void conv2d_nchw_core_generic( int soc = _g * ocpg; // Populate all the output channels in the group for (int _oc = soc; _oc < soc + ocpg; ++_oc) { + int32_t weight_zero_point = 0; + float bias_scale = 1.f; + if (quantized) { + weight_zero_point = + p_weight_zero_point[_oc * weight_zero_point_stride]; + bias_scale = p_bias_scale[_oc * bias_scale_stride]; + } OT* out_plane = out_batch + _oc * oh * ow; const WT* weight_batch = p_weight + _oc * wc * wh * ww; // We compute one output channel at a time. The computation can be @@ -193,9 +204,13 @@ __attribute__((noinline)) void conv2d_nhwc_core_generic( // Optional args that are only relevant for quantized convolution // input zero point IT in_zero_point = 0, - // weight zero point - int32_t weight_zero_point = 0, - float bias_scale = 1, + // weight zero point and bias scale. A stride of 0 keeps every output + // channel on the same qparam (per-tensor); a stride of 1 walks one entry + // per output channel (per-channel). + const int32_t* __restrict__ p_weight_zero_point = nullptr, + int32_t weight_zero_point_stride = 0, + const float* __restrict__ p_bias_scale = nullptr, + int32_t bias_scale_stride = 0, float out_scale = 1, OT out_zero_point = 0, // Whether this is a depthwise conv with [KH, KW, OC] weight layout @@ -222,6 +237,13 @@ __attribute__((noinline)) void conv2d_nhwc_core_generic( int soc = _g * ocpg; // Populate all the output channels in the group for (int _oc = soc; _oc < soc + ocpg; ++_oc) { + int32_t weight_zero_point = 0; + float bias_scale = 1.f; + if (quantized) { + weight_zero_point = + p_weight_zero_point[_oc * weight_zero_point_stride]; + bias_scale = p_bias_scale[_oc * bias_scale_stride]; + } float acc = p_bias[_oc]; if (zero_pad_unit_dilation) { for (int _wh = 0; _wh < wh; ++_wh) { @@ -291,8 +313,10 @@ void quantized_conv2d_nchw( IntArrayRef dilation, int16_t groups, int32_t in_zero_point, - int32_t weight_zero_point, - float bias_scale, + const int32_t* __restrict__ p_weight_zero_point, + int32_t weight_zero_point_stride, + const float* __restrict__ p_bias_scale, + int32_t bias_scale_stride, float output_scale, int32_t output_zero_point, Tensor& out) { @@ -311,12 +335,6 @@ void quantized_conv2d_nchw( const int oh = conv1d ? 1 : out.size(2); const int ow = conv1d ? out.size(2) : out.size(3); - ET_CHECK_MSG( - weight_zero_point >= -128 && weight_zero_point <= 127, - "weight_zero_point %" PRId32 - " must be in range [-128, 127] for int8 cast", - weight_zero_point); - // Handle W8A16 heterogeneous type (int16_t activations, int8_t weights) if (out.scalar_type() == ScalarType::Short && input.scalar_type() == ScalarType::Short && @@ -344,8 +362,10 @@ void quantized_conv2d_nchw( dilation[1], groups, static_cast(in_zero_point), - static_cast(weight_zero_point), - bias_scale, + p_weight_zero_point, + weight_zero_point_stride, + p_bias_scale, + bias_scale_stride, output_scale, static_cast(output_zero_point)); return; @@ -376,8 +396,10 @@ void quantized_conv2d_nchw( dilation[1], \ groups, \ in_zero_point, \ - weight_zero_point, \ - bias_scale, \ + p_weight_zero_point, \ + weight_zero_point_stride, \ + p_bias_scale, \ + bias_scale_stride, \ output_scale, \ (ctype)output_zero_point); \ break; \ @@ -407,8 +429,10 @@ void quantized_conv2d_nhwc_depthwise( IntArrayRef dilation, int16_t groups, int32_t in_zero_point, - int32_t weight_zero_point, - float bias_scale, + const int32_t* __restrict__ p_weight_zero_point, + int32_t weight_zero_point_stride, + const float* __restrict__ p_bias_scale, + int32_t bias_scale_stride, float output_scale, int32_t output_zero_point, Tensor& out) { @@ -452,6 +476,10 @@ void quantized_conv2d_nhwc_depthwise( for (int _g = 0; _g < groups; ++_g) { \ int soc = _g * ocpg; \ for (int _oc = soc; _oc < soc + ocpg; ++_oc) { \ + const int32_t weight_zero_point = \ + p_weight_zero_point[_oc * weight_zero_point_stride]; \ + const float bias_scale = \ + p_bias_scale[_oc * bias_scale_stride]; \ float acc = p_bias[_oc]; \ for (int _kh = 0; _kh < kh; ++_kh) { \ for (int _kw = 0; _kw < kw; ++_kw) { \ @@ -497,8 +525,10 @@ void quantized_conv2d_nhwc( IntArrayRef dilation, int16_t groups, int32_t in_zero_point, - int32_t weight_zero_point, - float bias_scale, + const int32_t* __restrict__ p_weight_zero_point, + int32_t weight_zero_point_stride, + const float* __restrict__ p_bias_scale, + int32_t bias_scale_stride, float output_scale, int32_t output_zero_point, Tensor& out) { @@ -556,8 +586,10 @@ void quantized_conv2d_nhwc( dilation[1], groups, static_cast(in_zero_point), - static_cast(weight_zero_point), - bias_scale, + p_weight_zero_point, + weight_zero_point_stride, + p_bias_scale, + bias_scale_stride, output_scale, static_cast(output_zero_point), is_depthwise); @@ -589,8 +621,10 @@ void quantized_conv2d_nhwc( dilation[1], \ groups, \ in_zero_point, \ - weight_zero_point, \ - bias_scale, \ + p_weight_zero_point, \ + weight_zero_point_stride, \ + p_bias_scale, \ + bias_scale_stride, \ output_scale, \ (ctype)output_zero_point, \ is_depthwise); \ @@ -607,26 +641,25 @@ void quantized_conv2d_nhwc( #undef typed_quantized_conv2d_nhwc } -Tensor& quantized_conv2d_nchw_out( - ET_UNUSED KernelRuntimeContext& ctx, +void quantized_conv2d_nchw( const Tensor& input, const Tensor& weight, const Tensor& bias, IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, - int64_t groups, - int64_t in_zero_point, - const Tensor& weight_zero_point, - const Tensor& bias_scale, - double output_scale, - int64_t output_zero_point, - ET_UNUSED const Tensor& out_multiplier, - ET_UNUSED const Tensor& out_shift, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, Tensor& out) { - const float bias_scale_float = bias_scale.const_data_ptr()[0]; - const int32_t weight_zero_point_int = - weight_zero_point.const_data_ptr()[0]; + ET_CHECK_MSG( + weight_zero_point >= -128 && weight_zero_point <= 127, + "weight_zero_point %" PRId32 + " must be in range [-128, 127] for int8 cast", + weight_zero_point); quantized_conv2d_nchw( input, weight, @@ -636,34 +669,61 @@ Tensor& quantized_conv2d_nchw_out( dilation, groups, in_zero_point, - weight_zero_point_int, - bias_scale_float, + &weight_zero_point, + 0, + &bias_scale, + 0, output_scale, output_zero_point, out); - return out; } -Tensor& quantized_conv2d_nhwc_out( - ET_UNUSED KernelRuntimeContext& ctx, +void quantized_conv2d_nhwc_depthwise( const Tensor& input, const Tensor& weight, const Tensor& bias, IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, - int64_t groups, - int64_t in_zero_point, - const Tensor& weight_zero_point, - const Tensor& bias_scale, - double output_scale, - int64_t output_zero_point, - ET_UNUSED const Tensor& out_multiplier, - ET_UNUSED const Tensor& out_shift, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + quantized_conv2d_nhwc_depthwise( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + &weight_zero_point, + 0, + &bias_scale, + 0, + output_scale, + output_zero_point, + out); +} + +void quantized_conv2d_nhwc( + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, Tensor& out) { - const float bias_scale_float = bias_scale.const_data_ptr()[0]; - const int32_t weight_zero_point_int = - weight_zero_point.const_data_ptr()[0]; quantized_conv2d_nhwc( input, weight, @@ -673,14 +733,17 @@ Tensor& quantized_conv2d_nhwc_out( dilation, groups, in_zero_point, - weight_zero_point_int, - bias_scale_float, + &weight_zero_point, + 0, + &bias_scale, + 0, output_scale, output_zero_point, out); - return out; } + + Tensor& quantized_conv2d_nchw_per_tensor_out( ET_UNUSED KernelRuntimeContext& ctx, const Tensor& input, @@ -1192,6 +1255,89 @@ Tensor& quantized_conv2d_nhwc_dilated_asym8uxsym8u_asym8u_per_tensor_out( return out; } + +// weight_zero_point and bias_scale carry one entry per output channel. +// out_multiplier/out_shift stay unused, matching the reference +// implementation, which folds the accumulator scale into bias_scale. +Tensor& quantized_conv2d_nchw_out( + ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + const Tensor& weight_zero_point, + const Tensor& bias_scale, + double output_scale, + int64_t output_zero_point, + ET_UNUSED const Tensor& out_multiplier, + ET_UNUSED const Tensor& out_shift, + Tensor& out) { + quantized_conv2d_nchw( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point.const_data_ptr(), + weight_zero_point.numel() > 1 ? 1 : 0, + bias_scale.const_data_ptr(), + bias_scale.numel() > 1 ? 1 : 0, + output_scale, + output_zero_point, + out); + return out; +} + +// weight_zero_point and bias_scale carry one entry per output channel. +// out_multiplier/out_shift stay unused, matching the reference +// implementation, which folds the accumulator scale into bias_scale. +Tensor& quantized_conv2d_nhwc_out( + ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + const Tensor& weight_zero_point, + const Tensor& bias_scale, + double output_scale, + int64_t output_zero_point, + ET_UNUSED const Tensor& out_multiplier, + ET_UNUSED const Tensor& out_shift, + Tensor& out) { + quantized_conv2d_nhwc( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point.const_data_ptr(), + weight_zero_point.numel() > 1 ? 1 : 0, + bias_scale.const_data_ptr(), + bias_scale.numel() > 1 ? 1 : 0, + output_scale, + output_zero_point, + out); + return out; +} + +// weight_zero_point and bias_scale carry one entry per output channel. +// out_multiplier/out_shift stay unused, matching the reference +// implementation, which folds the accumulator scale into bias_scale. + } // namespace native } // namespace generic } // namespace impl diff --git a/backends/cadence/generic/operators/op_quantized_conv2d.h b/backends/cadence/generic/operators/op_quantized_conv2d.h index 02740d3afec..3c5b747562d 100644 --- a/backends/cadence/generic/operators/op_quantized_conv2d.h +++ b/backends/cadence/generic/operators/op_quantized_conv2d.h @@ -20,23 +20,6 @@ using ::executorch::aten::Tensor; using ::executorch::runtime::KernelRuntimeContext; // Quantized Conv2D operators - NCHW layout -::executorch::aten::Tensor& quantized_conv2d_nchw_out( - KernelRuntimeContext& ctx, - const Tensor& input, - const Tensor& weight, - const Tensor& bias, - IntArrayRef stride, - IntArrayRef padding, - IntArrayRef dilation, - int64_t groups, - int64_t in_zero_point, - const Tensor& weight_zero_point, - const Tensor& bias_scale, - double output_scale, - int64_t output_zero_point, - const Tensor& out_multiplier, - const Tensor& out_shift, - Tensor& out); ::executorch::aten::Tensor& quantized_conv2d_nchw_per_tensor_out( KernelRuntimeContext& ctx, @@ -171,23 +154,6 @@ quantized_conv2d_nchw_dilated_asym8uxsym8u_asym8u_per_tensor_out( Tensor& out); // Quantized Conv2D operators - NHWC layout -::executorch::aten::Tensor& quantized_conv2d_nhwc_out( - KernelRuntimeContext& ctx, - const Tensor& input, - const Tensor& weight, - const Tensor& bias, - IntArrayRef stride, - IntArrayRef padding, - IntArrayRef dilation, - int64_t groups, - int64_t in_zero_point, - const Tensor& weight_zero_point, - const Tensor& bias_scale, - double output_scale, - int64_t output_zero_point, - const Tensor& out_multiplier, - const Tensor& out_shift, - Tensor& out); ::executorch::aten::Tensor& quantized_conv2d_nhwc_per_tensor_out( KernelRuntimeContext& ctx, @@ -340,6 +306,43 @@ quantized_conv2d_nhwc_dilated_asym8uxsym8u_asym8u_per_tensor_out( int64_t out_shift, Tensor& out); +::executorch::aten::Tensor& quantized_conv2d_nchw_out( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + const Tensor& weight_zero_point, + const Tensor& bias_scale, + double output_scale, + int64_t output_zero_point, + const Tensor& out_multiplier, + const Tensor& out_shift, + Tensor& out); + +::executorch::aten::Tensor& quantized_conv2d_nhwc_out( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + const Tensor& weight_zero_point, + const Tensor& bias_scale, + double output_scale, + int64_t output_zero_point, + const Tensor& out_multiplier, + const Tensor& out_shift, + Tensor& out); + + } // namespace native } // namespace generic } // namespace impl diff --git a/backends/cadence/generic/operators/op_quantized_depthwise_conv1d_ncl.cpp b/backends/cadence/generic/operators/op_quantized_depthwise_conv1d_ncl.cpp index ddc8ff44b3a..577173deddb 100644 --- a/backends/cadence/generic/operators/op_quantized_depthwise_conv1d_ncl.cpp +++ b/backends/cadence/generic/operators/op_quantized_depthwise_conv1d_ncl.cpp @@ -60,6 +60,44 @@ ::executorch::aten::Tensor& quantized_depthwise_conv1d_ncl_per_tensor_out( out); } +// Per-channel (Tensor qparams) form. Delegates to the regular conv1d entry +// point, which resolves weight_zero_point and bias_scale per output channel. +::executorch::aten::Tensor& quantized_depthwise_conv1d_ncl_out( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t input_zero_point, + const Tensor& weight_zero_point, + const Tensor& bias_scale, + double output_scale, + int64_t output_zero_point, + const Tensor& out_multiplier, + const Tensor& out_shift, + Tensor& out) { + return quantized_conv1d_ncl_out( + ctx, + input, + weight, + bias, + stride, + padding, + dilation, + groups, + input_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out_multiplier, + out_shift, + out); +} + } // namespace native } // namespace generic } // namespace impl diff --git a/backends/cadence/generic/operators/op_quantized_depthwise_conv1d_nlc.cpp b/backends/cadence/generic/operators/op_quantized_depthwise_conv1d_nlc.cpp index 05fb809cd51..a3ff819042c 100644 --- a/backends/cadence/generic/operators/op_quantized_depthwise_conv1d_nlc.cpp +++ b/backends/cadence/generic/operators/op_quantized_depthwise_conv1d_nlc.cpp @@ -61,6 +61,44 @@ ::executorch::aten::Tensor& quantized_depthwise_conv1d_nlc_per_tensor_out( out); } +// Per-channel (Tensor qparams) form. Delegates to the regular conv1d entry +// point, which resolves weight_zero_point and bias_scale per output channel. +::executorch::aten::Tensor& quantized_depthwise_conv1d_nlc_out( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t input_zero_point, + const Tensor& weight_zero_point, + const Tensor& bias_scale, + double output_scale, + int64_t output_zero_point, + const Tensor& out_multiplier, + const Tensor& out_shift, + Tensor& out) { + return quantized_conv1d_nlc_out( + ctx, + input, + weight, + bias, + stride, + padding, + dilation, + groups, + input_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out_multiplier, + out_shift, + out); +} + } // namespace native } // namespace generic } // namespace impl diff --git a/backends/cadence/generic/operators/op_quantized_fully_connected.cpp b/backends/cadence/generic/operators/op_quantized_fully_connected.cpp index ce74b5b8b7f..13a18685cf2 100644 --- a/backends/cadence/generic/operators/op_quantized_fully_connected.cpp +++ b/backends/cadence/generic/operators/op_quantized_fully_connected.cpp @@ -173,6 +173,7 @@ Tensor& quantized_fully_connected_asym8uxasym8u_asym8u_per_tensor_out( return out; } + } // namespace native } // namespace generic } // namespace impl diff --git a/backends/cadence/generic/operators/op_quantized_fully_connected.h b/backends/cadence/generic/operators/op_quantized_fully_connected.h index 408fbabe726..fe562b55215 100644 --- a/backends/cadence/generic/operators/op_quantized_fully_connected.h +++ b/backends/cadence/generic/operators/op_quantized_fully_connected.h @@ -69,6 +69,7 @@ quantized_fully_connected_asym8uxasym8u_asym8u_per_tensor_out( const std::optional<::executorch::aten::Tensor>& offset, ::executorch::aten::Tensor& out); + } // namespace native } // namespace generic } // namespace impl diff --git a/backends/cadence/generic/operators/op_quantized_linear.cpp b/backends/cadence/generic/operators/op_quantized_linear.cpp index 02ff97de74d..dd06824be5b 100644 --- a/backends/cadence/generic/operators/op_quantized_linear.cpp +++ b/backends/cadence/generic/operators/op_quantized_linear.cpp @@ -215,6 +215,7 @@ Tensor& quantized_linear_asym8uxasym8u_asym8u_per_tensor_out( return out; } + } // namespace native } // namespace generic } // namespace impl diff --git a/backends/cadence/generic/operators/op_quantized_linear.h b/backends/cadence/generic/operators/op_quantized_linear.h index 517357d5bf9..864e57dabab 100644 --- a/backends/cadence/generic/operators/op_quantized_linear.h +++ b/backends/cadence/generic/operators/op_quantized_linear.h @@ -69,6 +69,7 @@ quantized_linear_asym8uxasym8u_asym8u_per_tensor_out( const std::optional<::executorch::aten::Tensor>& offset, ::executorch::aten::Tensor& out); + } // namespace native } // namespace generic } // namespace impl