[Mlir-commits] [mlir] b70dde5 - [mlir][linalg] Fix typo in ExtractSliceOfPadTensorSwapPattern

[Matthias Springer]

Author: Matthias Springer
Date: 2021-07-14T22:30:32+09:00
New Revision: b70dde522d50c4aa5471ed7270fb9191c6b35424

URL: https://github.com/llvm/llvm-project/commit/b70dde522d50c4aa5471ed7270fb9191c6b35424
DIFF: https://github.com/llvm/llvm-project/commit/b70dde522d50c4aa5471ed7270fb9191c6b35424.diff

LOG: [mlir][linalg] Fix typo in ExtractSliceOfPadTensorSwapPattern

Differential Revision: https://reviews.llvm.org/D105607

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
    mlir/test/Dialect/Linalg/subtensor-of-padtensor.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index 82d259ed05722..59a858f57d731 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 d4c8e61d122ad..362de8ef1300d 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>
+}