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

[mlir][Linalg] emitLoopRanges and emitLoopRangesWithSymbols merged into one 

Right now there is a branching for 2 functions based on whether target map has
symbols or not. In this commit these functions are merged into one.
Furthermore, emitting does not require inverse and map applying as it computes
the correct Range in a single step and thus reduces unnecessary overhead.

Differential Revision: https://reviews.llvm.org/D83756
This commit is contained in:
Jakub Lichman 2020-07-23 12:27:46 +02:00 committed by Alex Zinenko
parent 919922b0c2
commit 20c3386f4a
1 changed files with 12 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
View File

@ -58,30 +58,16 @@ static SmallVector<Value, 4> permuteIvs(ArrayRef<Value> ivs,
: SmallVector<Value, 4>(ivs.begin(), ivs.end()); : SmallVector<Value, 4>(ivs.begin(), ivs.end());
} }
// Creates a number of ranges equal to the number of results in `map`.
// The returned ranges correspond to the loop ranges, in the proper order, for
// which new loops will be created.
static SmallVector<SubViewOp::Range, 4>
emitLoopRanges(OpBuilder &b, Location loc, AffineMap map,
ArrayRef<Value> allViewSizes) {
// Apply `map` to get view sizes in loop order.
auto sizes = applyMapToValues(b, loc, map, allViewSizes);
// Create a new range with the applied tile sizes.
ScopedContext scope(b, loc);
SmallVector<SubViewOp::Range, 4> res;
for (unsigned idx = 0, e = map.getNumResults(); idx < e; ++idx) {
res.push_back(SubViewOp::Range{std_constant_index(0), sizes[idx],
std_constant_index(1)});
}
return res;
}
/// Creates a number of ranges equal to the number of dimensions in the `map`. /// Creates a number of ranges equal to the number of dimensions in the `map`.
/// The function supports for now only limited number of expressions inside /// The returned ranges correspond to the loop ranges, in the proper order, for
/// map results. It expects a non-inverted, concatenated map and last values in /// which new loops will be created.
/// viewSizes will be applied to the symbols in the map. /// The function supports only maps that are invertible and have results of type
static SmallVector<SubViewOp::Range, 4> /// DimExpr or (DimExpr + DimExpr - SymbolExpr floordiv ConstExpr).
emitLoopRangesWithSymbols(OpBuilder &b, Location loc, AffineMap map, /// It expects a non-inverted, concatenated map and last values in
/// allViewSizes will be applied to the symbols in the map if it contains any.
static SmallVector<SubViewOp::Range, 4> emitLoopRanges(OpBuilder &b,
Location loc,
AffineMap map,
ValueRange viewSizes) { ValueRange viewSizes) {
unsigned numDims = map.getNumDims(), numRes = map.getNumResults(); unsigned numDims = map.getNumDims(), numRes = map.getNumResults();
unsigned numSym = map.getNumSymbols(); unsigned numSym = map.getNumSymbols();
@ -537,23 +523,8 @@ Optional<LinalgLoops> linalgOpToLoopsImpl(Operation *op, OpBuilder &builder) {
llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); })); llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
SmallVector<Value, 8> sizes = getViewSizes(builder, linalgOp); SmallVector<Value, 8> sizes = getViewSizes(builder, linalgOp);
AffineMap map = concatAffineMaps(maps); AffineMap map = concatAffineMaps(maps);
SmallVector<SubViewOp::Range, 4> loopRanges; auto loopRanges = emitLoopRanges(scope.getBuilderRef(), scope.getLocation(),
map, getViewSizes(builder, linalgOp));
if (map.getNumSymbols()) {
loopRanges = emitLoopRangesWithSymbols(scope.getBuilderRef(),
scope.getLocation(), map, sizes);
} else {
AffineMap invertedMap = inversePermutation(map);
if (!invertedMap)
return {};
if (invertedMap.isEmpty()) {
emitScalarImplementation<IndexedValueTy>({}, linalgOp);
return LinalgLoops();
}
loopRanges = emitLoopRanges(scope.getBuilderRef(), scope.getLocation(),
invertedMap, sizes);
}
SmallVector<Value, 4> allIvs; SmallVector<Value, 4> allIvs;
GenerateLoopNest<LoopTy>::doit( GenerateLoopNest<LoopTy>::doit(
loopRanges, linalgOp.iterator_types().getValue(), [&](ValueRange ivs) { loopRanges, linalgOp.iterator_types().getValue(), [&](ValueRange ivs) {