[Mlir-commits] [mlir] 813fa79 - Don't drop in_bounds when vector-transfer-collapse-inner-most-dims

[Ahmed Taei]

Author: Ahmed Taei
Date: 2021-10-29T09:07:57-07:00
New Revision: 813fa79c151cdc8543622f38d287fa4e08a0a99f

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

LOG: Don't drop in_bounds when vector-transfer-collapse-inner-most-dims

When operand is a subview we don't infer in_bounds and some default cases (e.g case in the tests) will crash with `operand is NULL` when converting to LLVM

Reviewed By: ThomasRaoux

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

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/VectorTransforms.cpp
    mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/VectorTransforms.cpp
index cabeb58971d05..c6f29e0a641a2 100644
--- a/mlir/lib/Dialect/Vector/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/VectorTransforms.cpp
@@ -3573,15 +3573,23 @@ class DropInnerMostUnitDims : public OpRewritePattern<vector::TransferReadOp> {
     auto loc = readOp.getLoc();
     SmallVector<int64_t> offsets(srcType.getRank(), 0);
     SmallVector<int64_t> strides(srcType.getRank(), 1);
+
+    ArrayAttr inBounds =
+        readOp.in_bounds()
+            ? rewriter.getArrayAttr(
+                  readOp.in_boundsAttr().getValue().drop_back(dimsToDrop))
+            : ArrayAttr();
     Value rankedReducedView = rewriter.create<memref::SubViewOp>(
         loc, resultMemrefType, readOp.source(), offsets, srcType.getShape(),
         strides);
+    auto permMap = getTransferMinorIdentityMap(
+        rankedReducedView.getType().cast<ShapedType>(), resultTargetVecType);
     Value result = rewriter.create<vector::TransferReadOp>(
         loc, resultTargetVecType, rankedReducedView,
-        readOp.indices().drop_back(dimsToDrop));
+        readOp.indices().drop_back(dimsToDrop), permMap, readOp.padding(),
+        inBounds);
     rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(readOp, targetType,
                                                      result);
-
     return success();
   }
 };

diff  --git a/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir b/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir
index 6ebfeebd81314..9fa5468f50acc 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir
@@ -31,3 +31,19 @@ func @contiguous_inner_most_dim(%A: memref<16x1xf32>, %i:index, %j:index) -> (ve
 //      CHECK:   %[[V:.+]] = vector.transfer_read %[[SRC_0]]
 //      CHECK:   %[[RESULT]] = vector.shape_cast %[[V]] : vector<8xf32> to vector<8x1xf32>
 //      CHECK:   return %[[RESULT]]
+
+// -----
+
+func @contiguous_inner_most_dim_bounds(%A: memref<1000x1xf32>, %i:index, %ii:index) -> (vector<4x1xf32>) {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.0 : f32
+  %0 = memref.subview %A[%i, 0] [40, 1] [1, 1] : memref<1000x1xf32> to memref<40x1xf32, affine_map<(d0, d1)[s0] -> (d0 + s0 + d1)>>
+  %1 = vector.transfer_read %0[%ii, %c0], %cst {in_bounds = [true, true]} : memref<40x1xf32, affine_map<(d0, d1)[s0] -> (d0 + s0 + d1)>>, vector<4x1xf32>
+  return %1 : vector<4x1xf32>
+}
+//      CHECK: func @contiguous_inner_most_dim_bounds(%[[SRC:.+]]: memref<1000x1xf32>, %[[II:.+]]: index, %[[J:.+]]: index) -> vector<4x1xf32>
+//      CHECK:   %[[SRC_0:.+]] = memref.subview %[[SRC]]
+//      CHECK:   %[[SRC_1:.+]] = memref.subview %[[SRC_0]]
+//      CHECK:   %[[V:.+]] = vector.transfer_read %[[SRC_1]]
+// CHECK-SAME:       {in_bounds = [true]}
+// CHECK-SAME:       vector<4xf32>