maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[MLIR] Improve support for 0-dimensional Affine Maps. 

Summary:
Modified AffineMap::get to remove support for the overload which allowed
an ArrayRef of AffineExpr but no context (and gathered the context from a
presumed first entry, resulting in bugs when there were 0 results).

Instead, we support only a ArrayRef and a context, and a version which
takes a single AffineExpr.

Additionally, removed some now needless case logic which previously
special cased which call to AffineMap::get to use.

Reviewers: flaub, bondhugula, rriddle!, nicolasvasilache, ftynse, ulysseB, mravishankar, antiagainst, aartbik

Subscribers: mehdi_amini, jpienaar, burmako, shauheen, antiagainst, arpith-jacob, mgester, lucyrfox, liufengdb, Joonsoo, bader, grosul1, frgossen, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D78226
This commit is contained in:
Jeremy Bruestle 2020-04-15 11:12:47 -07:00
parent 639b8da8dc
commit 9f3ab92ec8
25 changed files with 124 additions and 115 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
View File

@ -188,9 +188,9 @@ def DotOp : LinalgStructured_Op<"dot", [NInputs<2>, NOutputs<1>]> {
MLIRContext *context = getContext();
auto r_i = getAffineDimExpr(0, context);
return SmallVector<AffineMap, 8>{
AffineMap::get(1, 0, {r_i}),
AffineMap::get(1, 0, {r_i}),
AffineMap::get(1, 0, context)};
AffineMap::get(1, 0, {r_i}, context),
AffineMap::get(1, 0, {r_i}, context),
AffineMap::get(1, 0, {}, context)};
}
}];
@ -215,8 +215,10 @@ def MatvecOp : LinalgStructured_Op<"matvec", [NInputs<2>, NOutputs<1>]> {
AffineExpr i, r_j;
bindDims(context, i, r_j);
return SmallVector<AffineMap, 8>{
AffineMap::get(2, 0, {i, r_j}), AffineMap::get(2, 0, {r_j}),
AffineMap::get(2, 0, {i})};
AffineMap::get(2, 0, {i, r_j}, context),
AffineMap::get(2, 0, {r_j}, context),
AffineMap::get(2, 0, {i}, context)
};
}
}];
@ -242,9 +244,11 @@ def MatmulOp : LinalgStructured_Op<"matmul", [NInputs<2>, NOutputs<1>]> {
MLIRContext *context = getContext();
AffineExpr i, j, r_k;
bindDims(context, i, j, r_k);
return SmallVector<AffineMap, 8>{AffineMap::get(3, 0, {i, r_k}),
AffineMap::get(3, 0, {r_k, j}),
AffineMap::get(3, 0, {i, j})};
return SmallVector<AffineMap, 8>{
AffineMap::get(3, 0, {i, r_k}, context),
AffineMap::get(3, 0, {r_k, j},context),
AffineMap::get(3, 0, {i, j}, context)
};
}
}];
@ -403,15 +407,15 @@ def ConvOp : PoolingBase_Op<"conv", [NInputs<2>, NOutputs<1>]> {
auto ws = weightedPoolingInputIndex(*this, xs, zs);
return SmallVector<AffineMap, 8>{
// filter[z[0], ..., z[N-1], q, k]
AffineMap::get(idx, 0, concat(concat(zs, qs), ks)),
AffineMap::get(idx, 0, concat(concat(zs, qs), ks), context),
// input[b,
// x[0]*s[0] + d[0]*z[0] - pad_low[0],
// ...
// x[N-1]*s[N-1] + d[N-1]*z[N-1] - pad_low[N-1],
// q]
AffineMap::get(idx, 0, concat(concat(bs, ws), qs)),
AffineMap::get(idx, 0, concat(concat(bs, ws), qs), context),
// output[b, x[0], ..., x[N-1], k]
AffineMap::get(idx, 0, concat(concat(bs, xs), ks))};
AffineMap::get(idx, 0, concat(concat(bs, xs), ks), context)};
}
}];
@ -465,11 +469,11 @@ class SingleInputPoolingBase_Op<string mnemonic>
weightedPoolingInputIndex(*this, outputDims, windowDims);
return SmallVector<AffineMap, 8>{
// input
AffineMap::get(idx, 0, inputDims),
AffineMap::get(idx, 0, inputDims, context),
// windowDims
AffineMap::get(idx, 0, windowDims),
AffineMap::get(idx, 0, windowDims, context),
// output
AffineMap::get(idx, 0, outputDims)
AffineMap::get(idx, 0, outputDims, context)
};
}
}];

