From 1e431c6a909f85da21ad998cabcffa85637b8ebb Mon Sep 17 00:00:00 2001
From: zjgarvey <47986913+zjgarvey@users.noreply.github.com>
Date: Mon, 14 Oct 2024 14:41:31 -0500
Subject: [PATCH] Add AtenSliceTOp Canonicalization to
 SimplifyShapeCalculations pass (#3791)

Some ops were failing to infer the static component of partially dynamic
shapes, and the cause was a missing aten.slice.t pattern.

The lit test included here is an IR dump created before
DropAbstractInterpCalculations for an unflatten op that was failing to
infer shapes before the change.
---
 .../Transforms/SimplifyShapeCalculations.cpp  |  1 +
 .../Torch/simplify-shape-calculations.mlir    | 39 +++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/lib/Dialect/Torch/Transforms/SimplifyShapeCalculations.cpp b/lib/Dialect/Torch/Transforms/SimplifyShapeCalculations.cpp
index f63fb4eb9..edf936bf3 100644
--- a/lib/Dialect/Torch/Transforms/SimplifyShapeCalculations.cpp
+++ b/lib/Dialect/Torch/Transforms/SimplifyShapeCalculations.cpp
@@ -198,6 +198,7 @@ class SimplifyShapeCalculationsPass
     AtenSizeOp::getCanonicalizationPatterns(patterns, context);
     AtenLenTOp::getCanonicalizationPatterns(patterns, context);
     AtenAddTOp::getCanonicalizationPatterns(patterns, context);
+    AtenSliceTOp::getCanonicalizationPatterns(patterns, context);
 
     // TODO: Debug visitation order to make this more efficient.
     // A single linear scan should suffice.
diff --git a/test/Dialect/Torch/simplify-shape-calculations.mlir b/test/Dialect/Torch/simplify-shape-calculations.mlir
index b7e7cf17b..59884616f 100644
--- a/test/Dialect/Torch/simplify-shape-calculations.mlir
+++ b/test/Dialect/Torch/simplify-shape-calculations.mlir
@@ -489,3 +489,42 @@ func.func @shape_calc_with_two_uses(%arg0: !torch.vtensor<[2],f32>) -> !torch.vt
 
   return %arg0 : !torch.vtensor<[2],f32>
 }
+
+// CHECK-LABEL: func.func @unflat_shape_partial_dyn
+// CHECK-DAG:  %[[INT768:.*]] = torch.constant.int 768
+// CHECK-DAG:  %[[INT0:.*]] = torch.constant.int 0
+// CHECK-DAG:  %[[INT1:.*]] = torch.constant.int 1
+// CHECK-DAG:  %[[INT4:.*]] = torch.constant.int 4
+// CHECK :     } shapes {
+// CHECK :       %[[SZE0:.*]] = torch.aten.size.int %arg0, %[[INT0]] : !torch.vtensor<[?,?,3072],f32>, !torch.int -> !torch.int
+// CHECK :       %[[SZE1:.*]] = torch.aten.size.int %arg0, %[[INT1]] : !torch.vtensor<[?,?,3072],f32>, !torch.int -> !torch.int
+// CHECK :       %[[LIST:.*]] = torch.prim.ListConstruct %[[SZE0]], %[[SZE1]], %[[INT4]], %[[INT768]] : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
+// CHECK :       torch.shape.calculate.yield.shapes %[[LIST]] : !torch.list<int>
+// CHECK :     } : !torch.vtensor<[?,?,4,768],f32>
+func.func @unflat_shape_partial_dyn(%arg0: !torch.vtensor<[?,?,3072],f32>) -> !torch.vtensor<[?,?,4,?],f32> {
+  %int768 = torch.constant.int 768
+  %int3072 = torch.constant.int 3072
+  %int0 = torch.constant.int 0
+  %int3 = torch.constant.int 3
+  %int1 = torch.constant.int 1
+  %none = torch.constant.none
+  %int-1 = torch.constant.int -1
+  %int2 = torch.constant.int 2
+  %int4 = torch.constant.int 4
+  %0 = torch.prim.ListConstruct %int4, %int-1 : (!torch.int, !torch.int) -> !torch.list<int>
+  %1 = torch.shape.calculate {
+    %2 = torch.aten.unflatten.int %arg0, %int2, %0 : !torch.vtensor<[?,?,3072],f32>, !torch.int, !torch.list<int> -> !torch.vtensor<[?,?,4,?],f32>
+    torch.shape.calculate.yield %2 : !torch.vtensor<[?,?,4,?],f32>
+  } shapes {
+    %2 = torch.aten.size.int %arg0, %int0 : !torch.vtensor<[?,?,3072],f32>, !torch.int -> !torch.int
+    %3 = torch.aten.size.int %arg0, %int1 : !torch.vtensor<[?,?,3072],f32>, !torch.int -> !torch.int
+    %4 = torch.prim.ListConstruct %2, %3, %int3072 : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
+    %5 = torch.prim.ListConstruct %int4, %int768 : (!torch.int, !torch.int) -> !torch.list<int>
+    %6 = torch.aten.slice.t %4, %none, %int2, %int1 : !torch.list<int>, !torch.none, !torch.int, !torch.int -> !torch.list<int>
+    %7 = torch.aten.add.t %6, %5 : !torch.list<int>, !torch.list<int> -> !torch.list<int>
+    %8 = torch.aten.slice.t %4, %int3, %none, %int1 : !torch.list<int>, !torch.int, !torch.none, !torch.int -> !torch.list<int>
+    %9 = torch.aten.add.t %7, %8 : !torch.list<int>, !torch.list<int> -> !torch.list<int>
+    torch.shape.calculate.yield.shapes %9 : !torch.list<int>
+  } : !torch.vtensor<[?,?,4,?],f32>
+  return %1 : !torch.vtensor<[?,?,4,?],f32>
+}