diff --git a/backends/cadence/aot/reorder_ops.py b/backends/cadence/aot/reorder_ops.py index b72c0c6df8e..9542704c8ab 100644 --- a/backends/cadence/aot/reorder_ops.py +++ b/backends/cadence/aot/reorder_ops.py @@ -71,6 +71,12 @@ exir_ops.edge.aten.select_copy, } +supported_quantize_ops_overloadpkt = { + exir_ops.edge.quantized_decomposed.quantize_per_tensor, + exir_ops.edge.quantized_decomposed.quantize_per_channel, + exir_ops.edge.cadence.quantize_per_tensor, +} + class AdvanceQuantizeOpAboveDefInBranchPass(ExportPass): """ @@ -520,46 +526,43 @@ def postponing_feasible(self, dequant_node: torch.fx.Node): ): return False - dequant_shape = get_shape(self.graph_module, dequant_node) - slice_shapes = [ - shape - for user in users - if (shape := get_shape(self.graph_module, user)) - and ( - # skip slices that are the size of the sliced tensor itself. - # They should technically get removed in the later passes as nop. - shape is None - or dequant_shape is None - or prod(list(shape)) != prod(list(dequant_shape)) + quantized_branches = [] + for user in users: + slice_users = list(user.users) + quantized_branches.append( + bool(slice_users) + and all( + slice_user.op == "call_function" + and get_overload_packet(slice_user.target) + in supported_quantize_ops_overloadpkt + for slice_user in slice_users + ) ) - ] - if dequant_shape is not None and all( - shape is not None for shape in slice_shapes - ): - dequant_bytes = num_bytes_from_shape_and_dtype(dequant_shape, torch.float32) - slice_bytes = sum( - [ - num_bytes_from_shape_and_dtype(shape, torch.float32) - for shape in slice_shapes - ] + # Preserve the existing fallback for forks whose branches all requantize. + if all(quantized_branches): + return True + + dequant_shape = get_shape(self.graph_module, dequant_node) + if dequant_shape is None: + return False + + dequant_bytes = num_bytes_from_shape_and_dtype(dequant_shape, torch.float32) + surviving_dequant_bytes = 0 + for user, quantized_branch in zip(users, quantized_branches): + if quantized_branch: + continue + slice_shape = get_shape(self.graph_module, user) + if slice_shape is None: + return False + # Nop slices are removed later and do not add another materialized copy. + if prod(list(slice_shape)) == prod(list(dequant_shape)): + continue + surviving_dequant_bytes += num_bytes_from_shape_and_dtype( + slice_shape, torch.float32 ) - if slice_bytes <= dequant_bytes: - return True - - # If the users of each slice op is quantize op, then we can postpone - # dequantize, and convert slice -> dequantize -> quantize to - # slice -> requantize. - users = [x for y in users for x in y.users if x.op != "output"] - return all( - get_overload_packet(x.target) - in { - exir_ops.edge.quantized_decomposed.quantize_per_tensor, - exir_ops.edge.quantized_decomposed.quantize_per_channel, - exir_ops.edge.cadence.quantize_per_tensor, - } - for x in users - ) + + return surviving_dequant_bytes <= dequant_bytes def postpone_dequantize_op(self, graph_module: torch.fx.GraphModule) -> bool: # Different supported dequant ops have their own default variants @@ -598,7 +601,8 @@ def postpone_dequantize_op(self, graph_module: torch.fx.GraphModule) -> bool: graph.erase_node(node) modified = True - graph_module.recompile() + if modified: + graph_module.recompile() return modified def call(self, graph_module: torch.fx.GraphModule) -> PassResult: diff --git a/backends/cadence/aot/tests/test_reorder_ops_passes.py b/backends/cadence/aot/tests/test_reorder_ops_passes.py index 1decb01684c..be44d133a8c 100644 --- a/backends/cadence/aot/tests/test_reorder_ops_passes.py +++ b/backends/cadence/aot/tests/test_reorder_ops_passes.py @@ -608,6 +608,132 @@ def test_postpone_dequantize_branched(self) -> None: ), ) + def _build_mixed_slice_quant_graph( + self, third_branch_quantized: bool + ) -> tuple[torch.fx.GraphModule, tuple[torch.Tensor, ...]]: + builder = GraphBuilder() + x_data = torch.randint(0, 255, [12, 4], dtype=torch.uint8) + weights_data = torch.randn([4, 4], dtype=torch.float32) + + x = builder.placeholder("x", x_data) + weights = builder.placeholder("weights", weights_data) + dequant = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(x, 0.1, 10, 0, 255, torch.uint8), + ) + + outputs = [] + branch_specs = [ + (0, 8, (0.2, 5)), + (2, 10, None), + (4, 12, (0.3, 7) if third_branch_quantized else None), + ] + for start, end, target_qparams in branch_specs: + slice_node = builder.call_operator( + op=exir_ops.edge.aten.slice_copy.Tensor, + args=(dequant, 0, start, end), + ) + if target_qparams is not None: + scale, zero_point = target_qparams + output = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(slice_node, scale, zero_point, 0, 255, torch.uint8), + ) + else: + output = builder.call_operator( + op=exir_ops.edge.aten.mm.default, + args=(slice_node, weights), + ) + outputs.append(output) + + builder.output(outputs) + return builder.get_graph_module(), (x_data, weights_data) + + def test_postpone_dequantize_mixed_slice_quant_branches(self) -> None: + original_graph, inputs = self._build_mixed_slice_quant_graph(True) + result = transform_and_check_numerics( + original_graph, + inputs, + PostponeDequantizeOpBelowUseChainPass(), + ) + self.assertTrue(result.modified) + + converted_graph = result.graph_module + graph_nodes = list(converted_graph.graph.nodes) + slice_nodes = [ + node + for node in graph_nodes + if node.op == "call_function" + and node.target == exir_ops.edge.aten.slice_copy.Tensor + ] + dequant_nodes = [ + node + for node in graph_nodes + if node.op == "call_function" + and node.target + == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default + ] + self.assertEqual(len(slice_nodes), 3) + self.assertEqual(len(dequant_nodes), 3) + self.assertTrue( + all(node.meta["val"].dtype == torch.uint8 for node in slice_nodes) + ) + self.assertTrue( + all(tuple(node.meta["val"].shape) == (8, 4) for node in dequant_nodes) + ) + self.assertTrue( + all(node.meta["val"].dtype == torch.float32 for node in dequant_nodes) + ) + for slice_node in slice_nodes: + source = slice_node.args[0] + self.assertIsInstance(source, torch.fx.Node) + self.assertEqual(cast(torch.fx.Node, source).meta["val"].dtype, torch.uint8) + downstream_dequants = [ + user for user in slice_node.users if user in dequant_nodes + ] + self.assertEqual(len(downstream_dequants), 1) + self.assertLess( + graph_nodes.index(slice_node), + graph_nodes.index(downstream_dequants[0]), + ) + + fused_result = cast( + PassResult, FuseQuantDequantToRequantizePass()(converted_graph) + ) + self.assertTrue(fused_result.modified) + self.assertEqual( + count_node( + fused_result.graph_module, + exir_ops.edge.cadence.requantize.per_tensor, + ), + 2, + ) + self.assertEqual( + count_node( + fused_result.graph_module, + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ), + 1, + ) + + def test_postpone_dequantize_rejects_expensive_mixed_float_branches( + self, + ) -> None: + original_graph, inputs = self._build_mixed_slice_quant_graph(False) + result = transform_and_check_numerics( + original_graph, + inputs, + PostponeDequantizeOpBelowUseChainPass(), + ) + self.assertFalse(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + ), + 1, + ) + # 4d -> permute -> 4d -> view -> 3d def test_permute3_view4_chains(self) -> None: builder = GraphBuilder()