22
mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h
View File

@ -25,22 +25,24 @@
namespace mlir {
inline bool isRowMajorMatmul(ArrayAttr indexingMaps) {
auto context = indexingMaps.getContext();
AffineExpr m, n, k;
bindDims(indexingMaps.getContext(), m, n, k);
auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}));
auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}));
auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {m, n}));
auto maps = ArrayAttr::get({mapA, mapB, mapC}, indexingMaps.getContext());
bindDims(context, m, n, k);
auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}, context));
auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}, context));
auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {m, n}, context));
auto maps = ArrayAttr::get({mapA, mapB, mapC}, context);
return indexingMaps == maps;
}
inline bool isColumnMajorMatmul(ArrayAttr indexingMaps) {
auto context = indexingMaps.getContext();
AffineExpr m, n, k;
bindDims(indexingMaps.getContext(), m, n, k);
auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}));
auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}));
auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {n, m}));
auto maps = ArrayAttr::get({mapA, mapB, mapC}, indexingMaps.getContext());
bindDims(context, m, n, k);
auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}, context));
auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}, context));
auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {n, m}, context));
auto maps = ArrayAttr::get({mapA, mapB, mapC}, context);
return indexingMaps == maps;
}

8
mlir/include/mlir/IR/AffineMap.h
View File

@ -49,13 +49,13 @@ public:
static AffineMap get(unsigned dimCount, unsigned symbolCount,
MLIRContext *context);
/// Returns an affine map with `dimCount` dimensions and `symbolCount` symbols
/// mapping to the given results. The array of results cannot be empty.
/// Returns an affine map with `dimCount` dimensions and `symbolCount` mapping
/// to a single output dimension
static AffineMap get(unsigned dimCount, unsigned symbolCount,
ArrayRef<AffineExpr> results);
AffineExpr result);
/// Returns an affine map with `dimCount` dimensions and `symbolCount` mapping
/// to the given results, where the number of results can be zero.
/// to the given results.
static AffineMap get(unsigned dimCount, unsigned symbolCount,
ArrayRef<AffineExpr> results, MLIRContext *context);

7
mlir/lib/Analysis/AffineStructures.cpp
View File

@ -1464,11 +1464,8 @@ std::pair<AffineMap, AffineMap> FlatAffineConstraints::getLowerAndUpperBound(
lbExprs.push_back(expr);
}
auto lbMap = lbExprs.empty() ? AffineMap()
: AffineMap::get(dimCount, symCount, lbExprs);
auto ubMap = ubExprs.empty() ? AffineMap()
: AffineMap::get(dimCount, symCount, ubExprs);
auto lbMap = AffineMap::get(dimCount, symCount, lbExprs, context);
auto ubMap = AffineMap::get(dimCount, symCount, ubExprs, context);
return {lbMap, ubMap};
}

4
mlir/lib/Analysis/LoopAnalysis.cpp
View File

@ -62,8 +62,8 @@ void mlir::buildTripCountMapAndOperands(
SmallVector<AffineExpr, 4> lbSplatExpr(ubValueMap.getNumResults(),
lbMap.getResult(0));
auto lbMapSplat =
AffineMap::get(lbMap.getNumDims(), lbMap.getNumSymbols(), lbSplatExpr);
auto lbMapSplat = AffineMap::get(lbMap.getNumDims(), lbMap.getNumSymbols(),
lbSplatExpr, b.getContext());
AffineValueMap lbSplatValueMap(lbMapSplat, forOp.getLowerBoundOperands());
AffineValueMap tripCountValueMap;

4
mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp
View File

@ -99,11 +99,11 @@ SingleWorkgroupReduction::matchAsPerformingReduction(
genericOp.indexing_maps().getValue()[1].cast<AffineMapAttr>();
// The indexing map for the input should be `(i) -> (i)`.
if (inputMap.getValue() !=
AffineMap::get(1, 0, {getAffineDimExpr(0, op->getContext())}))
AffineMap::get(1, 0, getAffineDimExpr(0, op->getContext())))
return llvm::None;
// The indexing map for the input should be `(i) -> (0)`.
if (outputMap.getValue() !=
AffineMap::get(1, 0, {getAffineConstantExpr(0, op->getContext())}))
AffineMap::get(1, 0, getAffineConstantExpr(0, op->getContext())))
return llvm::None;
return linalg::RegionMatcher::matchAsScalarBinaryOp(genericOp);

2
mlir/lib/Dialect/Affine/EDSC/Builders.cpp
View File

@ -129,7 +129,7 @@ static ValueHandle createBinaryIndexHandle(
if (v1) {
operands.push_back(v1);
}
auto map = AffineMap::get(numDims, numSymbols, {affCombiner(d0, d1)});
auto map = AffineMap::get(numDims, numSymbols, affCombiner(d0, d1));
// TODO: createOrFold when available.
Operation *op =
makeComposedAffineApply(ScopedContext::getBuilder(),

23
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
View File

@ -522,7 +522,8 @@ AffineApplyNormalizer::AffineApplyNormalizer(AffineMap map,
"Unexpected number of concatenated symbols");
auto numDims = dimValueToPosition.size();
auto numSymbols = concatenatedSymbols.size() - map.getNumSymbols();
auto auxiliaryMap = AffineMap::get(numDims, numSymbols, auxiliaryExprs);
auto auxiliaryMap =
AffineMap::get(numDims, numSymbols, auxiliaryExprs, map.getContext());
LLVM_DEBUG(map.print(dbgs() << "\nCompose map: "));
LLVM_DEBUG(auxiliaryMap.print(dbgs() << "\nWith map: "));
@ -2163,19 +2164,13 @@ LogicalResult AffinePrefetchOp::fold(ArrayRef<Attribute> cstOperands,
void AffineParallelOp::build(Builder *builder, OperationState &result,
ArrayRef<int64_t> ranges) {
// Default initialize empty maps.
auto lbMap = AffineMap::get(builder->getContext());
auto ubMap = AffineMap::get(builder->getContext());
// If there are ranges, set each to [0, N).
if (ranges.size()) {
SmallVector<AffineExpr, 8> lbExprs(ranges.size(),
builder->getAffineConstantExpr(0));
lbMap = AffineMap::get(0, 0, lbExprs);
SmallVector<AffineExpr, 8> ubExprs;
for (int64_t range : ranges)
ubExprs.push_back(builder->getAffineConstantExpr(range));
ubMap = AffineMap::get(0, 0, ubExprs);
}
SmallVector<AffineExpr, 8> lbExprs(ranges.size(),
builder->getAffineConstantExpr(0));
auto lbMap = AffineMap::get(0, 0, lbExprs, builder->getContext());
SmallVector<AffineExpr, 8> ubExprs;
for (int64_t range : ranges)
ubExprs.push_back(builder->getAffineConstantExpr(range));
auto ubMap = AffineMap::get(0, 0, ubExprs, builder->getContext());
build(builder, result, lbMap, {}, ubMap, {});
}

5
mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp
View File

@ -51,8 +51,9 @@ void AffineValueMap::difference(const AffineValueMap &a,
diffExprs.push_back(normalizer.getAffineMap().getResult(i) -
bMap.getResult(i));
auto diffMap = AffineMap::get(normalizer.getNumDims(),
normalizer.getNumSymbols(), diffExprs);
auto diffMap =
AffineMap::get(normalizer.getNumDims(), normalizer.getNumSymbols(),
diffExprs, aMap.getContext());
canonicalizeMapAndOperands(&diffMap, &bOperands);
diffMap = simplifyAffineMap(diffMap);
res->reset(diffMap, bOperands);

5
mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
View File

@ -143,8 +143,9 @@ constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
boundExprs.push_back(dim + tileSizes[i]);
boundExprs.append(origUbMap.getResults().begin(),
origUbMap.getResults().end());
auto ubMap = AffineMap::get(origUbMap.getNumDims() + 1,
origUbMap.getNumSymbols(), boundExprs);
auto ubMap =
AffineMap::get(origUbMap.getNumDims() + 1, origUbMap.getNumSymbols(),
boundExprs, b.getContext());
newLoops[width + i].setUpperBound(/*operands=*/ubOperands, ubMap);
} else {
// No need of the min expression.

6
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
View File

@ -523,8 +523,10 @@ getSymbolLessAffineMaps(ArrayRef<ArrayRef<AffineExpr>> reassociation) {
"Expected symbol-less expressions");
SmallVector<AffineMap, 4> maps;
maps.reserve(reassociation.size());
for (auto exprs : reassociation)
maps.push_back(AffineMap::get(maxDim + 1, 0, exprs));
for (auto exprs : reassociation) {
assert(exprs.size() != 0);
maps.push_back(AffineMap::get(maxDim + 1, 0, exprs, exprs[0].getContext()));
}
return maps;
}

3
mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp
View File

@ -252,7 +252,8 @@ public:
// so having a max op is enough.
auto maxMap = AffineMap::get(/*dimCount=*/1, 0,
{getAffineDimExpr(/*position=*/0, context),
getAffineConstantExpr(0, context)});
getAffineConstantExpr(0, context)},
context);
clampedImIdx.push_back(
affine_max(dim.getType(), maxMap, ValueRange{dim}));
}

