Add synchronize and check for contiguous

Author: feihugis

tests/optimizer/jit/test_einsum_rewriter.py

@@ -1,16 +1,15 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-import functools
 import logging
-import timeit
+import time
 
 from absl.testing import absltest, parameterized
 import torch
 from torch import Tensor
 
 from fastseq.logging import get_logger
-from fastseq.optimizer.jit.einsum_rewriter import rewrite_einsum
+from fastseq.optimizer.jit.einsum_rewriter import rewrite_einsum, einsum_rewrite_pattern_0
 from fastseq.utils.test_utils import TestCaseBase
 
 logger = get_logger(__name__, logging.INFO)
@@ -39,11 +38,15 @@ def run_einsum(eqn: str, t0: Tensor, t1: Tensor):
 
         t0 = torch.randn(shape0, dtype=torch.float32).cuda()
         t1 = torch.randn(shape1, dtype=torch.float32).cuda()
-        repeat_times = 1000
+        repeat_times = 1024
 
         r0 = run_einsum(eqn, t0, t1)
-        time0 = timeit.Timer(functools.partial(run_einsum, eqn, t0, t1))
-        s0 = time0.timeit(repeat_times)
+        torch.cuda.synchronize()
+        start0 = time.time()
+        for _ in range(repeat_times):
+            run_einsum(eqn, t0, t1)
+        torch.cuda.synchronize()
+        end0 = time.time()
 
         script_run_einsum = torch.jit.script(run_einsum)
         logger.debug(f"Original graph: \n{script_run_einsum.graph.str()}")
@@ -52,13 +55,28 @@ def run_einsum(eqn: str, t0: Tensor, t1: Tensor):
         self.assertTrue('bmm' in script_run_einsum.graph.str())
 
         r1 = script_run_einsum(eqn, t0, t1)
-        time1 = timeit.Timer(
-            functools.partial(script_run_einsum, eqn, t0, t1))
-        s1 = time1.timeit(repeat_times)
+        torch.cuda.synchronize()
+        start1 = time.time()
+        for _ in range(repeat_times):
+            script_run_einsum(eqn, t0, t1)
+        torch.cuda.synchronize()
+        end1 = time.time()
+
+        r2 = einsum_rewrite_pattern_0(eqn, [t0, t1])
+        torch.cuda.synchronize()
+        start2 = time.time()
+        for _ in range(repeat_times):
+            einsum_rewrite_pattern_0(eqn, [t0, t1])
+        torch.cuda.synchronize()
+        end2 = time.time()
 
         self.assertTrue(torch.equal(r0, r1))
-        logger.info(f"einsum took: {s0}; optimized einsum torchscript took: "
-                    f"{s1};")
+        self.assertTrue(torch.equal(r0, r2))
+        self.assertEqual(
+            r0.is_contiguous(), r1.is_contiguous(), r2.is_contiguous())
+        logger.info(f"einsum took: {end0 - start0};"
+                    f"optimized einsum torchscript took: {end1 - start1};"
+                    f"optimized einsum python took: {end2 - start2};")
 
 
 if __name__ == "__main__":