diff --git a/backends/vulkan/patterns/quantized_linear.py b/backends/vulkan/patterns/quantized_linear.py index 86a35298fa4..7050d060404 100644 --- a/backends/vulkan/patterns/quantized_linear.py +++ b/backends/vulkan/patterns/quantized_linear.py @@ -236,6 +236,27 @@ def is_weight_perchannel_quantized(self) -> bool: # scales should have same size as weight's output channels dim return scales_shape[0] == weight_shape[-2] + def get_fp_input_node(self) -> Optional[torch.fx.Node]: + """ + Returns the floating point tensor that the pattern's input quantization + consumes, or None if it is not available in the graph. + + Some custom ops (i.e. et_vk.linear_q8ta_q8csw) quantize the activation + tensor themselves, so they need the floating point input rather than the + quantized one. For dynamically quantized inputs pattern_input_node is + already the floating point tensor; for statically quantized inputs it is + the quantized output of a quantize node, so step back past it. + """ + if self.dequantize_input_node is None or self.quantize_input_node is not None: + return self.pattern_input_node + + if utils.is_quant_node(self.pattern_input_node): + return self.pattern_input_node.args[0] # pyre-ignore[7] + + # The quantized tensor is produced outside of the graph, so there is no + # floating point tensor to use. + return None + def is_input_static_per_tensor_quantized(self) -> bool: if self.dequantize_input_node is None: return False @@ -503,12 +524,34 @@ def make_linear_q8ta_q8csw_custom_op( data=sum_per_output_channel, ) + # This op quantizes the activation tensor itself, so it takes the floating + # point input rather than the quantized one. + fp_input_node = match.get_fp_input_node() + if fp_input_node is None: + # The input is only available as a quantized tensor, i.e. it is produced + # by a preceding quantized op. Dequantize it so that the op has the + # floating point tensor it expects. + with graph_module.graph.inserting_before(match.output_node): + fp_input_node = graph_module.graph.create_node( + "call_function", + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=( + match.pattern_input_node, + match.input_scales_node, + match.input_zeros_node, + -128, + 127, + torch.int8, + ), + ) + fp_input_node.meta["val"] = match.dequantize_input_node.meta["val"] + with graph_module.graph.inserting_before(match.output_node): qlinear_node = graph_module.graph.create_node( "call_function", exir_ops.edge.et_vk.linear_q8ta_q8csw.default, args=( - match.pattern_input_node, + fp_input_node, match.input_scales_node, match.input_zeros_node, match.weight_node, diff --git a/backends/vulkan/test/test_vulkan_passes.py b/backends/vulkan/test/test_vulkan_passes.py index f030b9268a1..5d7a9a1290f 100644 --- a/backends/vulkan/test/test_vulkan_passes.py +++ b/backends/vulkan/test/test_vulkan_passes.py @@ -3,6 +3,8 @@ import torch +import executorch.backends.vulkan.utils as utils + from executorch.backends.vulkan._passes.fuse_patterns import FusePatternsPass from executorch.exir import EdgeCompileConfig, EdgeProgramManager, to_edge @@ -637,6 +639,100 @@ def forward(self, x): f"q8ta_linear[{i}].output_zero_point should equal q8ta_linear[{i + 1}].input_zero_point", ) + def test_linear_q8ta_q8csw_takes_floating_point_input(self): + """et_vk.linear_q8ta_q8csw quantizes its activation itself, so it must be + given the floating point input rather than the quantized one. + + QuantizedLinear.cpp names that argument fp_input and passes it to + add_quantize_and_pack_4h4w_node. Before this was fixed the fusion passed + the int8 output of the input quantize node instead, and the delegate + failed at inference looking for a shader variant that takes an already + quantized input, e.g. clone_buffer_to_image_int8_int32. + """ + # The pattern is built directly rather than with a quantizer: + # XNNPACKQuantizer also quantizes the linear's output, which produces + # q8ta_linear instead, and VulkanQuantizer offers no static activation + # quantization mode. + qd = torch.ops.quantized_decomposed + in_features, out_features, batch = 256, 128, 4 + act_scale = 0.05 + + weight = torch.randn(out_features, in_features) + weight_scales = weight.abs().amax(dim=1).clamp(min=1e-8) / 127.0 + weight_zeros = torch.zeros(out_features, dtype=torch.int64) + weight_q = qd.quantize_per_channel.default( + weight, weight_scales, weight_zeros, 0, -127, 127, torch.int8 + ) + + class StaticQuantLinear(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight_q", weight_q) + self.register_buffer("weight_scales", weight_scales) + self.register_buffer("weight_zeros", weight_zeros) + + def forward(self, x): + xq = qd.quantize_per_tensor.default( + x, act_scale, 0, -128, 127, torch.int8 + ) + xdq = qd.dequantize_per_tensor.default( + xq, act_scale, 0, -128, 127, torch.int8 + ) + wdq = qd.dequantize_per_channel.default( + self.weight_q, + self.weight_scales, + self.weight_zeros, + 0, + -127, + 127, + torch.int8, + ) + return torch.nn.functional.linear(xdq, wdq) + + model = StaticQuantLinear().eval() + sample_inputs = (torch.randn(batch, in_features),) + + edge_program = to_edge( + torch.export.export(model, sample_inputs, strict=True), + compile_config=EdgeCompileConfig( + _skip_dim_order=False, + _check_ir_validity=False, + ), + ) + + ep = edge_program._edge_programs["forward"] + fuse_pass = FusePatternsPass() + fuse_pass._exported_program = ep + self.assertTrue(fuse_pass.call(ep.graph_module).modified) + + gm = ep.graph_module + + # With no output quantization the linear becomes linear_q8ta_q8csw + # rather than q8ta_linear. + q8csw_nodes = [ + node + for node in gm.graph.nodes + if get_target_canonical_name(node) == "linear_q8ta_q8csw.default" + ] + self.assertEqual( + len(q8csw_nodes), + 1, + "Expected the output-unquantized linear to fuse to linear_q8ta_q8csw", + ) + + input_node = q8csw_nodes[0].args[0] + self.assertIsInstance(input_node, torch.fx.Node) + self.assertFalse( + utils.is_quant_node(input_node), + "linear_q8ta_q8csw was given the quantized activation; it expects the " + "floating point one", + ) + self.assertNotEqual( + input_node.meta["val"].dtype, + torch.int8, + "linear_q8ta_q8csw input must not be int8", + ) + def test_fuse_q8ta_linear_gemv_non_aligned_oc(self): """Test that quantized linear with non-aligned output channels (not multiple of 4) fuses correctly.""" from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (