[llvm-branch-commits] [mlir] 2373491 - Revert "[mlir][memref] `memref.subview`: Verify result strides"
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Wed Jan 31 00:32:59 PST 2024
Author: Matthias Springer
Date: 2024-01-31T09:32:54+01:00
New Revision: 23734917da16bdfbadf4d75e657f7ff6deabd618
URL: https://github.com/llvm/llvm-project/commit/23734917da16bdfbadf4d75e657f7ff6deabd618
DIFF: https://github.com/llvm/llvm-project/commit/23734917da16bdfbadf4d75e657f7ff6deabd618.diff
LOG: Revert "[mlir][memref] `memref.subview`: Verify result strides"
Added:
Modified:
mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
mlir/test/Dialect/GPU/decompose-memrefs.mlir
mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir
mlir/test/Dialect/MemRef/invalid.mlir
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_rewrite_sort_coo.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index f43217f6f27ae..8b5765b7f8dba 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -917,7 +917,7 @@ static std::map<int64_t, unsigned> getNumOccurences(ArrayRef<int64_t> vals) {
/// This accounts for cases where there are multiple unit-dims, but only a
/// subset of those are dropped. For MemRefTypes these can be disambiguated
/// using the strides. If a dimension is dropped the stride must be dropped too.
-static FailureOr<llvm::SmallBitVector>
+static std::optional<llvm::SmallBitVector>
computeMemRefRankReductionMask(MemRefType originalType, MemRefType reducedType,
ArrayRef<OpFoldResult> sizes) {
llvm::SmallBitVector unusedDims(originalType.getRank());
@@ -941,7 +941,7 @@ computeMemRefRankReductionMask(MemRefType originalType, MemRefType reducedType,
getStridesAndOffset(originalType, originalStrides, originalOffset)) ||
failed(
getStridesAndOffset(reducedType, candidateStrides, candidateOffset)))
- return failure();
+ return std::nullopt;
// For memrefs, a dimension is truly dropped if its corresponding stride is
// also dropped. This is particularly important when more than one of the dims
@@ -976,22 +976,22 @@ computeMemRefRankReductionMask(MemRefType originalType, MemRefType reducedType,
candidateStridesNumOccurences[originalStride]) {
// This should never happen. Cant have a stride in the reduced rank type
// that wasnt in the original one.
- return failure();
+ return std::nullopt;
}
}
if ((int64_t)unusedDims.count() + reducedType.getRank() !=
originalType.getRank())
- return failure();
+ return std::nullopt;
return unusedDims;
}
llvm::SmallBitVector SubViewOp::getDroppedDims() {
MemRefType sourceType = getSourceType();
MemRefType resultType = getType();
- FailureOr<llvm::SmallBitVector> unusedDims =
+ std::optional<llvm::SmallBitVector> unusedDims =
computeMemRefRankReductionMask(sourceType, resultType, getMixedSizes());
- assert(succeeded(unusedDims) && "unable to find unused dims of subview");
+ assert(unusedDims && "unable to find unused dims of subview");
return *unusedDims;
}
@@ -2745,7 +2745,7 @@ void SubViewOp::build(OpBuilder &b, OperationState &result, Value source,
/// For ViewLikeOpInterface.
Value SubViewOp::getViewSource() { return getSource(); }
-/// Return true if `t1` and `t2` have equal offsets (both dynamic or of same
+/// Return true if t1 and t2 have equal offsets (both dynamic or of same
/// static value).
static bool haveCompatibleOffsets(MemRefType t1, MemRefType t2) {
int64_t t1Offset, t2Offset;
@@ -2755,41 +2755,56 @@ static bool haveCompatibleOffsets(MemRefType t1, MemRefType t2) {
return succeeded(res1) && succeeded(res2) && t1Offset == t2Offset;
}
-/// Return true if `t1` and `t2` have equal strides (both dynamic or of same
-/// static value).
-static bool haveCompatibleStrides(MemRefType t1, MemRefType t2) {
- int64_t t1Offset, t2Offset;
- SmallVector<int64_t> t1Strides, t2Strides;
- auto res1 = getStridesAndOffset(t1, t1Strides, t1Offset);
- auto res2 = getStridesAndOffset(t2, t2Strides, t2Offset);
- if (failed(res1) || failed(res2))
- return false;
- for (auto [s1, s2] : llvm::zip_equal(t1Strides, t2Strides))
- if (s1 != s2)
- return false;
- return true;
+/// Checks if `original` Type type can be rank reduced to `reduced` type.
+/// This function is slight variant of `is subsequence` algorithm where
+/// not matching dimension must be 1.
+static SliceVerificationResult
+isRankReducedMemRefType(MemRefType originalType,
+ MemRefType candidateRankReducedType,
+ ArrayRef<OpFoldResult> sizes) {
+ auto partialRes = isRankReducedType(originalType, candidateRankReducedType);
+ if (partialRes != SliceVerificationResult::Success)
+ return partialRes;
+
+ auto optionalUnusedDimsMask = computeMemRefRankReductionMask(
+ originalType, candidateRankReducedType, sizes);
+
+ // Sizes cannot be matched in case empty vector is returned.
+ if (!optionalUnusedDimsMask)
+ return SliceVerificationResult::LayoutMismatch;
+
+ if (originalType.getMemorySpace() !=
+ candidateRankReducedType.getMemorySpace())
+ return SliceVerificationResult::MemSpaceMismatch;
+
+ // No amount of stride dropping can reconcile incompatible offsets.
+ if (!haveCompatibleOffsets(originalType, candidateRankReducedType))
+ return SliceVerificationResult::LayoutMismatch;
+
+ return SliceVerificationResult::Success;
}
+template <typename OpTy>
static LogicalResult produceSubViewErrorMsg(SliceVerificationResult result,
- Operation *op, Type expectedType) {
+ OpTy op, Type expectedType) {
auto memrefType = llvm::cast<ShapedType>(expectedType);
switch (result) {
case SliceVerificationResult::Success:
return success();
case SliceVerificationResult::RankTooLarge:
- return op->emitError("expected result rank to be smaller or equal to ")
+ return op.emitError("expected result rank to be smaller or equal to ")
<< "the source rank. ";
case SliceVerificationResult::SizeMismatch:
- return op->emitError("expected result type to be ")
+ return op.emitError("expected result type to be ")
<< expectedType
<< " or a rank-reduced version. (mismatch of result sizes) ";
case SliceVerificationResult::ElemTypeMismatch:
- return op->emitError("expected result element type to be ")
+ return op.emitError("expected result element type to be ")
<< memrefType.getElementType();
case SliceVerificationResult::MemSpaceMismatch:
- return op->emitError("expected result and source memory spaces to match.");
+ return op.emitError("expected result and source memory spaces to match.");
case SliceVerificationResult::LayoutMismatch:
- return op->emitError("expected result type to be ")
+ return op.emitError("expected result type to be ")
<< expectedType
<< " or a rank-reduced version. (mismatch of result layout) ";
}
@@ -2811,46 +2826,13 @@ LogicalResult SubViewOp::verify() {
if (!isStrided(baseType))
return emitError("base type ") << baseType << " is not strided";
- // Compute the expected result type, assuming that there are no rank
- // reductions.
- auto expectedType = cast<MemRefType>(SubViewOp::inferResultType(
- baseType, getStaticOffsets(), getStaticSizes(), getStaticStrides()));
-
- // Verify all properties of a shaped type: rank, element type and dimension
- // sizes. This takes into account potential rank reductions.
- auto shapedTypeVerification = isRankReducedType(
- /*originalType=*/expectedType, /*candidateReducedType=*/subViewType);
- if (shapedTypeVerification != SliceVerificationResult::Success)
- return produceSubViewErrorMsg(shapedTypeVerification, *this, expectedType);
-
- // Make sure that the memory space did not change.
- if (expectedType.getMemorySpace() != subViewType.getMemorySpace())
- return produceSubViewErrorMsg(SliceVerificationResult::MemSpaceMismatch,
- *this, expectedType);
-
- // Verify the offset of the layout map.
- if (!haveCompatibleOffsets(expectedType, subViewType))
- return produceSubViewErrorMsg(SliceVerificationResult::LayoutMismatch,
- *this, expectedType);
-
- // The only thing that's left to verify now are the strides. First, compute
- // the unused dimensions due to rank reductions. We have to look at sizes and
- // strides to decide which dimensions were dropped. This function also
- // partially verifies strides in case of rank reductions.
- auto unusedDims = computeMemRefRankReductionMask(expectedType, subViewType,
- getMixedSizes());
- if (failed(unusedDims))
- return produceSubViewErrorMsg(SliceVerificationResult::LayoutMismatch,
- *this, expectedType);
-
- // Strides must match if there are no rank reductions.
- // TODO: Verify strides when there are rank reductions. Strides are partially
- // checked in `computeMemRefRankReductionMask`.
- if (unusedDims->none() && !haveCompatibleStrides(expectedType, subViewType))
- return produceSubViewErrorMsg(SliceVerificationResult::LayoutMismatch,
- *this, expectedType);
+ // Verify result type against inferred type.
+ auto expectedType = SubViewOp::inferResultType(
+ baseType, getStaticOffsets(), getStaticSizes(), getStaticStrides());
- return success();
+ auto result = isRankReducedMemRefType(llvm::cast<MemRefType>(expectedType),
+ subViewType, getMixedSizes());
+ return produceSubViewErrorMsg(result, *this, expectedType);
}
raw_ostream &mlir::operator<<(raw_ostream &os, const Range &range) {
@@ -2900,9 +2882,11 @@ static MemRefType getCanonicalSubViewResultType(
ArrayRef<OpFoldResult> mixedSizes, ArrayRef<OpFoldResult> mixedStrides) {
auto nonRankReducedType = llvm::cast<MemRefType>(SubViewOp::inferResultType(
sourceType, mixedOffsets, mixedSizes, mixedStrides));
- FailureOr<llvm::SmallBitVector> unusedDims = computeMemRefRankReductionMask(
- currentSourceType, currentResultType, mixedSizes);
- if (failed(unusedDims))
+ std::optional<llvm::SmallBitVector> unusedDims =
+ computeMemRefRankReductionMask(currentSourceType, currentResultType,
+ mixedSizes);
+ // Return nullptr as failure mode.
+ if (!unusedDims)
return nullptr;
auto layout = llvm::cast<StridedLayoutAttr>(nonRankReducedType.getLayout());
diff --git a/mlir/test/Dialect/GPU/decompose-memrefs.mlir b/mlir/test/Dialect/GPU/decompose-memrefs.mlir
index 56fc9a66b7ace..d714010d0f254 100644
--- a/mlir/test/Dialect/GPU/decompose-memrefs.mlir
+++ b/mlir/test/Dialect/GPU/decompose-memrefs.mlir
@@ -119,7 +119,7 @@ func.func @decompose_subview(%arg0 : memref<?x?x?xf32>) {
// CHECK: %[[IDX1:.*]] = affine.apply #[[MAP1]]()[%[[STRIDES]]#1]
// CHECK: %[[IDX2:.*]] = affine.apply #[[MAP2]]()[%[[TX]], %[[STRIDES]]#0, %[[TY]], %[[STRIDES]]#1, %[[TZ]]]
// CHECK: %[[PTR:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[IDX2]]], sizes: [%{{.*}}, %{{.*}}, %{{.*}}], strides: [%[[IDX]], %[[IDX1]], 4]
-// CHECK: "test.test"(%[[PTR]]) : (memref<?x?x?xf32, strided<[?, ?, 4], offset: ?>>) -> ()
+// CHECK: "test.test"(%[[PTR]]) : (memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>) -> ()
func.func @decompose_subview_strided(%arg0 : memref<?x?x?xf32>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
@@ -129,8 +129,8 @@ func.func @decompose_subview_strided(%arg0 : memref<?x?x?xf32>) {
%block_dim2 = memref.dim %arg0, %c2 : memref<?x?x?xf32>
gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
threads(%tx, %ty, %tz) in (%block_x = %block_dim0, %block_y = %block_dim1, %block_z = %block_dim2) {
- %res = memref.subview %arg0[%tx, %ty, %tz] [%c2, %c2, %c2] [2, 3, 4] : memref<?x?x?xf32> to memref<?x?x?xf32, strided<[?, ?, 4], offset: ?>>
- "test.test"(%res) : (memref<?x?x?xf32, strided<[?, ?, 4], offset: ?>>) -> ()
+ %res = memref.subview %arg0[%tx, %ty, %tz] [%c2, %c2, %c2] [2, 3, 4] : memref<?x?x?xf32> to memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
+ "test.test"(%res) : (memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>) -> ()
gpu.terminator
}
return
diff --git a/mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir b/mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir
index 3407bdbc7c8f9..96b72e042b9e0 100644
--- a/mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir
+++ b/mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir
@@ -595,9 +595,9 @@ func.func @subview_of_subview(%m: memref<1x1024xf32, 3>, %pos: index)
{
%0 = memref.subview %m[3, %pos] [1, 2] [1, 1]
: memref<1x1024xf32, 3>
- to memref<1x2xf32, strided<[1024, 1], offset: ?>, 3>
+ to memref<1x2xf32, strided<[1024, 2], offset: ?>, 3>
%1 = memref.subview %0[1, 2] [1, 1] [1, 1]
- : memref<1x2xf32, strided<[1024, 1], offset: ?>, 3>
+ : memref<1x2xf32, strided<[1024, 2], offset: ?>, 3>
to memref<f32, strided<[], offset: ?>, 3>
return %1 : memref<f32, strided<[], offset: ?>, 3>
}
@@ -675,9 +675,9 @@ func.func @fold_gpu_subgroup_mma_store_matrix_1d(%dst: memref<?xvector<4xf32>>,
// CHECK-LABEL: func.func @fold_gpu_subgroup_mma_load_matrix_2d
// CHECK-SAME: %[[SRC:.+]]: memref<128x128xf32>
func.func @fold_gpu_subgroup_mma_load_matrix_2d(%arg0 : memref<128x128xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index) -> !gpu.mma_matrix<16x16xf16, "COp"> {
- %subview = memref.subview %arg0[%arg1, %arg2][64, 32][2, 1] : memref<128x128xf32> to memref<64x32xf32, strided<[256, 1], offset: ?>>
+ %subview = memref.subview %arg0[%arg1, %arg2][64, 32][2, 1] : memref<128x128xf32> to memref<64x32xf32, strided<[64, 1], offset: ?>>
// CHECK: gpu.subgroup_mma_load_matrix %[[SRC]][{{.+}}] {leadDimension = 32 : index} : memref<128x128xf32> -> !gpu.mma_matrix<16x16xf16, "COp">
- %matrix = gpu.subgroup_mma_load_matrix %subview[%arg3, %arg4] {leadDimension = 32 : index} : memref<64x32xf32, strided<[256, 1], offset: ?>> -> !gpu.mma_matrix<16x16xf16, "COp">
+ %matrix = gpu.subgroup_mma_load_matrix %subview[%arg3, %arg4] {leadDimension = 32 : index} : memref<64x32xf32, strided<[64, 1], offset: ?>> -> !gpu.mma_matrix<16x16xf16, "COp">
return %matrix : !gpu.mma_matrix<16x16xf16, "COp">
}
@@ -686,9 +686,9 @@ func.func @fold_gpu_subgroup_mma_load_matrix_2d(%arg0 : memref<128x128xf32>, %ar
// CHECK-LABEL: func.func @fold_gpu_subgroup_mma_load_matrix_2d
// CHECK-SAME: %[[DST:.+]]: memref<128x128xf32>
func.func @fold_gpu_subgroup_mma_load_matrix_2d(%arg0 : memref<128x128xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %matrix: !gpu.mma_matrix<16x16xf16, "COp">) {
- %subview = memref.subview %arg0[%arg1, %arg2][64, 32][2, 1] : memref<128x128xf32> to memref<64x32xf32, strided<[256, 1], offset: ?>>
+ %subview = memref.subview %arg0[%arg1, %arg2][64, 32][2, 1] : memref<128x128xf32> to memref<64x32xf32, strided<[64, 1], offset: ?>>
// CHECK: gpu.subgroup_mma_store_matrix %{{.+}}, %[[DST]][{{.+}}] {leadDimension = 32 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<128x128xf32>
- gpu.subgroup_mma_store_matrix %matrix, %subview[%arg3, %arg4] {leadDimension = 32 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<64x32xf32, strided<[256, 1], offset: ?>>
+ gpu.subgroup_mma_store_matrix %matrix, %subview[%arg3, %arg4] {leadDimension = 32 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<64x32xf32, strided<[64, 1], offset: ?>>
return
}
diff --git a/mlir/test/Dialect/MemRef/invalid.mlir b/mlir/test/Dialect/MemRef/invalid.mlir
index be60a3dcb1b20..7bb7a2affcbd1 100644
--- a/mlir/test/Dialect/MemRef/invalid.mlir
+++ b/mlir/test/Dialect/MemRef/invalid.mlir
@@ -1073,12 +1073,3 @@ func.func @dim_0_ranked(%arg : memref<f32>, %arg1 : index) {
memref.dim %arg, %arg1 : memref<f32> // expected-error {{'memref.dim' op operand #0 must be unranked.memref of any type values or non-0-ranked.memref of any type values, but got 'memref<f32>'}}
return
}
-
-// -----
-
-func.func @subview_invalid_strides(%m: memref<7x22x333x4444xi32>) {
- // expected-error @below{{expected result type to be 'memref<7x11x333x4444xi32, strided<[32556744, 2959704, 4444, 1]>>' or a rank-reduced version. (mismatch of result layout)}}
- %subview = memref.subview %m[0, 0, 0, 0] [7, 11, 333, 4444] [1, 2, 1, 1]
- : memref<7x22x333x4444xi32> to memref<7x11x333x4444xi32>
- return
-}
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_rewrite_sort_coo.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_rewrite_sort_coo.mlir
index e7dd0ad32a243..3773cca9c8d69 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_rewrite_sort_coo.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_rewrite_sort_coo.mlir
@@ -88,10 +88,10 @@ module {
// Prepare a buffer for x0, x1, x2, y0 and a buffer for y1.
%xys = memref.alloc() : memref<20xi32>
%xy = memref.cast %xys : memref<20xi32> to memref<?xi32>
- %x0 = memref.subview %xy[%i0][%i5][4] : memref<?xi32> to memref<?xi32, strided<[4], offset: ?>>
- %x1 = memref.subview %xy[%i1][%i5][4] : memref<?xi32> to memref<?xi32, strided<[4], offset: ?>>
- %x2 = memref.subview %xy[%i2][%i5][4] : memref<?xi32> to memref<?xi32, strided<[4], offset: ?>>
- %y0 = memref.subview %xy[%i3][%i5][4] : memref<?xi32> to memref<?xi32, strided<[4], offset: ?>>
+ %x0 = memref.subview %xy[%i0][%i5][%i4] : memref<?xi32> to memref<?xi32, strided<[4], offset: ?>>
+ %x1 = memref.subview %xy[%i1][%i5][%i4] : memref<?xi32> to memref<?xi32, strided<[4], offset: ?>>
+ %x2 = memref.subview %xy[%i2][%i5][%i4] : memref<?xi32> to memref<?xi32, strided<[4], offset: ?>>
+ %y0 = memref.subview %xy[%i3][%i5][%i4] : memref<?xi32> to memref<?xi32, strided<[4], offset: ?>>
%y1s = memref.alloc() : memref<7xi32>
%y1 = memref.cast %y1s : memref<7xi32> to memref<?xi32>
More information about the llvm-branch-commits
mailing list