# Owner(s): ["module: fx"] import operator import torch import torch.fx from torch.fx.experimental import const_fold from torch.fx.passes.shape_prop import _extract_tensor_metadata, ShapeProp from torch.testing._internal.common_utils import TestCase class TestConstFold(TestCase): def _get_attr(self, node): mod = node.graph.owning_module target = str(node.target) target_atoms = target.split(".") curr_obj = mod for i, atom in enumerate(target_atoms): if not hasattr(curr_obj, atom): raise RuntimeError( f"Node referenced nonexistent target '{'.'.join(target_atoms[:i])}'; " f" original whole target: '{target}'" ) curr_obj = getattr(curr_obj, atom) return curr_obj def _verify_const_fold_mod(self, mod_folded: const_fold.FoldedGraphModule): self.assertTrue(mod_folded.const_subgraph_module is not None) # Check that we don't have the const or non-const fold graphs in the gm, and # that we do have the const folded get_attr. found_folded_attrs = False for n in mod_folded.graph.nodes: if n.op == "get_attr" and n.target.startswith("_FX_CONST_FOLDED_ATTRS"): found_folded_attrs = True elif n.op == "call_module": self.assertTrue(n.target not in {"submod_0", "submod_1"}) self.assertTrue(found_folded_attrs) def test_const_fold_basic_one_attr_no_name_collision(self): r""" Perform constant folding conversion, from original mod to split constant folding module with two split subgraphs, where there's a single attr to fold and a single output attr result to replace. attr1 attr1 | | | | x add add \ / | sub y output (becomes attr add_1) \ / ==> -------+------- (const/base subgraph split) mul attr2 x / (input from previous subgraph \ / \ / is attr) add sub y | \ / output mul attr2 \ / add | output """ class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.attr_1 = torch.nn.Parameter(torch.tensor([[-0.9]])) self.attr_2 = torch.nn.Parameter(torch.tensor([[17.1]])) def forward(self, x, y): a = self.attr_1 + self.attr_1 x = x - a return x * y + self.attr_2 mod = ConstFoldTestModule() mod_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(mod) self._verify_const_fold_mod(mod_folded) # Now run both folded and non-folded to check results equal. in_x, in_y = torch.tensor([[-0.45]]), torch.tensor([0.9]) base_result = mod(in_x, in_y) fold_result = mod_folded(in_x, in_y) self.assertTrue(torch.equal(fold_result, base_result)) def test_const_fold_basic_one_attr_name_collision(self): r""" Perform constant folding conversion, from original mod to split constant folding module with two split subgraphs, where there's a single attr to fold and a single output attr result to replace. Name the attrs such that they will collide by name with folded attrs. add_1 add_1 | | | | x add add \ / | sub y output (becomes attr add_1) \ / ==> -------+------- (const/base subgraph split) mul add_2 x / (input from previous subgraph \ / \ / is attr) add sub y | \ / output mul add_2 \ / add | output """ class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() # Note: Named as such to result in name collision. self.add_1__CF = torch.nn.Parameter(torch.tensor([[1.0]])) self.add_2__CF = torch.nn.Parameter(torch.tensor([[17.1]])) def forward(self, x, y): a = self.add_1__CF + self.add_1__CF x = x - a return x * y + self.add_2__CF mod = ConstFoldTestModule() mod_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(mod) self._verify_const_fold_mod(mod_folded) # Now run both folded and non-folded to check results equal. in_x, in_y = torch.tensor([[5.0]]), torch.tensor([4.0]) base_result = mod(in_x, in_y) fold_result = mod_folded(in_x, in_y) self.assertTrue(torch.equal(fold_result, base_result)) def test_const_fold_basic_placeholder_reordered(self): """ Test code path where placeholder comes after normal op node in FX """ class ConstFoldTestModule(torch.nn.Module): def forward(self, x, y): return x * 2 + y mod = ConstFoldTestModule() mod = torch.fx.symbolic_trace(mod) yy = None for n in mod.graph.nodes: if n.op == "placeholder" and n.target == "y": yy = n elif yy is not None and n.op == "call_function": yy.prepend(n) break mod_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(mod) self.assertTrue(mod_folded.const_subgraph_module is None) # Now run both folded and non-folded to check results equal. in_x = torch.tensor([[-0.45]]) in_y = torch.tensor([[0.45]]) base_result = mod(in_x, in_y) fold_result = mod_folded(in_x, in_y) self.assertTrue(torch.equal(fold_result, base_result)) def test_const_fold_noop(self): r""" Check that a graph with no constant folding is handled correctly. x attr1 \ / sub | output """ class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.attr1 = torch.nn.Parameter(torch.tensor([[-0.9]])) def forward(self, x): return x - self.attr1 mod = ConstFoldTestModule() mod_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(mod) # Check that the folded graph module is None, since there was no folding to do. self.assertTrue(mod_folded.const_subgraph_module is None) # Now run both folded and non-folded to check results equal. in_x = torch.tensor([[-0.45]]) base_result = mod(in_x) fold_result = mod_folded(in_x) self.assertTrue(torch.equal(fold_result, base_result)) def test_const_fold_basic_two_attr_three_input(self): r""" Perform constant folding conversion, from original mod to split constant folding module with two split subgraphs, where there are two attrs to fold into a single output, and there are three placeholder inputs. attr1 attr2 attr1 attr2 \ / \ / x add add \ / | sub y output (becomes attr add_1) \ / ==> -------+------- (const/base subgraph split) mul z x / (input from previous subgraph \ / \ / is attr) div sub y | \ / output mul z \ / div | output """ class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.attr1 = torch.nn.Parameter(torch.tensor([[-0.9]])) self.attr1 = torch.nn.Parameter(torch.tensor([[1.32]])) def forward(self, x, y, z): a = self.attr1 + self.attr1 sub = x - a mul = sub * y return mul / z mod = ConstFoldTestModule() mod_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(mod) self._verify_const_fold_mod(mod_folded) # Now run both folded and non-folded to check results equal. in_x, in_y, in_z = ( torch.tensor([[-0.45]]), torch.tensor([0.9]), torch.tensor([1.1]), ) base_result = mod(in_x, in_y, in_z) fold_result = mod_folded(in_x, in_y, in_z) self.assertTrue(torch.equal(fold_result, base_result)) def test_const_fold_basic_two_attr(self): r""" Perform constant folding conversion, from original mod to split constant folding module with two split subgraphs, where there are two attrs to fold into a single output. attr1 attr2 attr1 attr2 \ / \ / x add add (becomes attr add_1) \ / ==> -------+------- (const/base subgraph split) sub x | (input from previous subgraph is attr) | \ / output sub | output """ class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.attr1 = torch.nn.Parameter(torch.randn(2, 3)) self.attr2 = torch.nn.Parameter(torch.randn(2, 3)) def forward(self, x): y = self.attr1 + self.attr2 return x + y mod = ConstFoldTestModule() mod_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(mod) self._verify_const_fold_mod(mod_folded) # Now run both folded and non-folded to check results equal. in_x = torch.randn(2, 3) fold_result = mod_folded(in_x) base_result = mod(in_x) self.assertTrue(torch.equal(fold_result, base_result)) def test_const_fold_multi_const_folded_attrs(self): r""" Perform constant folding conversion, from original mod to split constant folding module with two split subgraphs, where there are two attrs to fold into two new attrs. attr1 attr2 attr1 attr2 / \ | / \ | permute | sum permute | sum \ / / \ / | x add y / add | \ / \ / | | sub add output output (become attrs add_1 and mul_1) \ / ==> --------+-------+------ (const/base subgraph split) \ / x | y | (inputs from previous subgraph add \ / \ / are attrs) | sub add linear \ / | add sigmoid | | linear output | sigmoid | output """ class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.attr1 = torch.nn.Parameter(torch.randn(4, 4)) self.attr2 = torch.nn.Parameter(torch.randn(4, 4)) self.lin = torch.nn.Linear(4, 4) def forward(self, x, y): a = self.attr1 + self.attr1.permute(1, 0) x = x - a amax = torch.sum(self.attr2, dim=1) y = y + amax return torch.sigmoid(self.lin(x + y)) mod = ConstFoldTestModule() mod_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(mod) self._verify_const_fold_mod(mod_folded) # Now run both folded and non-folded to check results equal. in_x, in_y = torch.randn(4, 4), torch.randn(4) fold_result = mod_folded(in_x, in_y) base_result = mod(in_x, in_y) self.assertTrue(torch.equal(fold_result, base_result)) def test_const_fold_submod_hierarchy(self): r""" Perform constant folding conversion, from original mod to split constant folding module where one of the folded attrs comes from a submod deeper in the hierarchy of the base module. """ class TracedThroughModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.internal_attr = torch.nn.Parameter(torch.randn(2, 3)) def forward(self): return self.internal_attr class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.my_mod = TracedThroughModule() self.attr = torch.nn.Parameter(torch.randn(2, 3)) def forward(self, x): return self.attr + self.my_mod() + x mod = ConstFoldTestModule() mod_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(mod) self._verify_const_fold_mod(mod_folded) # Now run both folded and non-folded to check results equal. in_x = torch.randn(2, 3) fold_result = mod_folded(in_x) base_result = mod(in_x) self.assertTrue(torch.equal(fold_result, base_result)) def test_retain_node_meta(self): r""" Perform constant folding conversion, and validate that node meta is retained. """ class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.attr = torch.nn.Parameter(torch.randn(2, 3)) def forward(self, x): a = self.attr + self.attr return x - a mod = ConstFoldTestModule() gm = torch.fx.symbolic_trace(mod) # Add a count for each node to check after we const fold. for idx, node in enumerate(gm.graph.nodes): if node.op != "output": node.meta["meta_idx"] = idx # Pre-folding: # idx 0: placeholder # idx 1: get_attr (will no longer be used, hence removed) # idx 2: add (will be folded into a get_attr) # idx 3: sub gm_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(gm) self._verify_const_fold_mod(gm_folded) # Post-folding: # idx 0: placeholder # idx 2: get_attr (replaced original add; original get_attr was removed) # idx 3: sub # Check the expected indices are still here. for node in gm_folded.graph.nodes: if node.op == "placeholder": self.assertEqual(node.meta["meta_idx"], 0) elif node.op == "get_attr": self.assertEqual(node.meta["meta_idx"], 2) elif node.op == "call_function" and node.target == operator.sub: self.assertEqual(node.meta["meta_idx"], 3) else: self.assertEqual(node.op, "output") # Now run both folded and non-folded to check results equal. in_x = torch.randn(2, 3) fold_result = gm_folded(in_x) base_result = mod(in_x) self.assertTrue(torch.equal(fold_result, base_result)) def test_const_fold_has_inlined_call_module_node(self): class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.attr = torch.nn.Parameter(torch.randn(2, 3)) self.mod = torch.nn.Identity() self.mod.relu = torch.nn.ReLU() def forward(self, x): a = self.attr + self.attr return self.mod.relu(x - a) mod = ConstFoldTestModule() gm_folded = const_fold.split_const_subgraphs(mod) # Now run both folded and non-folded to check results equal. in_x = torch.randn(2, 3) fold_result = gm_folded(in_x) base_result = mod(in_x) self.assertTrue(torch.equal(fold_result, base_result)) def test_const_fold_module_attr(self): class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.const = torch.nn.Parameter(torch.randn(2, 3)) self.mod = torch.nn.Identity() self.mod.attr = torch.nn.Parameter(torch.randn(2, 3)) def forward(self, x): a = self.const + self.mod.attr x = x + a return x + self.mod.attr mod = ConstFoldTestModule() gm_folded = const_fold.split_const_subgraphs(mod) # Now run both folded and non-folded to check results equal. in_x = torch.randn(2, 3) fold_result = gm_folded(in_x) base_result = mod(in_x) self.assertTrue(torch.equal(fold_result, base_result)) def test_const_fold_unused_placeholder(self): class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.const = torch.nn.Parameter(torch.randn(2, 3)) def forward(self, x, y, z): a = self.const + self.const return y + a mod = ConstFoldTestModule() gm_folded = const_fold.split_const_subgraphs(mod) # Now run both folded and non-folded to check results equal. in_x = torch.randn(2, 3) fold_result = gm_folded(in_x, in_x, in_x) base_result = mod(in_x, in_x, in_x) self.assertTrue(torch.equal(fold_result, base_result)) def test_dict_output(self): class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.const = torch.nn.Parameter(torch.randn(2, 3)) def forward(self, x): a = self.const + self.const return {"result": x + a} mod = ConstFoldTestModule() gm_folded = const_fold.split_const_subgraphs(mod) # Now run both folded and non-folded to check results equal. in_x = torch.randn(2, 3) fold_result = gm_folded(in_x) base_result = mod(in_x) self.assertTrue(torch.equal(fold_result["result"], base_result["result"])) def test_two_outputs(self): class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.const = torch.nn.Parameter(torch.randn(2, 3)) def forward(self, x): a = self.const + self.const return x, x + a mod = ConstFoldTestModule() gm_folded = const_fold.split_const_subgraphs(mod) # Now run both folded and non-folded to check results equal. in_x = torch.randn(2, 3) fold_result = gm_folded(in_x) base_result = mod(in_x) self.assertTrue(torch.equal(fold_result[0], base_result[0])) self.assertTrue(torch.equal(fold_result[1], base_result[1])) def test_three_outputs(self): class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.const = torch.nn.Parameter(torch.randn(2, 3)) def forward(self, x): a = self.const + self.const return x, x + a, x + a mod = ConstFoldTestModule() gm_folded = const_fold.split_const_subgraphs(mod) # Now run both folded and non-folded to check results equal. in_x = torch.randn(2, 3) fold_result = gm_folded(in_x) base_result = mod(in_x) self.assertTrue(torch.equal(fold_result[0], base_result[0])) self.assertTrue(torch.equal(fold_result[1], base_result[1])) self.assertTrue(torch.equal(fold_result[2], base_result[2])) def test_check_inline_non_const(self): r""" Perform constant folding conversion and check that the non-const module is inlined correctly. """ class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.attr = torch.nn.Parameter(torch.randn(2, 3)) def forward(self, x): a = self.attr + self.attr return (x - a * x) / 2 mod = ConstFoldTestModule() gm = torch.fx.symbolic_trace(mod) gm_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(gm) self._verify_const_fold_mod(gm_folded) # Check there are no call modules, because they've been inlined or extracted for # const folding. for node in gm_folded.graph.nodes: self.assertNotEqual(node.op, "call_module") # Now run both folded and non-folded to check results equal. in_x = torch.randn(2, 3) fold_result = gm_folded(in_x) base_result = mod(in_x) self.assertTrue(torch.equal(fold_result, base_result)) def test_check_inline_non_const_mult_return(self): r""" Perform constant folding conversion and check that the non-const module is inlined correctly. """ class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.attr = torch.nn.Parameter(torch.randn(2, 3)) def forward(self, x): a = self.attr + self.attr return x - a, x / 2 mod = ConstFoldTestModule() gm = torch.fx.symbolic_trace(mod) gm_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(gm) self._verify_const_fold_mod(gm_folded) # Check there are no call modules, because they've been inlined or extracted for # const folding. for node in gm_folded.graph.nodes: self.assertNotEqual(node.op, "call_module") # Now run both folded and non-folded to check results equal. in_x = torch.randn(2, 3) fold_result = gm_folded(in_x) base_result = mod(in_x) self.assertTrue(torch.equal(fold_result[0], base_result[0])) self.assertTrue(torch.equal(fold_result[1], base_result[1])) def test_check_skip_folding_quant_dequant_pattern(self): r""" Set up skip_folding_quant_dequant function to skip quant/dequant pattern. This example shows how to use skip_folding_node_fn. """ class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.weight = torch.nn.Parameter(torch.randn(4, 4)) self.bias = torch.nn.Parameter(torch.randn(4)) self.relu = torch.nn.ReLU() def forward(self, x): quant_weight = torch.quantize_per_tensor( self.weight, 0.5, 3, torch.quint8 ) dequant_weight = torch.dequantize(quant_weight) output = torch.nn.functional.linear(x, dequant_weight, self.bias) return self.relu(output) mod = ConstFoldTestModule() in_x = torch.randn(2, 4) gm = torch.fx.symbolic_trace(mod) def skip_folding_quant_dequant(node: torch.fx.Node): if node.target != torch.quantize_per_tensor: return False # If quantize_per_node -> dequantize, then skip folding. for user in node.users: if user.target == torch.dequantize: return True return False gm_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs( gm, skip_folding_node_fn=skip_folding_quant_dequant ) # Check that the folded graph module is None, since there was no folding to do. self.assertTrue(gm_folded.const_subgraph_module is None) # Now run both folded and non-folded to check results equal. fold_result = gm_folded(in_x) base_result = mod(in_x) self.assertTrue(torch.equal(fold_result, base_result)) def test_fold_module(self): r""" Perform constant folding with a call_module node. """ class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.lin_input = torch.nn.Parameter(torch.randn(4, 4)) self.lin = torch.nn.Linear(4, 4) def forward(self, x): return self.lin(self.lin_input) + x mod = ConstFoldTestModule() mod_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(mod) self._verify_const_fold_mod(mod_folded) # Now run both folded and non-folded to check results equal. inp = torch.randn(4, 4) self.assertTrue(torch.equal(mod_folded(inp), mod(inp))) def test_const_fold_tensor_meta(self): self._test_const_fold_tensor_meta(True) self._test_const_fold_tensor_meta(False) def _test_const_fold_tensor_meta(self, requires_grad): """ Verify tensor_meta is handled correctly. """ class ConstFoldTestModule(torch.nn.Module): def __init__(self) -> None: super().__init__() self.attr_1 = torch.nn.Parameter(torch.tensor([[-0.9]]), requires_grad) self.attr_2 = torch.nn.Parameter(torch.tensor([[17.1]]), requires_grad) def forward(self, x, y): a = self.attr_1 + self.attr_1 x = x - a return x * y + self.attr_2 mod = ConstFoldTestModule() gm = torch.fx.symbolic_trace(mod) in_x, in_y = torch.tensor([[-0.45]]), torch.tensor([0.9]) ShapeProp(gm).propagate(in_x, in_y) mod_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs( gm, device_for_folded_attrs="cpu" ) self._verify_const_fold_mod(mod_folded) mod_folded.run_folding() for n in mod_folded.graph.nodes: if n.op == "get_attr": attr = self._get_attr(n) self.assertEqual(_extract_tensor_metadata(attr), n.meta["tensor_meta"]) # Now run both folded and non-folded to check results equal. base_result = mod(in_x, in_y) fold_result = mod_folded(in_x, in_y) self.assertTrue(torch.equal(fold_result, base_result))