[Mlir-commits] [mlir] [MLIR] Fix crash in AffineMap::replace for zero result maps (PR #80930)
Uday Bondhugula
llvmlistbot at llvm.org
Tue Feb 6 19:32:50 PST 2024
https://github.com/bondhugula created https://github.com/llvm/llvm-project/pull/80930
Fix obvious bug in AffineMap::replace for the case of zero result maps.
Extend/complete inferExprsFromList to work with empty expression lists.
>From ef13c688c909694e58a8a1459e340fbddcf1d24d Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday at polymagelabs.com>
Date: Tue, 6 Feb 2024 17:37:58 +0530
Subject: [PATCH] [MLIR] Fix crash in AffineMap::replace for zero result maps
Fix obvious bug in AffineMap::replace for the case of zero result maps.
Extend/complete inferExprsFromList to work with empty expression lists.
---
.../mlir/Dialect/Affine/IR/AffineOps.td | 3 ++-
.../mlir/Dialect/Utils/StructuredOpsUtils.h | 4 +++-
mlir/include/mlir/IR/AffineMap.h | 6 +++--
.../Conversion/TosaToLinalg/TosaToLinalg.cpp | 14 +++++++----
.../Conversion/VectorToGPU/VectorToGPU.cpp | 8 +++++--
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 8 +++++--
mlir/lib/Dialect/Linalg/Transforms/Split.cpp | 4 +++-
mlir/lib/Dialect/Linalg/Utils/Utils.cpp | 17 ++++++-------
.../Transforms/SparseGPUCodegen.cpp | 4 +++-
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 7 +++---
.../Vector/Transforms/LowerVectorContract.cpp | 4 +++-
.../VectorTransferSplitRewritePatterns.cpp | 2 +-
.../Vector/Transforms/VectorTransforms.cpp | 9 ++++---
mlir/lib/IR/AffineMap.cpp | 24 ++++++++++++-------
mlir/lib/IR/BuiltinTypes.cpp | 2 +-
15 files changed, 75 insertions(+), 41 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 225e4d3194e23..edcfcfd830c44 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -67,7 +67,8 @@ def AffineApplyOp : Affine_Op<"apply", [Pure]> {
OpBuilder<(ins "ArrayRef<AffineExpr> ":$exprList,"ValueRange":$mapOperands),
[{
build($_builder, $_state, $_builder.getIndexType(),
- AffineMap::inferFromExprList(exprList).front(), mapOperands);
+ AffineMap::inferFromExprList(exprList, $_builder.getContext())
+ .front(), mapOperands);
}]>
];
diff --git a/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h b/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h
index 134c5569fbb2f..929a2a7d39649 100644
--- a/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h
@@ -121,7 +121,9 @@ class StructuredGenerator {
}
bool layout(MapList l) {
- auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
+ auto infer = [&](MapList m) {
+ return AffineMap::inferFromExprList(m, ctx);
+ };
return maps == infer(l);
}
diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h
index cd751af5bb255..cce141253989e 100644
--- a/mlir/include/mlir/IR/AffineMap.h
+++ b/mlir/include/mlir/IR/AffineMap.h
@@ -122,9 +122,11 @@ class AffineMap {
/// `exprs.size()`, as many dims as the largest dim in `exprs` and as many
/// symbols as the largest symbol in `exprs`.
static SmallVector<AffineMap, 4>
- inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList);
+ inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList,
+ MLIRContext *context);
static SmallVector<AffineMap, 4>
- inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList);
+ inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList,
+ MLIRContext *context);
MLIRContext *getContext() const;
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 1eb5678b41755..f4f6dadfb3716 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -2010,7 +2010,8 @@ class ArgMaxConverter : public OpRewritePattern<tosa::ArgMaxOp> {
}
bool didEncounterError = false;
- auto maps = AffineMap::inferFromExprList({srcExprs, dstExprs, dstExprs});
+ auto maps = AffineMap::inferFromExprList({srcExprs, dstExprs, dstExprs},
+ rewriter.getContext());
auto linalgOp = rewriter.create<linalg::GenericOp>(
loc, ArrayRef<Type>({resultTy, resultMaxTy}), input,
ValueRange({filledTensorIdx, filledTensorMax}), maps, iteratorTypes,
@@ -2351,9 +2352,11 @@ struct RFFT2dConverter final : public OpRewritePattern<RFFT2dOp> {
createZeroTensor(rewriter, loc, outputType, dynamicSizes)};
// Indexing maps for input and output tensors
- auto indexingMaps = AffineMap::inferFromExprList(llvm::ArrayRef{
- affineDimsExpr(rewriter, 0, 3, 4), affineDimsExpr(rewriter, 0, 1, 2),
- affineDimsExpr(rewriter, 0, 1, 2)});
+ auto indexingMaps = AffineMap::inferFromExprList(
+ llvm::ArrayRef{affineDimsExpr(rewriter, 0, 3, 4),
+ affineDimsExpr(rewriter, 0, 1, 2),
+ affineDimsExpr(rewriter, 0, 1, 2)},
+ rewriter.getContext());
// Width and height dimensions of the original input.
auto dimH = rewriter.createOrFold<tensor::DimOp>(loc, input, 1);
@@ -2463,7 +2466,8 @@ struct FFT2dConverter final : OpRewritePattern<FFT2dOp> {
ArrayRef{RFFT2dConverter::affineDimsExpr(rewriter, 0, 3, 4),
RFFT2dConverter::affineDimsExpr(rewriter, 0, 3, 4),
RFFT2dConverter::affineDimsExpr(rewriter, 0, 1, 2),
- RFFT2dConverter::affineDimsExpr(rewriter, 0, 1, 2)});
+ RFFT2dConverter::affineDimsExpr(rewriter, 0, 1, 2)},
+ rewriter.getContext());
// Width and height dimensions of the original input.
auto dimH = rewriter.createOrFold<tensor::DimOp>(loc, input_real, 1);
diff --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index b63baf330c864..85fb8a539912f 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -77,7 +77,9 @@ static void getXferIndices(RewriterBase &rewriter, TransferOpType xferOp,
static bool contractSupportsMMAMatrixType(vector::ContractionOp contract,
bool useNvGpu) {
using MapList = ArrayRef<ArrayRef<AffineExpr>>;
- auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
+ auto infer = [&](MapList m) {
+ return AffineMap::inferFromExprList(m, contract.getContext());
+ };
AffineExpr m, n, k;
bindDims(contract.getContext(), m, n, k);
auto iteratorTypes = contract.getIteratorTypes().getValue();
@@ -394,7 +396,9 @@ struct PrepareContractToGPUMMA
// Set up the parallel/reduction structure in right form.
using MapList = ArrayRef<ArrayRef<AffineExpr>>;
- auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
+ auto infer = [&](MapList m) {
+ return AffineMap::inferFromExprList(m, op.getContext());
+ };
AffineExpr m, n, k;
bindDims(rewriter.getContext(), m, n, k);
static constexpr std::array<int64_t, 2> perm = {1, 0};
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index adb56ab36438b..c4b13193f4e77 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -1145,7 +1145,9 @@ AffineApplyOp
mlir::affine::makeComposedAffineApply(OpBuilder &b, Location loc, AffineExpr e,
ArrayRef<OpFoldResult> operands) {
return makeComposedAffineApply(
- b, loc, AffineMap::inferFromExprList(ArrayRef<AffineExpr>{e}).front(),
+ b, loc,
+ AffineMap::inferFromExprList(ArrayRef<AffineExpr>{e}, b.getContext())
+ .front(),
operands);
}
@@ -1220,7 +1222,9 @@ mlir::affine::makeComposedFoldedAffineApply(OpBuilder &b, Location loc,
AffineExpr expr,
ArrayRef<OpFoldResult> operands) {
return makeComposedFoldedAffineApply(
- b, loc, AffineMap::inferFromExprList(ArrayRef<AffineExpr>{expr}).front(),
+ b, loc,
+ AffineMap::inferFromExprList(ArrayRef<AffineExpr>{expr}, b.getContext())
+ .front(),
operands);
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Split.cpp b/mlir/lib/Dialect/Linalg/Transforms/Split.cpp
index 0174db45a83db..47b5fcd4014a0 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Split.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Split.cpp
@@ -83,7 +83,9 @@ linalg::splitOp(RewriterBase &rewriter, TilingInterface op, unsigned dimension,
bindDims(rewriter.getContext(), d0, d1, d2);
OpFoldResult minSplitPoint = affine::makeComposedFoldedAffineMin(
rewriter, op.getLoc(),
- AffineMap::inferFromExprList(ArrayRef<AffineExpr>{d0, d1 + d2}).front(),
+ AffineMap::inferFromExprList(ArrayRef<AffineExpr>{d0, d1 + d2},
+ rewriter.getContext())
+ .front(),
{splitPoint, offsets[dimension], sizes[dimension]});
// Compute the size of the second part. Return early if the second part would
diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
index 986b5f3e1fb60..5d220c6cdd7e5 100644
--- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
@@ -670,7 +670,8 @@ computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile,
<< ": make sure in bound with affine.min\n");
AffineExpr dim0, dim1, dim2;
- bindDims(builder.getContext(), dim0, dim1, dim2);
+ MLIRContext *context = builder.getContext();
+ bindDims(context, dim0, dim1, dim2);
// Get the dimension size for this dimension. We need to first calculate
// the max index and then plus one. This is important because for
@@ -678,12 +679,12 @@ computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile,
// form `(d0 * s0 + d1)`, where `d0`/`d1 is an output/filter window
// dimension and `s0` is stride. Directly use the dimension size of
// output/filer window dimensions will cause incorrect calculation.
- AffineMap minusOneMap =
- AffineMap::inferFromExprList({ArrayRef<AffineExpr>{dim0 - 1}})
- .front();
- AffineMap plusOneMap =
- AffineMap::inferFromExprList({ArrayRef<AffineExpr>{dim0 + 1}})
- .front();
+ AffineMap minusOneMap = AffineMap::inferFromExprList(
+ {ArrayRef<AffineExpr>{dim0 - 1}}, context)
+ .front();
+ AffineMap plusOneMap = AffineMap::inferFromExprList(
+ {ArrayRef<AffineExpr>{dim0 + 1}}, context)
+ .front();
SmallVector<OpFoldResult> maxIndices =
llvm::to_vector(llvm::map_range(ubs, [&](OpFoldResult ub) {
return makeComposedFoldedAffineApply(rewriter, loc, minusOneMap,
@@ -696,7 +697,7 @@ computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile,
// Compute min(dim - offset, size) to avoid out-of-bounds accesses.
AffineMap minMap = AffineMap::inferFromExprList(
- {ArrayRef<AffineExpr>{dim1 - dim2, dim0}})
+ {ArrayRef<AffineExpr>{dim1 - dim2, dim0}}, context)
.front();
size =
makeComposedFoldedAffineMin(rewriter, loc, minMap, {size, d, offset});
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
index 87a37a7926e9e..dd3af9d835412 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
@@ -1263,7 +1263,9 @@ struct LinalgOpRewriter : public OpRewritePattern<linalg::GenericOp> {
SmallVector<AffineMap, 4> maps = op.getIndexingMapsArray();
using MapList = ArrayRef<ArrayRef<AffineExpr>>;
- auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
+ auto infer = [&](MapList m) {
+ return AffineMap::inferFromExprList(m, op.getContext());
+ };
AffineExpr i, j, k;
bindDims(getContext(), i, j, k);
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 452354413e883..5be6a628904cd 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -675,9 +675,10 @@ void vector::ContractionOp::build(OpBuilder &builder, OperationState &result,
ArrayRef<IteratorType> iteratorTypes) {
result.addOperands({lhs, rhs, acc});
result.addTypes(acc.getType());
- result.addAttribute(getIndexingMapsAttrName(result.name),
- builder.getAffineMapArrayAttr(
- AffineMap::inferFromExprList(indexingExprs)));
+ result.addAttribute(
+ getIndexingMapsAttrName(result.name),
+ builder.getAffineMapArrayAttr(
+ AffineMap::inferFromExprList(indexingExprs, builder.getContext())));
result.addAttribute(
getIteratorTypesAttrName(result.name),
builder.getArrayAttr(llvm::to_vector(llvm::map_range(
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index 446eb853d2e92..0eaf9f71a37d2 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -695,7 +695,9 @@ ContractionOpToDotLowering::matchAndRewrite(vector::ContractionOp op,
Value lhs = op.getLhs(), rhs = op.getRhs();
using MapList = ArrayRef<ArrayRef<AffineExpr>>;
- auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
+ auto infer = [&](MapList m) {
+ return AffineMap::inferFromExprList(m, op.getContext());
+ };
AffineExpr m, n, k;
bindDims(rewriter.getContext(), m, n, k);
SmallVector<AffineMap> maps = op.getIndexingMapsArray();
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransferSplitRewritePatterns.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransferSplitRewritePatterns.cpp
index f1a27168bd4e5..b844c2bfa837c 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransferSplitRewritePatterns.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransferSplitRewritePatterns.cpp
@@ -209,7 +209,7 @@ createSubViewIntersection(RewriterBase &b, VectorTransferOpInterface xferOp,
AffineExpr i, j, k;
bindDims(xferOp.getContext(), i, j, k);
SmallVector<AffineMap, 4> maps =
- AffineMap::inferFromExprList(MapList{{i - j, k}});
+ AffineMap::inferFromExprList(MapList{{i - j, k}}, b.getContext());
// affine_min(%dimMemRef - %index, %dimAlloc)
Value affineMin = b.create<affine::AffineMinOp>(
loc, index.getType(), maps[0], ValueRange{dimMemRef, index, dimAlloc});
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 4034dc40685a4..53ae138d1e43a 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -160,8 +160,9 @@ struct MultiReduceToContract
iteratorTypes.push_back(vector::IteratorType::reduction);
}
}
- auto dstMap = AffineMap::get(/*dimCount=*/reductionMask.size(),
- /*symCount=*/0, exprs, reduceOp.getContext());
+ auto dstMap =
+ AffineMap::get(/*dimCount=*/reductionMask.size(),
+ /*symbolCount=*/0, exprs, reduceOp.getContext());
rewriter.replaceOpWithNewOp<mlir::vector::ContractionOp>(
reduceOp, mulOp->getOperand(0), mulOp->getOperand(1), reduceOp.getAcc(),
rewriter.getAffineMapArrayAttr({srcMap, srcMap, dstMap}),
@@ -1399,7 +1400,9 @@ struct CanonicalizeContractMatmulToMMT final
// Set up the parallel/reduction structure in right form.
using MapList = ArrayRef<ArrayRef<AffineExpr>>;
- auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
+ auto infer = [&](MapList m) {
+ return AffineMap::inferFromExprList(m, op.getContext());
+ };
AffineExpr m;
AffineExpr n;
AffineExpr k;
diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index c280462663594..4aa0d4f34a09f 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -272,12 +272,16 @@ AffineMap AffineMap::getMultiDimMapWithTargets(unsigned numDims,
return result;
}
+/// Creates an affine map each for each list of AffineExpr's in `exprsList`
+/// while inferring the right number of dimensional and symbolic inputs needed
+/// based on the maximum dimensional and symbolic identifier appearing in the
+/// expressions.
template <typename AffineExprContainer>
static SmallVector<AffineMap, 4>
-inferFromExprList(ArrayRef<AffineExprContainer> exprsList) {
- assert(!exprsList.empty());
- assert(!exprsList[0].empty());
- auto context = exprsList[0][0].getContext();
+inferFromExprList(ArrayRef<AffineExprContainer> exprsList,
+ MLIRContext *context) {
+ if (exprsList.empty())
+ return {};
int64_t maxDim = -1, maxSym = -1;
getMaxDimAndSymbol(exprsList, maxDim, maxSym);
SmallVector<AffineMap, 4> maps;
@@ -289,13 +293,15 @@ inferFromExprList(ArrayRef<AffineExprContainer> exprsList) {
}
SmallVector<AffineMap, 4>
-AffineMap::inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList) {
- return ::inferFromExprList(exprsList);
+AffineMap::inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList,
+ MLIRContext *context) {
+ return ::inferFromExprList(exprsList, context);
}
SmallVector<AffineMap, 4>
-AffineMap::inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList) {
- return ::inferFromExprList(exprsList);
+AffineMap::inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList,
+ MLIRContext *context) {
+ return ::inferFromExprList(exprsList, context);
}
uint64_t AffineMap::getLargestKnownDivisorOfMapExprs() {
@@ -521,7 +527,7 @@ AffineMap::replace(const DenseMap<AffineExpr, AffineExpr> &map) const {
newResults.reserve(getNumResults());
for (AffineExpr e : getResults())
newResults.push_back(e.replace(map));
- return AffineMap::inferFromExprList(newResults).front();
+ return AffineMap::inferFromExprList(newResults, getContext()).front();
}
AffineMap AffineMap::dropResults(const llvm::SmallBitVector &positions) const {
diff --git a/mlir/lib/IR/BuiltinTypes.cpp b/mlir/lib/IR/BuiltinTypes.cpp
index 9b8ee3d452803..1794b38478a72 100644
--- a/mlir/lib/IR/BuiltinTypes.cpp
+++ b/mlir/lib/IR/BuiltinTypes.cpp
@@ -921,7 +921,7 @@ AffineExpr mlir::makeCanonicalStridedLayoutExpr(ArrayRef<int64_t> sizes,
return getAffineConstantExpr(0, context);
assert(!exprs.empty() && "expected exprs");
- auto maps = AffineMap::inferFromExprList(exprs);
+ auto maps = AffineMap::inferFromExprList(exprs, context);
assert(!maps.empty() && "Expected one non-empty map");
unsigned numDims = maps[0].getNumDims(), nSymbols = maps[0].getNumSymbols();
More information about the Mlir-commits
mailing list