Skip to content

Commit

Permalink
[mlir] Add basic support for dynamic tensor results in TensorToBuffer…
Browse files Browse the repository at this point in the history 
…s.cpp.

The simplest case is when the indexing maps are DimIds in every component. This covers cwise ops.

Also:
* Expose populateConvertLinalgOnTensorsToBuffersPatterns in Transforms.h
* Expose emitLoopRanges in Transforms.h

Differential Revision: https://reviews.llvm.org/D88781
  • Loading branch information
@pifon2a
pifon2a committed Oct 8, 2020
1 parent aa47962 commit c1fd430
Show file tree
Hide file tree
Showing 4 changed files with 288 additions and 154 deletions.
19 changes: 19 additions & 0 deletions mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,9 @@
#include "llvm/ADT/SmallBitVector.h"

namespace mlir {

class BufferAssignmentTypeConverter;

namespace linalg {

struct LinalgFusionOptions;
Expand Down Expand Up @@ -45,6 +48,12 @@ void populateConvVectorizationPatterns(
MLIRContext *context, SmallVectorImpl<OwningRewritePatternList> &patterns,
ArrayRef<int64_t> tileSizes);

/// Populates the given list with patterns to convert Linalg operations on
/// tensors to buffers.
void populateConvertLinalgOnTensorsToBuffersPatterns(
MLIRContext *context, BufferAssignmentTypeConverter *converter,
OwningRewritePatternList *patterns);

/// Performs standalone tiling of a single LinalgOp by `tileSizes`.
/// and permute the loop nest according to `interchangeVector`
/// The permutation is expressed as a list of integers that specify
Expand Down Expand Up @@ -246,6 +255,16 @@ Optional<LinalgOp> promoteSubViews(OpBuilder &b, LinalgOp op,
LinalgPromotionOptions options,
OperationFolder *folder = nullptr);

/// Creates a number of ranges equal to the number of dimensions in the `map`.
/// The returned ranges correspond to the loop ranges, in the proper order, for
/// which new loops will be created.
/// The function supports only maps that are invertible and have results of type
/// DimExpr or (DimExpr + DimExpr - SymbolExpr floordiv ConstExpr).
/// 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.
SmallVector<Range, 4> emitLoopRanges(OpBuilder &b, Location loc, AffineMap map,
ValueRange viewSizes);

/// Emit a suitable vector form for a Linalg op with fully static shape.
void vectorizeLinalgOp(OpBuilder &builder, Operation *op);

Expand Down
135 changes: 64 additions & 71 deletions mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,77 +58,6 @@ static SmallVector<Value, 4> permuteIvs(ArrayRef<Value> ivs,
: SmallVector<Value, 4>(ivs.begin(), ivs.end());
}

/// Creates a number of ranges equal to the number of dimensions in the `map`.
/// The returned ranges correspond to the loop ranges, in the proper order, for
/// which new loops will be created.
/// The function supports only maps that are invertible and have results of type
/// DimExpr or (DimExpr + DimExpr - SymbolExpr floordiv ConstExpr).
/// 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<Range, 4> emitLoopRanges(OpBuilder &b, Location loc,
AffineMap map,
ValueRange viewSizes) {
unsigned numDims = map.getNumDims(), numRes = map.getNumResults();
unsigned numSym = map.getNumSymbols();
assert(viewSizes.size() == numRes + numSym &&
"viewSizes must contain sizes of all views and values for symbols");
SmallVector<Range, 4> res(numDims);
for (unsigned idx = 0; idx < numRes; ++idx) {
auto result = map.getResult(idx);
if (auto d = result.dyn_cast<AffineDimExpr>()) {
if (res[d.getPosition()].offset)
continue;
res[d.getPosition()] =
Range{std_constant_index(0), viewSizes[idx], std_constant_index(1)};
}

// If the access pattern is of form (m, n)[s] -> (m + n - s floordiv 2),
// then the bounds are:
// (s floordiv 2) <= m <= (size(m) + s floordiv 2 - s + 1).
// where size(n) is applied to the symbol s.
// This is done statically now.
if (auto binOp = result.dyn_cast<AffineBinaryOpExpr>()) {
auto lhs = binOp.getLHS().dyn_cast<AffineBinaryOpExpr>();
auto rhs = binOp.getRHS().dyn_cast<AffineBinaryOpExpr>();
if (!lhs || !rhs || binOp.getKind() != AffineExprKind::Add ||
lhs.getKind() != AffineExprKind::Add ||
rhs.getKind() != mlir::AffineExprKind::Mul)
continue;

auto m = lhs.getLHS().dyn_cast<AffineDimExpr>();
auto n = lhs.getRHS().dyn_cast<AffineDimExpr>();
auto fDiv = rhs.getLHS().dyn_cast<AffineBinaryOpExpr>();
auto minusOne = rhs.getRHS().dyn_cast<AffineConstantExpr>();
if (!m || !n || !fDiv || !minusOne ||
fDiv.getKind() != AffineExprKind::FloorDiv ||
fDiv.getLHS().getKind() != AffineExprKind::SymbolId ||
fDiv.getRHS().getKind() != AffineExprKind::Constant)
continue;

auto s = fDiv.getLHS().dyn_cast<AffineSymbolExpr>();
if (minusOne.getValue() != -1)
continue;

int mPos = m.getPosition();
AffineExpr one = getAffineConstantExpr(1, s.getContext());
AffineExpr sizeOfM = getAffineSymbolExpr(numSym, s.getContext());
// Construction of upper bound (size(m) + s floordiv 2 - s + 1).
AffineExpr upperOffsetExpr = sizeOfM + fDiv + one - s;
AffineMap fromMap = AffineMap::get(numDims, numSym + 1, fDiv);
AffineMap toMap = AffineMap::get(numDims, numSym + 1, upperOffsetExpr);
SmallVector<Value, 8> values(viewSizes.begin(),
viewSizes.begin() + numDims);
values.insert(values.end(), viewSizes.begin() + numRes, viewSizes.end());
values.push_back(viewSizes[mPos]);
// Construction of the lower bound (s floordiv 2).
Value from = applyMapToValues(b, loc, fromMap, values).front();
Value to = applyMapToValues(b, loc, toMap, values).front();
res[mPos] = Range{from, to, std_constant_index(1)};
}
}
return res;
}

template <typename IndexedValueType, typename OpType>
static void inlineRegionAndEmitStore(OpType op, ArrayRef<Value> indexedValues,
ArrayRef<SmallVector<Value, 8>> indexing,
Expand Down Expand Up @@ -708,6 +637,70 @@ static Optional<LinalgLoops> linalgOpToLoopsImplSwitch(Operation *op,
llvm_unreachable("Unexpected op in linalgOpToLoopsImpl");
}

SmallVector<Range, 4> mlir::linalg::emitLoopRanges(OpBuilder &b, Location loc,
AffineMap map,
ValueRange viewSizes) {
unsigned numDims = map.getNumDims(), numRes = map.getNumResults();
unsigned numSym = map.getNumSymbols();
assert(viewSizes.size() == numRes + numSym &&
"viewSizes must contain sizes of all views and values for symbols");
SmallVector<Range, 4> res(numDims);
for (unsigned idx = 0; idx < numRes; ++idx) {
auto result = map.getResult(idx);
if (auto d = result.dyn_cast<AffineDimExpr>()) {
if (res[d.getPosition()].offset)
continue;
res[d.getPosition()] =
Range{std_constant_index(0), viewSizes[idx], std_constant_index(1)};
}

// If the access pattern is of form (m, n)[s] -> (m + n - s floordiv 2),
// then the bounds are:
// (s floordiv 2) <= m <= (size(m) + s floordiv 2 - s + 1).
// where size(n) is applied to the symbol s.
// This is done statically now.
if (auto binOp = result.dyn_cast<AffineBinaryOpExpr>()) {
auto lhs = binOp.getLHS().dyn_cast<AffineBinaryOpExpr>();
auto rhs = binOp.getRHS().dyn_cast<AffineBinaryOpExpr>();
if (!lhs || !rhs || binOp.getKind() != AffineExprKind::Add ||
lhs.getKind() != AffineExprKind::Add ||
rhs.getKind() != mlir::AffineExprKind::Mul)
continue;

auto m = lhs.getLHS().dyn_cast<AffineDimExpr>();
auto n = lhs.getRHS().dyn_cast<AffineDimExpr>();
auto fDiv = rhs.getLHS().dyn_cast<AffineBinaryOpExpr>();
auto minusOne = rhs.getRHS().dyn_cast<AffineConstantExpr>();
if (!m || !n || !fDiv || !minusOne ||
fDiv.getKind() != AffineExprKind::FloorDiv ||
fDiv.getLHS().getKind() != AffineExprKind::SymbolId ||
fDiv.getRHS().getKind() != AffineExprKind::Constant)
continue;

auto s = fDiv.getLHS().dyn_cast<AffineSymbolExpr>();
if (minusOne.getValue() != -1)
continue;

int mPos = m.getPosition();
AffineExpr one = getAffineConstantExpr(1, s.getContext());
AffineExpr sizeOfM = getAffineSymbolExpr(numSym, s.getContext());
// Construction of upper bound (size(m) + s floordiv 2 - s + 1).
AffineExpr upperOffsetExpr = sizeOfM + fDiv + one - s;
AffineMap fromMap = AffineMap::get(numDims, numSym + 1, fDiv);
AffineMap toMap = AffineMap::get(numDims, numSym + 1, upperOffsetExpr);
SmallVector<Value, 8> values(viewSizes.begin(),
viewSizes.begin() + numDims);
values.insert(values.end(), viewSizes.begin() + numRes, viewSizes.end());
values.push_back(viewSizes[mPos]);
// Construction of the lower bound (s floordiv 2).
Value from = applyMapToValues(b, loc, fromMap, values).front();
Value to = applyMapToValues(b, loc, toMap, values).front();
res[mPos] = Range{from, to, std_constant_index(1)};
}
}
return res;
}

/// Emits a loop nest with the proper body for `op`.
template <typename LoopTy>
Optional<LinalgLoops> mlir::linalg::linalgLowerOpToLoops(OpBuilder &builder,
Expand Down
Loading

0 comments on commit c1fd430

Please sign in to comment.