From b70dde522d50c4aa5471ed7270fb9191c6b35424 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm@google.com>
Date: Wed, 14 Jul 2021 22:14:05 +0900
Subject: [PATCH] [mlir][linalg] Fix typo in ExtractSliceOfPadTensorSwapPattern

Differential Revision: https://reviews.llvm.org/D105607
---
 mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp  |  2 +-
 .../Dialect/Linalg/subtensor-of-padtensor.mlir     | 24 ++++++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index 82d259e..59a858f 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -954,7 +954,7 @@ LogicalResult ExtractSliceOfPadTensorSwapPattern::matchAndRewrite(
     SmallVector<Value> dynDims;
     for (unsigned i = 0; i < type.getRank(); ++i) {
       if (type.isDynamicDim(i))
-        dynDims.push_back(asValue(rewriter, loc, sliceOp.getMixedOffsets()[i]));
+        dynDims.push_back(asValue(rewriter, loc, sliceOp.getMixedSizes()[i]));
     }
 
     // Create GenerateOp.
diff --git a/mlir/test/Dialect/Linalg/subtensor-of-padtensor.mlir b/mlir/test/Dialect/Linalg/subtensor-of-padtensor.mlir
index d4c8e61..362de8e 100644
--- a/mlir/test/Dialect/Linalg/subtensor-of-padtensor.mlir
+++ b/mlir/test/Dialect/Linalg/subtensor-of-padtensor.mlir
@@ -153,3 +153,27 @@ func @dynamic_high_pad(%arg0 : tensor<?x5xf32>, %h1: index, %pad : f32) -> tenso
   return %1 : tensor<3x4xf32>
 }
 
+// -----
+
+// CHECK-LABEL: @dynamic_extract_size
+//  CHECK-SAME:     %[[ARG0:.*]]: tensor<?x5xf32>, %[[ARG1:.*]]: index
+//   CHECK-NOT:   linalg.pad_tensor
+//       CHECK:   %[[C0:.*]] = constant 0 : index
+//       CHECK:   tensor.dim %[[ARG0]], %[[C0]]
+//       CHECK:   %[[RESULT:.*]] = scf.if %{{.*}} -> (tensor<?x4xf32>) {
+//       CHECK:     %[[GEN:.*]] = tensor.generate %[[ARG1]]
+//       CHECK:     scf.yield %[[GEN]]
+//       CHECK:   } else {
+//       CHECK:     %[[SUBTENSOR:.*]] = tensor.extract_slice %[[ARG0]][%{{.*}}, 4] [%{{.*}}, 1] [1, 1] : tensor<?x5xf32> to tensor<?x1xf32>
+//       CHECK:     %[[PADTENSOR:.*]] = linalg.pad_tensor %[[SUBTENSOR]] low[0, 0] high[%{{.*}}, 3]
+//       CHECK:     scf.yield %[[PADTENSOR]]
+//       CHECK:   }
+//       CHECK:   return %[[RESULT]]
+func @dynamic_extract_size(%arg0 : tensor<?x5xf32>, %s1: index, %pad : f32) -> tensor<?x4xf32> {
+  %0 = linalg.pad_tensor %arg0 low[0, 0] high[7, 8] {
+    ^bb0(%arg1: index, %arg2: index):
+      linalg.yield %pad : f32
+    } : tensor<?x5xf32> to tensor<?x13xf32>
+  %1 = tensor.extract_slice %0[2, 4] [%s1, 4] [1, 1] : tensor<?x13xf32> to tensor<?x4xf32>
+  return %1 : tensor<?x4xf32>
+}
-- 
2.7.4