[MLIR] Fix crash in AffineMap::replace for zero result maps #80930
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
bondhugula
requested review from
aartbik,
PeimingLiu,
yinying-lisa-li,
matthias-springer,
hanhanW,
nicolasvasilache and
dcaballe
as code owners
llvmbot
added
mlir:core
|
@llvm/pr-subscribers-mlir-memref @llvm/pr-subscribers-mlir-linalg Author: Uday Bondhugula (bondhugula) ChangesFix obvious bug in AffineMap::replace for the case of zero result maps. Full diff: https://github.com/llvm/llvm-project/pull/80930.diff 15 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 225e4d3194e230..edcfcfd830c443 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 134c5569fbb2f3..929a2a7d396496 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 cd751af5bb2558..cce141253989e5 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 1eb5678b417552..f4f6dadfb37166 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 b63baf330c8645..85fb8a539912f7 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 adb56ab36438bf..c4b13193f4e773 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 0174db45a83db2..47b5fcd4014a04 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 986b5f3e1fb604..5d220c6cdd7e58 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 87a37a7926e9e5..dd3af9d8354123 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 452354413e8833..5be6a628904cdf 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 446eb853d2e92d..0eaf9f71a37d21 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 f1a27168bd4e54..b844c2bfa837ce 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 4034dc40685a4b..53ae138d1e43a0 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 c2804626635947..4aa0d4f34a09fa 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 9b8ee3d4528035..1794b38478a72d 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();
|
|
@llvm/pr-subscribers-mlir Author: Uday Bondhugula (bondhugula) ChangesFix obvious bug in AffineMap::replace for the case of zero result maps. Full diff: https://github.com/llvm/llvm-project/pull/80930.diff 15 Files Affected:
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();
|
|
The other option here is to leave the current behavior of asserting in |
bondhugula
force-pushed
the
uday/fix_affine_map_replace
branch
from
ef13c68 to
7fbe9e7
Compare
Fix obvious bug in AffineMap::replace for the case of zero result maps. Extend/complete inferExprsFromList to work with empty expression lists.
bondhugula
force-pushed
the
uday/fix_affine_map_replace
branch
from
7fbe9e7 to
c0f4690
Compare
joker-eph
reviewed
Feb 7, 2024
joker-eph
approved these changes
Feb 7, 2024
matthias-springer
approved these changes
Feb 7, 2024
nicolasvasilache
previously requested changes
Feb 7, 2024
There was a problem hiding this comment.
I'd supersede the "important arguments" point by the "default arguments".
Do we expect this to support default arguments or not?
If not this is fine, otherwise we need to put the mandatory args first.
Putting a block to get an answer to this (didn't look at the code here).
There are no default arguments here now or later that I can think of. |
bondhugula
dismissed
nicolasvasilache’s stale review
Responded. Change request wasn't a request for changes.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Fix obvious bug in AffineMap::replace for the case of zero result maps.
Extend/complete inferExprsFromList to work with empty expression lists.