3
mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
View File

@ -294,7 +294,8 @@ makeTiledViews(OpBuilder &b, Location loc, LinalgOp linalgOp,
/*dimCount=*/3, /*symbolCount=*/0,
{getAffineDimExpr(/*position=*/0, b.getContext()),
getAffineDimExpr(/*position=*/1, b.getContext()) -
getAffineDimExpr(/*position=*/2, b.getContext())});
getAffineDimExpr(/*position=*/2, b.getContext())},
b.getContext());
auto d = folded_std_dim(folder, view, r);
size = folded_affine_min(folder, b.getIndexType(), minMap,
ValueRange{size, d, offset});

3
mlir/lib/Dialect/LoopOps/Transforms/ParallelLoopTiling.cpp
View File

@ -66,7 +66,8 @@ void mlir::loop::tileParallelLoop(ParallelOp op, ArrayRef<int64_t> tileSizes) {
/*dimCount=*/3, /*symbolCount=*/0,
{getAffineDimExpr(/*position=*/0, b.getContext()),
getAffineDimExpr(/*position=*/1, b.getContext()) -
getAffineDimExpr(/*position=*/2, b.getContext())});
getAffineDimExpr(/*position=*/2, b.getContext())},
b.getContext());
// Create the inner loop with adjusted bounds.
SmallVector<Value, 2> newBounds;

4
mlir/lib/Dialect/Vector/VectorTransforms.cpp
View File

@ -1348,9 +1348,7 @@ private:
auto targetExpr = getAffineDimExpr(idx < index ? idx : idx - 1, ctx);
results.push_back(targetExpr);
}
// The (...) -> () affine map has its own factory method.
return results.empty() ? AffineMap::get(map.getNumDims() - 1, 0, ctx)
: AffineMap::get(map.getNumDims() - 1, 0, results);
return AffineMap::get(map.getNumDims() - 1, 0, results, ctx);
}
// Helper to drop dimension from vector type.

2
mlir/lib/Dialect/Vector/VectorUtils.cpp
View File

