Fix and re-enable eliminate_permutations graph transformation

Summary:
Fix eliminate_permutations graph transformation in AITemplate to retain successive cancelling permutations on strided input Tensor.

Re-enabled the transformation as well.

Differential Revision: D47371853

fbshipit-source-id: 5aeffcb4823c1db25641c7a61dc039dea79fed47

python/aitemplate/compiler/transform/optimize_graph.py

@@ -51,6 +51,7 @@
 from aitemplate.compiler.transform.transform_odd_alignment import (
     transform_odd_alignment,
 )
+from aitemplate.compiler.transform.transform_permutations import eliminate_permutations
 from aitemplate.compiler.transform.transform_permute_to_reshape import (
     transform_permute_to_reshape,
 )
@@ -125,8 +126,7 @@ def optimize_graph(
         split_large_split_ops,
         transform_permute_to_reshape,
         transform_memory_ops,
-        # FIXME: temporarily disable this due to some accuracy issue
-        # eliminate_permutations,
+        eliminate_permutations,
     ]
 
     if not optimize:

python/aitemplate/compiler/transform/transform_permutations.py

@@ -76,8 +76,8 @@ def eliminate_permutations(
             input_accessors = cur_op._attrs.get("input_accessors", None)
             if (
                 input_accessors is not None
-                and hasattr(input_accessors[0], "strided_dim")
-                and input_accessors[0].strided_dim is not None
+                and hasattr(input_accessors[0], "stride_dim")
+                and input_accessors[0].stride_dim is not None
             ):
                 continue
             curr_op_output = cur_op._attrs["outputs"][0]

tests/unittest/compiler/test_eliminate_permutations.py

@@ -28,7 +28,6 @@
 )
 
 
-@unittest.skip("Skip until we fix the accuracy issue")
 class EliminatePermutationTestCase(unittest.TestCase):
     def test_eliminate_permutation(self):
         dtype = "float"
@@ -190,6 +189,41 @@ def test_eliminate_permutation_all_permutations(self):
         self.assertEqual(len(result_graph), 3)
         self.assertTrue(graph_has_op(result_graph, "permute"))
 
+    def test_do_not_eliminate_permutation_of_strided_input(self):
+        dtype = "float"
+        shape = [3, 2, 4]
+        new_shape = [3, 2 * 2]
+        target = detect_target()
+
+        x = Tensor(shape, name="x", dtype=dtype, is_input=True)
+        s1 = ops.dynamic_slice()(
+            x, start_indices=[0, 0, 2], end_indices=[2147483647, 2147483647, 4]
+        )
+        p1 = ops.permute()(s1, dims=[0, 2, 1])
+        p2 = ops.permute()(p1, dims=[0, 2, 1])
+        z = ops.reshape()(p2, new_shape)
+        z._attrs["is_output"] = True
+        z._attrs["name"] = "z"
+
+        with compile_model(
+            z, target, "./tmp", "test_do_not_eliminate_permutation_of_strided_input"
+        ) as module:
+            # Verify the generated graph.
+            sorted_graph = module.debug_sorted_graph
+            self.assertEqual(len(sorted_graph), 4)
+            self.assertTrue(graph_has_op(sorted_graph, "permute021"))
+
+            x_pt = get_random_torch_tensor(shape, dtype)
+            z_pt = get_torch_empty_tensor(new_shape, dtype)
+
+            module.run_with_tensors({"x": x_pt}, {"z": z_pt})
+
+            self.assertTrue(
+                torch.equal(
+                    torch.reshape(torch.split(x_pt, 2, dim=2)[1], new_shape), z_pt
+                )
+            )
+
 
 if __name__ == "__main__":
     unittest.main()