Add peephole optimization for type_as operators. (#9316)

Runtian Zhou · facebook-github-bot · commit 66fe3b5c068e · 2018-07-16T10:26:56.000-07:00
Summary: If the type_as operator takes in two values with the same type, remove that operator. Pull Request resolved: #9316 Reviewed By: zdevito Differential Revision: D8808355 fbshipit-source-id: 2d5710a6380b22f4568fc38a439061b5340c4eb1
diff --git a/test/expect/TestJit.test_peephole.expect b/test/expect/TestJit.test_peephole.expect
@@ -0,0 +1,4 @@
+graph(%0 : Double(1)
+      %1 : Double(1)) {
+  return (%0);
+}
diff --git a/test/expect/TestJit.test_peephole_cuda-different_device.expect b/test/expect/TestJit.test_peephole_cuda-different_device.expect
@@ -0,0 +1,5 @@
+graph(%0 : Double(1)
+      %1 : Double(1)) {
+  %2 : Double(1) = aten::type_as(%0, %1)
+  return (%2);
+}
diff --git a/test/expect/TestJit.test_peephole_cuda-same_device.expect b/test/expect/TestJit.test_peephole_cuda-same_device.expect
@@ -0,0 +1,4 @@
+graph(%0 : Double(1)
+      %1 : Double(1)) {
+  return (%0);
+}
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -256,6 +256,48 @@ def f(x, y):
         self.assertExpectedGraph(trace)
         self.assertExportImport(trace, (x, y))
 
+    def test_peephole(self):
+        a = torch.tensor([0.4], requires_grad=True)
+        b = torch.tensor([0.7], requires_grad=True)
+        c = torch.tensor([0], dtype=torch.int32)
+
+        def f(x, y):
+            return x.type_as(y)
+
+        trace, z = torch.jit.get_trace_graph(f, (a, b))
+        self.run_pass('peephole', trace)
+        self.assertExpectedGraph(trace)
+        trace, z = torch.jit.get_trace_graph(f, (a, c))
+        s = str(trace)
+        self.run_pass('peephole', trace)
+        self.assertEqual(s, str(trace))
+
+    def test_peephole_dynamic(self):
+        def f(x, y):
+            return x.type_as(y)
+
+        fn = torch.jit.script(f)
+        s = str(fn.graph)
+        torch._C._jit_pass_peephole(fn.graph)
+        self.assertEqual(s, str(fn.graph))
+
+    @unittest.skipIf(not RUN_CUDA, "cpp tests require CUDA")
+    def test_peephole_cuda(self):
+        a = torch.tensor([0.4], requires_grad=True, device='cpu')
+        b = torch.tensor([0.7], requires_grad=True, device='cuda')
+        c = torch.tensor([0.7], requires_grad=True, device='cuda')
+
+        def f(x, y):
+            return x.type_as(y)
+
+        trace, z = torch.jit.get_trace_graph(f, (a, c))
+        s = str(trace)
+        self.run_pass('peephole', trace)
+        self.assertEqual(s, str(trace))
+        trace, z = torch.jit.get_trace_graph(f, (b, c))
+        self.run_pass('peephole', trace)
+        self.assertExpectedGraph(trace, subname="same_device")
+
     def test_index(self):
         x = torch.tensor([0.4], requires_grad=True)
         y = torch.tensor([0], dtype=torch.int64)
diff --git a/torch/csrc/jit/passes/peephole.cpp b/torch/csrc/jit/passes/peephole.cpp
@@ -41,6 +41,21 @@ void PeepholeOptimize(Block * block) {
           // Let DCE clean up any unused nodes at this point
         }
       } break;
+      case aten::type_as: {
+        JIT_ASSERT(n->inputs().size() == 2);
+        Value *lhs = n->input(0);
+        Value *rhs = n->input(1);
+        // If LHS and RHS have the same static type, remove the type_as operator.
+        if (lhs->type()->kind() == TypeKind::TensorType &&
+            rhs->type()->kind() == TypeKind::TensorType) {
+           auto ltype = (*lhs->type()).cast<TensorType>();
+           auto rtype = (*rhs->type()).cast<TensorType>();
+           if(ltype->device() == rtype->device() &&
+              ltype->scalarType() == rtype->scalarType()) {
+              n->output()->replaceAllUsesWith(lhs);
+           }
+        }
+      } break;
       // Fuse mm + add into addmm
       case aten::add: {
         // Must have two inputs