@ -176,7 +176,7 @@ static AffineMap makePermutationMap(
"Vectorization prerequisite violated: at most 1 index may be "
"invariant wrt a vectorized loop");
}
return AffineMap::get(indices.size(), 0, perm);
return AffineMap::get(indices.size(), 0, perm, context);
}
/// Implementation detail that walks up the parents and records the ones with

35
mlir/lib/IR/AffineMap.cpp
View File

@ -104,7 +104,7 @@ AffineMap AffineMap::getPermutationMap(ArrayRef<unsigned> permutation,
for (auto index : permutation)
affExprs.push_back(getAffineDimExpr(index, context));
auto m = std::max_element(permutation.begin(), permutation.end());
auto permutationMap = AffineMap::get(*m + 1, 0, affExprs);
auto permutationMap = AffineMap::get(*m + 1, 0, affExprs, context);
assert(permutationMap.isPermutation() && "Invalid permutation vector");
return permutationMap;
}
@ -127,13 +127,16 @@ static void getMaxDimAndSymbol(ArrayRef<AffineExprContainer> exprsList,
template <typename AffineExprContainer>
static SmallVector<AffineMap, 4>
inferFromExprList(ArrayRef<AffineExprContainer> exprsList) {
assert(!exprsList.empty());
assert(!exprsList[0].empty());
auto context = exprsList[0][0].getContext();
int64_t maxDim = -1, maxSym = -1;
getMaxDimAndSymbol(exprsList, maxDim, maxSym);
SmallVector<AffineMap, 4> maps;
maps.reserve(exprsList.size());
for (const auto &exprs : exprsList)
maps.push_back(AffineMap::get(/*dimCount=*/maxDim + 1,
/*symbolCount=*/maxSym + 1, exprs));
/*symbolCount=*/maxSym + 1, exprs, context));
return maps;
}
@ -153,7 +156,7 @@ AffineMap AffineMap::getMultiDimIdentityMap(unsigned numDims,
dimExprs.reserve(numDims);
for (unsigned i = 0; i < numDims; ++i)
dimExprs.push_back(mlir::getAffineDimExpr(i, context));
return get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs);
return get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs, context);
}
MLIRContext *AffineMap::getContext() const { return map->context; }
@ -255,8 +258,7 @@ AffineMap AffineMap::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
results.push_back(
expr.replaceDimsAndSymbols(dimReplacements, symReplacements));
return results.empty() ? get(numResultDims, 0, getContext())
: get(numResultDims, numResultSyms, results);
return get(numResultDims, numResultSyms, results, getContext());
}
AffineMap AffineMap::compose(AffineMap map) {
@ -280,8 +282,7 @@ AffineMap AffineMap::compose(AffineMap map) {
exprs.reserve(getResults().size());
for (auto expr : getResults())
exprs.push_back(expr.compose(newMap));
return exprs.empty() ? AffineMap::get(numDims, 0, map.getContext())
: AffineMap::get(numDims, numSymbols, exprs);
return AffineMap::get(numDims, numSymbols, exprs, map.getContext());
}
bool AffineMap::isProjectedPermutation() {
@ -312,7 +313,7 @@ AffineMap AffineMap::getSubMap(ArrayRef<unsigned> resultPos) {
for (auto idx : resultPos) {
exprs.push_back(getResult(idx));
}
return AffineMap::get(getNumDims(), getNumSymbols(), exprs);
return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());
}
AffineMap mlir::simplifyAffineMap(AffineMap map) {
@ -321,7 +322,8 @@ AffineMap mlir::simplifyAffineMap(AffineMap map) {
exprs.push_back(
simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols()));
}
return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs);
return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs,
map.getContext());
}
AffineMap mlir::removeDuplicateExprs(AffineMap map) {
@ -354,7 +356,7 @@ AffineMap mlir::inversePermutation(AffineMap map) {
seenExprs.push_back(expr);
if (seenExprs.size() != map.getNumInputs())
return AffineMap();
return AffineMap::get(map.getNumResults(), 0, seenExprs);
return AffineMap::get(map.getNumResults(), 0, seenExprs, map.getContext());
}
AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
@ -369,9 +371,8 @@ AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
results.append(m.getResults().begin(), m.getResults().end());
numDims = std::max(m.getNumDims(), numDims);
}
return results.empty() ? AffineMap::get(numDims, /*numSymbols=*/0,
maps.front().getContext())
: AffineMap::get(numDims, /*numSymbols=*/0, results);
return AffineMap::get(numDims, /*numSymbols=*/0, results,
maps.front().getContext());
}
//===----------------------------------------------------------------------===//
@ -380,8 +381,7 @@ AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
MutableAffineMap::MutableAffineMap(AffineMap map)
: numDims(map.getNumDims()), numSymbols(map.getNumSymbols()),
// A map always has at least 1 result by construction
context(map.getResult(0).getContext()) {
context(map.getContext()) {
for (auto result : map.getResults())
results.push_back(result);
}
@ -390,8 +390,7 @@ void MutableAffineMap::reset(AffineMap map) {
results.clear();
numDims = map.getNumDims();
numSymbols = map.getNumSymbols();
// A map always has at least 1 result by construction
context = map.getResult(0).getContext();
context = map.getContext();
for (auto result : map.getResults())
results.push_back(result);
}
@ -416,5 +415,5 @@ void MutableAffineMap::simplify() {
}
AffineMap MutableAffineMap::getAffineMap() const {
return AffineMap::get(numDims, numSymbols, results);
return AffineMap::get(numDims, numSymbols, results, context);
}

15
mlir/lib/IR/Builders.cpp
View File

@ -293,12 +293,11 @@ AffineMap Builder::getEmptyAffineMap() { return AffineMap::get(context); }
AffineMap Builder::getConstantAffineMap(int64_t val) {
return AffineMap::get(/*dimCount=*/0, /*symbolCount=*/0,
{getAffineConstantExpr(val)});
getAffineConstantExpr(val));
}
AffineMap Builder::getDimIdentityMap() {
return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0,
{getAffineDimExpr(0)});
return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, getAffineDimExpr(0));
}
AffineMap Builder::getMultiDimIdentityMap(unsigned rank) {
@ -306,18 +305,19 @@ AffineMap Builder::getMultiDimIdentityMap(unsigned rank) {
dimExprs.reserve(rank);
for (unsigned i = 0; i < rank; ++i)
dimExprs.push_back(getAffineDimExpr(i));
return AffineMap::get(/*dimCount=*/rank, /*symbolCount=*/0, dimExprs);
return AffineMap::get(/*dimCount=*/rank, /*symbolCount=*/0, dimExprs,
context);
}
AffineMap Builder::getSymbolIdentityMap() {
return AffineMap::get(/*dimCount=*/0, /*symbolCount=*/1,
{getAffineSymbolExpr(0)});
getAffineSymbolExpr(0));
}
AffineMap Builder::getSingleDimShiftAffineMap(int64_t shift) {
// expr = d0 + shift.
auto expr = getAffineDimExpr(0) + shift;
return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, {expr});
return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr);
}
AffineMap Builder::getShiftedAffineMap(AffineMap map, int64_t shift) {
@ -325,7 +325,8 @@ AffineMap Builder::getShiftedAffineMap(AffineMap map, int64_t shift) {
shiftedResults.reserve(map.getNumResults());
for (auto resultExpr : map.getResults())
shiftedResults.push_back(resultExpr + shift);
return AffineMap::get(map.getNumDims(), map.getNumSymbols(), shiftedResults);
return AffineMap::get(map.getNumDims(), map.getNumSymbols(), shiftedResults,
context);
}
//===----------------------------------------------------------------------===//

6
mlir/lib/IR/MLIRContext.cpp
View File

@ -717,10 +717,8 @@ AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
}
AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
ArrayRef<AffineExpr> results) {
// The number of results can't be zero.
assert(!results.empty());
return getImpl(dimCount, symbolCount, results, results[0].getContext());
AffineExpr result) {
return getImpl(dimCount, symbolCount, {result}, result.getContext());
}
AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,

