| # Owner(s): ["oncall: jit"] |
| |
| import os |
| import sys |
| |
| import torch |
| |
| # Make the helper files in test/ importable |
| pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) |
| sys.path.append(pytorch_test_dir) |
| from torch.testing._internal.jit_utils import JitTestCase, warmup_backward, FileCheck |
| |
| if __name__ == '__main__': |
| raise RuntimeError("This test file is not meant to be run directly, use:\n\n" |
| "\tpython test/test_jit.py TESTNAME\n\n" |
| "instead.") |
| |
| class TestProfiler(JitTestCase): |
| def setUp(self): |
| self.prev_exec = torch._C._jit_set_profiling_executor(True) |
| self.prev_profiling = torch._C._get_graph_executor_optimize(True) |
| self.inline_autodiff = torch._C._debug_set_autodiff_subgraph_inlining(False) |
| self.texpr_fuser_state = torch._C._jit_texpr_fuser_enabled() |
| self.can_fuse_on_cpu = torch._C._jit_can_fuse_on_cpu() |
| torch._C._jit_set_texpr_fuser_enabled(True) |
| torch._C._jit_override_can_fuse_on_cpu(True) |
| self.default_dtype = torch.get_default_dtype() |
| self.old_reduction_enabled = torch._C._jit_set_texpr_reductions_enabled(True) |
| torch.set_default_dtype(torch.double) |
| self.old_fusion_inlining = torch._C._debug_get_fusion_group_inlining() |
| torch._C._debug_set_fusion_group_inlining(False) |
| self.old_te_must_use_llvm_cpu = torch._C._jit_get_te_must_use_llvm_cpu() |
| torch._C._jit_set_te_must_use_llvm_cpu(False) |
| |
| def tearDown(self): |
| torch._C._jit_set_profiling_executor(self.prev_exec) |
| torch._C._get_graph_executor_optimize(self.prev_profiling) |
| torch._C._debug_set_autodiff_subgraph_inlining(self.inline_autodiff) |
| torch._C._jit_set_texpr_fuser_enabled(self.texpr_fuser_state) |
| torch._C._jit_override_can_fuse_on_cpu(self.can_fuse_on_cpu) |
| torch.set_default_dtype(self.default_dtype) |
| torch._C._jit_set_texpr_reductions_enabled(self.old_reduction_enabled) |
| torch._C._debug_set_fusion_group_inlining(self.old_fusion_inlining) |
| torch._C._jit_set_te_must_use_llvm_cpu(self.old_te_must_use_llvm_cpu) |
| |
| def test_tensor_type_not_determined_by_inputs(self): |
| @torch.jit.script |
| def scalar_type_input(x, y, z): |
| return x + y + 4 + z.item() |
| |
| x = torch.tensor([2, 2]) |
| scalar_type_input(x, x, torch.tensor(1)) |
| scalar_type_input(x, x, torch.tensor(1)) |
| scalar_type_input(x, x, torch.tensor(1.0)) |
| g = torch.jit.last_executed_optimized_graph() |
| |
| # item & add should not get pulled into the fusion group - |
| # we expect to see Fusion Group (item / add) Fusion Group in ir dump |
| FileCheck().check("TensorExpr").check("Scalar = aten::item").check_next("Tensor = aten::add").check("TensorExpr").run(g) |
| |
| |
| @torch.jit.script |
| def non_const_dtype(x, y, cond: bool): |
| dtype = torch.int16 if cond else torch.int32 |
| return (x + y + 3).sum(dtype=dtype) |
| |
| non_const_dtype(x, x, True) |
| non_const_dtype(x, x, True) |
| g = torch.jit.last_executed_optimized_graph() |
| # because dtype is non-const, sum should not get pulled into the Fusion Group |
| FileCheck().check("TensorExpr").check("TensorExpr").check_not("aten::sum").run(g) |
| |
| def test_specialize_backward(self): |
| def test_fuse(a, b): |
| c = a * b |
| d = c * b |
| return d |
| |
| test_fuse.__disable_jit_function_caching__ = True |
| |
| scripted_f = torch.jit.script(test_fuse) |
| x = torch.ones(1, requires_grad=True) |
| y = torch.ones(1, requires_grad=True) |
| scripted_f(x, y) |
| b = scripted_f(x, y) |
| warmup_backward(b) |
| g = torch.jit.last_executed_optimized_graph() |
| # Backward has an if node guarding specializations, |
| # within the if node true block there is only one if node |
| # that guards a tensorexpr group |
| optimized_block = next(g.findNode("prim::If").blocks()) |
| if_nodes = list(optimized_block.findAllNodes("prim::If")) |
| |
| self.assertEqual(len(if_nodes), 1) |
| FileCheck().check("Group[Subgraph").run(str(if_nodes[0])) |
| # no broadcasts occurred, sum_to_size have been specialized out |
| self.assertIsNone(optimized_block.findNode("aten::_grad_sum_to_size")) |
| |
| broadcast_f = torch.jit.script(test_fuse) |
| x = torch.ones([2, 2], requires_grad=True) |
| y = torch.ones([1], requires_grad=True) |
| broadcast_f(x, y) |
| b = broadcast_f(x, y) |
| b.backward(torch.ones([2, 2], dtype=torch.float), retain_graph=True) |
| b.backward(torch.ones([2, 2], dtype=torch.float)) |
| # warmup_backward(b, torch.ones([2, 2], dtype=torch.float)) |
| g = torch.jit.last_executed_optimized_graph() |
| optimized_block = next(g.findNode("prim::If").blocks()) |
| # broadcasts occurred, currently expect to see aten::_grad_sum_to_size |
| self.assertIsNotNone(optimized_block.findNode("aten::_grad_sum_to_size")) |
| |
| def test_specialized_types(self): |
| @torch.jit.script |
| def test_fuse(a, b): |
| c = a * b |
| d = c * b |
| return d |
| |
| x = torch.tensor([.5]) |
| for _ in range(3): |
| test_fuse(x, x) |
| |
| g = torch.jit.last_executed_optimized_graph() |
| # Types should remain specialized for typecheck outputs & fusion outputs |
| FileCheck().check("Double(").check_same("prim::TypeCheck").check_same("\n").check("Double").check_same("TensorExpr").run(g) |
| |
| # other outputs should not be specialized |
| FileCheck().check("Tensor = prim::If").run(g) |
| |
| def test_aliasing_merge(self): |
| @torch.jit.script |
| def foo(a, b): |
| c = a * b |
| d = c * b |
| d.add_(b) |
| e = d * b |
| return d + e |
| |
| x = torch.ones(1) |
| y = torch.ones(1) |
| foo(x, y) |
| b = foo(x, y) |
| g = torch.jit.last_executed_optimized_graph() |
| self.assertEqual(len(list(g.findAllNodes("prim::TypeCheck"))), 2) |
| FileCheck().check("TensorExpr").check("aten::add_").check("TensorExpr").run(g) |
| |
| def test_use_not_profiled(self): |
| def foo(t1, t2, t3, t4, t: float): |
| h = t1 + t2 + t3 + t4 |
| if t > 0.5: |
| # Putting a use of t1 in a never-executed conditional prevents |
| return t1 + 1 |
| return h |
| |
| t = torch.rand(8, dtype=torch.float) |
| |
| foo_script = torch.jit.script(foo) |
| for _ in range(torch._C._jit_get_num_profiled_runs() + 1): |
| foo_script(t, t, t, t, 0.1) |
| |
| self.assertEqual(foo(t, t, t, t, 0.1), foo_script(t, t, t, t, 0.1)) |
| g = torch.jit.last_executed_optimized_graph() |
| # all adds fused |
| FileCheck().check("graph").check_not("aten::add").check("prim::If").run(g) |
| |
| def test_not_fusing_scalar_ops(self): |
| @torch.jit.script |
| def foo(x: int, y: int): |
| return x + y + 2 + 4 + 5 + 6 |
| |
| foo(1, 2) |
| foo(2, 3) |
| g = torch.jit.last_executed_optimized_graph() |
| FileCheck().check_not("TensorExpr").run(g) |
| |
| def test_not_optimizing_property(self): |
| @torch.jit.script |
| def foo(x, y): |
| return x + y + 1 + 2 + 3, x.size() |
| |
| x = torch.ones(1) |
| foo(x, x) |
| foo(x, x) |
| g = torch.jit.last_executed_optimized_graph() |
| FileCheck().check("aten::size").run(g) |
| x = torch.ones([2, 3, 5]) |
| self.assertEqual(foo(x, x), (x + x + 1 + 2 + 3, x.size())) |
| |
| def test_fallback_graph_not_specialized(self): |
| @torch.jit.script |
| def foo(a, b): |
| c = a * b |
| d = c * b |
| e = d * b |
| return d + e |
| |
| x = torch.ones(1) |
| y = torch.ones(1) |
| foo(x, y) |
| foo(x, y) |
| g = torch.jit.last_executed_optimized_graph() |
| FileCheck().check("CallFunction").check_next("Tensor = prim::TupleUnpack").run(g) |
| |
| def test_autograd_fallback_graph(self): |
| @torch.jit.script |
| def foo(a, b): |
| c = a * b |
| d = c * b |
| e = d * b |
| return d + e |
| |
| x = torch.ones(1, requires_grad=True) |
| y = torch.ones(1, requires_grad=True) |
| foo(x, y) |
| b = foo(x, y) |
| b.backward(torch.ones([1], dtype=torch.float), retain_graph=True) |
| b.backward(torch.ones([1], dtype=torch.float)) |
| |
| g = torch.jit.last_executed_optimized_graph() |
| FileCheck().check("fallback_function").check_next("CallFunction").run(g) |
| |
| def test_tensor_constant(self): |
| def foo(a, b): |
| return a + b + torch.tensor([2]) |
| |
| x = torch.ones(1, requires_grad=False) |
| foo_script = torch.jit.script(foo) |
| foo_script(x, x) |
| foo_script(x, x) |
| |
| self.assertEqual(foo_script(x, x), foo(x, x)) |
| g = torch.jit.last_executed_optimized_graph() |
| FileCheck().check_count("aten::add", 2, exactly=True).run(g) |
| |
| def test_local_fusion_strategy(self): |
| @torch.jit.script |
| def foo(x): |
| return x + x + x |
| |
| torch.jit.set_fusion_strategy([("STATIC", 1)]) |
| for _ in range(3): |
| foo(torch.rand([10])) |
| |
| torch.jit.set_fusion_strategy([("STATIC", 10)]) |
| |
| for i in range(10): |
| foo(torch.rand([i])) |
| foo(torch.rand([i])) |
| |
| g = torch.jit.last_executed_optimized_graph() |
| FileCheck().check_count(":TensorExprGroup", 2, exactly=True).run(g) |
| |
| def test_iterative_fusion(self): |
| @torch.jit.script |
| def foo(a, b, c, d): |
| a = a + b |
| b.add_(3) |
| c = c + b + d |
| a = a + 1 |
| return a, c |
| |
| x = torch.ones(1, requires_grad=False) |
| foo(x, x, x, x) |
| foo(x, x, x, x) |
| |
| # when we iterate through the block, we will start |
| # by fusing a = a + b with a = a + 1 |
| # if we were to continue iteration from that fusion point, |
| # would miss the fusion opportunity of c = c + d + b |
| |
| g = torch.jit.last_executed_optimized_graph() |
| self.assertEqual(len(list(g.findAllNodes("prim::TensorExprGroup"))), 2) |