2
mlir/lib/IR/StandardTypes.cpp
View File

@ -723,7 +723,7 @@ MemRefType mlir::canonicalizeStridedLayout(MemRefType t) {
simplifyAffineExpr(m.getResult(0), m.getNumDims(), m.getNumSymbols());
if (expr != simplifiedLayoutExpr)
return MemRefType::Builder(t).setAffineMaps({AffineMap::get(
m.getNumDims(), m.getNumSymbols(), {simplifiedLayoutExpr})});
m.getNumDims(), m.getNumSymbols(), simplifiedLayoutExpr)});
return MemRefType::Builder(t).setAffineMaps({});
}

13
mlir/lib/Parser/Parser.cpp
View File

@ -3132,12 +3132,8 @@ AffineParser::parseAffineMapOfSSAIds(AffineMap &map,
/*allowEmptyList=*/true))
return failure();
// Parsed a valid affine map.
if (exprs.empty())
map = AffineMap::get(numDimOperands, dimsAndSymbols.size() - numDimOperands,
getContext());
else
map = AffineMap::get(numDimOperands, dimsAndSymbols.size() - numDimOperands,
exprs);
map = AffineMap::get(numDimOperands, dimsAndSymbols.size() - numDimOperands,
exprs, getContext());
return success();
}
@ -3166,11 +3162,8 @@ AffineMap AffineParser::parseAffineMapRange(unsigned numDims,
if (parseCommaSeparatedListUntil(Token::r_paren, parseElt, true))
return AffineMap();
if (exprs.empty())
return AffineMap::get(numDims, numSymbols, getContext());
// Parsed a valid affine map.
return AffineMap::get(numDims, numSymbols, exprs);
return AffineMap::get(numDims, numSymbols, exprs, getContext());
}
/// Parse an affine constraint.

3
mlir/lib/Transforms/LoopFusion.cpp
View File

@ -942,7 +942,8 @@ static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
}
auto indexRemap = zeroOffsetCount == rank
? AffineMap()
: AffineMap::get(outerIVs.size() + rank, 0, remapExprs);
: AffineMap::get(outerIVs.size() + rank, 0, remapExprs,
forOp.getContext());
// Replace all users of 'oldMemRef' with 'newMemRef'.
LogicalResult res =
replaceAllMemRefUsesWith(oldMemRef, newMemRef, {}, indexRemap,

4
mlir/lib/Transforms/PipelineDataTransfer.cpp
View File

@ -98,8 +98,8 @@ static bool doubleBuffer(Value oldMemRef, AffineForOp forOp) {
// Create 'iv mod 2' value to index the leading dimension.
auto d0 = bInner.getAffineDimExpr(0);
int64_t step = forOp.getStep();
auto modTwoMap = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0,
{d0.floorDiv(step) % 2});
auto modTwoMap =
AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, d0.floorDiv(step) % 2);
auto ivModTwoOp = bInner.create<AffineApplyOp>(forOp.getLoc(), modTwoMap,
forOp.getInductionVar());

18
mlir/lib/Transforms/Utils/LoopUtils.cpp
View File

@ -103,7 +103,8 @@ static void getCleanupLoopLowerBound(AffineForOp forOp, unsigned unrollFactor,
operands.clear();
operands.push_back(lb);
operands.append(bumpValues.begin(), bumpValues.end());
map = AffineMap::get(1 + tripCountMap.getNumResults(), 0, newUbExprs);
map = AffineMap::get(1 + tripCountMap.getNumResults(), 0, newUbExprs,
b.getContext());
// Simplify the map + operands.
fullyComposeAffineMapAndOperands(&map, &operands);
map = simplifyAffineMap(map);
@ -485,7 +486,7 @@ LogicalResult mlir::loopUnrollByFactor(AffineForOp forOp,
if (!forOpIV.use_empty()) {
// iv' = iv + 1/2/3...unrollFactor-1;
auto d0 = builder.getAffineDimExpr(0);
auto bumpMap = AffineMap::get(1, 0, {d0 + i * step});
auto bumpMap = AffineMap::get(1, 0, d0 + i * step);
auto ivUnroll =
builder.create<AffineApplyOp>(forOp.getLoc(), bumpMap, forOpIV);
operandMap.map(forOpIV, ivUnroll);
@ -616,7 +617,7 @@ LogicalResult mlir::loopUnrollJamByFactor(AffineForOp forOp,
if (!forOpIV.use_empty()) {
// iv' = iv + i, i = 1 to unrollJamFactor-1.
auto d0 = builder.getAffineDimExpr(0);
auto bumpMap = AffineMap::get(1, 0, {d0 + i * step});
auto bumpMap = AffineMap::get(1, 0, d0 + i * step);
auto ivUnroll =
builder.create<AffineApplyOp>(forOp.getLoc(), bumpMap, forOpIV);
operandMap.map(forOpIV, ivUnroll);
@ -859,7 +860,8 @@ static void augmentMapAndBounds(OpBuilder &b, Value iv, AffineMap *map,
auto bounds = llvm::to_vector<4>(map->getResults());
bounds.push_back(b.getAffineDimExpr(map->getNumDims()) + offset);
operands->insert(operands->begin() + map->getNumDims(), iv);
*map = AffineMap::get(map->getNumDims() + 1, map->getNumSymbols(), bounds);
*map = AffineMap::get(map->getNumDims() + 1, map->getNumSymbols(), bounds,
b.getContext());
canonicalizeMapAndOperands(map, operands);
}
@ -1514,7 +1516,7 @@ generatePointWiseCopy(Location loc, Value memref, Value fastMemRef,
b = forOp.getBodyBuilder();
auto fastBufOffsetMap =
AffineMap::get(lbOperands.size(), 0, {fastBufOffsets[d]});
AffineMap::get(lbOperands.size(), 0, fastBufOffsets[d]);
auto offset = b.create<AffineApplyOp>(loc, fastBufOffsetMap, lbOperands);
// Construct the subscript for the fast memref being copied into/from:
@ -1529,7 +1531,8 @@ generatePointWiseCopy(Location loc, Value memref, Value fastMemRef,
memIndices.push_back(forOp.getInductionVar());
}
auto fastBufMap = AffineMap::get(2 * rank, /*symbolCount=*/0, fastBufExprs);
auto fastBufMap =
AffineMap::get(2 * rank, /*symbolCount=*/0, fastBufExprs, b.getContext());
fullyComposeAffineMapAndOperands(&fastBufMap, &fastBufMapOperands);
fastBufMap = simplifyAffineMap(fastBufMap);
canonicalizeMapAndOperands(&fastBufMap, &fastBufMapOperands);
@ -1837,7 +1840,8 @@ static LogicalResult generateCopy(
auto dimExpr = b.getAffineDimExpr(regionSymbols.size() + i);
remapExprs.push_back(dimExpr - fastBufOffsets[i]);
}
auto indexRemap = AffineMap::get(regionSymbols.size() + rank, 0, remapExprs);
auto indexRemap = AffineMap::get(regionSymbols.size() + rank, 0, remapExprs,
b.getContext());
// Record the begin since it may be invalidated by memref replacement.
Block::iterator prevOfBegin;

10
mlir/test/Dialect/Affine/simplify-affine-structures.mlir
View File

@ -271,3 +271,13 @@ func @affine.apply(%N : index) {
// CHECK-NEXT: addi
return
}
// -----
// CHECK-DAG: #[[MAP_0D:.*]] = affine_map<() -> ()>
// CHECK-LABEL: func @simplify_zero_dim_map
func @simplify_zero_dim_map(%in : memref<f32>) -> f32 {
%out = affine.load %in[] : memref<f32>
return %out : f32
}