0a8e3dd432
Rename interface functions as follows: * `hasTensorSemantics` -> `hasPureTensorSemantics` * `hasBufferSemantics` -> `hasPureBufferSemantics` These two functions return "true" if the op has tensor/buffer operands but not buffer/tensor operands. Also drop the "ranked" part from the interface, i.e., do not distinguish between ranked/unranked types. The new function names describe the functions more accurately. They also align their semantics with the notion of "tensor semantics" with the bufferization framework. (An op is supposed to be bufferized if it has tensor operands, and we don't care if it also has memref operands.) This change is in preparation of #75273, which adds `BufferizableOpInterface::hasTensorSemantics`. By renaming the functions in the `DestinationStyleOpInterface`, we can avoid name clashes between the two interfaces.
1388 lines
57 KiB
C++
1388 lines
57 KiB
C++
//===- Sparsification.cpp - Implementation of sparsification --------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file implements converting sparse tensor types to actual sparse code.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "Utils/CodegenEnv.h"
|
|
#include "Utils/CodegenUtils.h"
|
|
#include "Utils/LoopEmitter.h"
|
|
|
|
#include "mlir/Dialect/Affine/IR/AffineOps.h"
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
|
|
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
|
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
|
|
#include "mlir/Dialect/Linalg/IR/Linalg.h"
|
|
#include "mlir/Dialect/Linalg/Utils/Utils.h"
|
|
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
|
#include "mlir/Dialect/SCF/IR/SCF.h"
|
|
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
|
|
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
|
|
#include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
|
|
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
|
|
#include "mlir/Dialect/SparseTensor/Utils/Merger.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/IR/AffineExprVisitor.h"
|
|
#include "mlir/IR/Matchers.h"
|
|
#include "mlir/IR/TensorEncoding.h"
|
|
#include "llvm/ADT/SmallBitVector.h"
|
|
|
|
#include <optional>
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::sparse_tensor;
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Sparsifier analysis methods.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
/// Returns true iff affine expression is invariant. Sets the
|
|
/// parameter `isCurrentLoop` when expression just became invariant.
|
|
static bool isInvariantAffine(AffineExpr a, LoopId curr, bool &isCurrentLoop) {
|
|
switch (a.getKind()) {
|
|
case AffineExprKind::DimId: {
|
|
const LoopId i = cast<AffineDimExpr>(a).getPosition();
|
|
if (i + 1 == curr) {
|
|
isCurrentLoop = true;
|
|
return true; // becomes invariant at current loop
|
|
}
|
|
return i < curr; // invariant when already generated
|
|
}
|
|
case AffineExprKind::Add:
|
|
case AffineExprKind::Mul: {
|
|
auto binOp = cast<AffineBinaryOpExpr>(a);
|
|
return isInvariantAffine(binOp.getLHS(), curr, isCurrentLoop) &&
|
|
isInvariantAffine(binOp.getRHS(), curr, isCurrentLoop);
|
|
}
|
|
default: {
|
|
assert(isa<AffineConstantExpr>(a));
|
|
return true;
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Helper method to inspect affine expressions. Rejects cases where the
|
|
/// same index is used more than once. Also rejects compound affine
|
|
/// expressions in sparse dimensions.
|
|
static bool findAffine(Merger &merger, TensorId tid, Level lvl, AffineExpr a,
|
|
LevelType lt, bool setLvlFormat = true) {
|
|
switch (a.getKind()) {
|
|
case AffineExprKind::DimId: {
|
|
const LoopId idx = merger.makeLoopId(cast<AffineDimExpr>(a).getPosition());
|
|
if (!isUndefLT(merger.getLvlType(tid, idx)))
|
|
return false; // used more than once
|
|
if (setLvlFormat)
|
|
merger.setLevelAndType(tid, idx, lvl, lt);
|
|
return true;
|
|
}
|
|
case AffineExprKind::Add:
|
|
case AffineExprKind::Mul:
|
|
case AffineExprKind::Constant: {
|
|
assert(isDenseLT(lt));
|
|
if (auto binOp = dyn_cast<AffineBinaryOpExpr>(a)) {
|
|
// We do not set dim level format for affine expression like d0 + d1 on
|
|
// either loop index at d0 or d1. We continue the recursion merely to
|
|
// check whether current affine is admissible or not.
|
|
return findAffine(merger, tid, lvl, binOp.getLHS(), lt, false) &&
|
|
findAffine(merger, tid, lvl, binOp.getRHS(), lt, false);
|
|
}
|
|
// Falls through when it is a constant Affine
|
|
return true;
|
|
}
|
|
default:
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/// Helper method to inspect affine expressions for index variable reduction
|
|
/// based codegen. It finds the dependent index set for all tensor levels in the
|
|
/// current expression we are generating.
|
|
///
|
|
/// For example, when handling A[i+j][j+k], we build the two way mapping in
|
|
/// merger between (tensor, level) pairs and their dependent index variable set:
|
|
/// A_0 <=> [i, j] and A_1 <=> [j, k]
|
|
///
|
|
/// It rejects cases (returns false)
|
|
/// 1st, when the same index is used more than once, e.g., A[i+j][i]
|
|
/// 2nd, when multiplication is used in the non-trivial index expression.
|
|
/// 3rd, when a constant operand is used in the non-trivial index expression.
|
|
///
|
|
/// TODO: constant should be easy to handle.
|
|
static bool findDepIdxSet(Merger &merger, TensorId tensor, Level lvl,
|
|
AffineExpr a, LevelType lt, bool isSubExp = false,
|
|
int64_t coefficient = 1) {
|
|
switch (a.getKind()) {
|
|
case AffineExprKind::DimId: {
|
|
// Only allow positive coefficients on AffineDimExpr.
|
|
if (coefficient <= 0)
|
|
return false;
|
|
|
|
const LoopId idx = merger.makeLoopId(cast<AffineDimExpr>(a).getPosition());
|
|
if (!isUndefLT(merger.getLvlType(tensor, idx)))
|
|
return false; // used more than once, e.g., A[i][i]
|
|
|
|
// TODO: Generalizes the following two cases. A[i] (with trivial index
|
|
// expression) can be treated as a special affine index expression. We do
|
|
// not necessarily need to differentiate them.
|
|
if (!isSubExp) {
|
|
assert(coefficient == 1);
|
|
merger.setLevelAndType(tensor, idx, lvl, lt);
|
|
}
|
|
|
|
if (isSubExp) {
|
|
// The current loops appears in more than one affine expressions on the
|
|
// same tensor. We can not handle this case. e.g., A[i+j][i+k], `i` is
|
|
// used twice.
|
|
if (merger.hasDependentLvl(idx, tensor)) {
|
|
// TODO: This can be supported by coiterate slices if the loop idx is
|
|
// appeared on affine index for different tensor, or take slice on
|
|
// multiple dimensions when it is on the same tensor.
|
|
// E.g.,
|
|
// `d0 + d1` for indexing t0[lvl0] and `d0 + d2` for indexing t1[lvl0]
|
|
// d0_1 = getNextSliceOffset t0 along lvl0
|
|
// d0_2 = getNextSliceOffset t1 along lvl0
|
|
// if d0_1 == d0_2 then d0 = d0_1 = d0_1
|
|
// else increase min(d0_1, d0_2).
|
|
return false;
|
|
}
|
|
merger.setLoopDependentTensorLevel(idx, tensor, lvl, lt, coefficient);
|
|
}
|
|
return true;
|
|
}
|
|
case AffineExprKind::Constant:
|
|
case AffineExprKind::Mul: {
|
|
// TODO: Support index expression like `2 * d0`, we now only support more
|
|
// complicated cases like `2 * d0 + d1`.
|
|
if (!isSubExp)
|
|
return false;
|
|
|
|
// TODO: Support Constant AffineExp for slice-based codegen
|
|
if (isa<AffineConstantExpr>(a))
|
|
llvm_unreachable("Not yet implemented");
|
|
|
|
auto binOp = cast<AffineBinaryOpExpr>(a);
|
|
auto lhs = binOp.getLHS(), rhs = binOp.getRHS();
|
|
if (isa<AffineConstantExpr>(rhs))
|
|
std::swap(lhs, rhs);
|
|
// Must be in form of `constant * d`.
|
|
assert(isa<AffineConstantExpr>(lhs) && isa<AffineDimExpr>(rhs));
|
|
int64_t coefficient = cast<AffineConstantExpr>(lhs).getValue();
|
|
return findDepIdxSet(merger, tensor, lvl, rhs, lt, isSubExp, coefficient);
|
|
}
|
|
case AffineExprKind::Add: {
|
|
auto binOp = cast<AffineBinaryOpExpr>(a);
|
|
return findDepIdxSet(merger, tensor, lvl, binOp.getLHS(), lt, true) &&
|
|
findDepIdxSet(merger, tensor, lvl, binOp.getRHS(), lt, true);
|
|
}
|
|
default:
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/// Gets the total number of compound affine expressions in the
|
|
/// `getMatchingIndexingMap` for the given tensor. For the following inputs:
|
|
///
|
|
/// map = (d0, d1, d2) => (d0 + d1 : compressed, d2 : compressed)
|
|
///
|
|
/// Returns 1 (because the first level is compressed and its corresponding
|
|
/// indexing-expression is `d0 + d1`)
|
|
static unsigned getNumNonTrivialIdxExpOnSparseLvls(AffineMap map,
|
|
Value tensor) {
|
|
// The `tensor` is not guaranteed to have `RankedTensorType`, therefore
|
|
// we can't use `getRankedTensorType`/`getSparseTensorType` here.
|
|
// However, we don't need to handle `StorageSpecifierType`, so we
|
|
// can use `SparseTensorType` once we guard against non-tensors.
|
|
const auto rtp = dyn_cast<RankedTensorType>(tensor.getType());
|
|
if (!rtp)
|
|
return 0;
|
|
const SparseTensorType stt(rtp);
|
|
|
|
const Level lvlRank = stt.getLvlRank();
|
|
const auto exprs = map.getResults();
|
|
assert(static_cast<Dimension>(exprs.size()) == lvlRank &&
|
|
"AffineMap does not have dimension-rank many results");
|
|
unsigned num = 0;
|
|
for (Level l = 0; l < lvlRank; l++) {
|
|
if (!isa<AffineDimExpr>(exprs[l]) && !stt.isDenseLvl(l))
|
|
num++;
|
|
}
|
|
return num;
|
|
}
|
|
|
|
/// Gets the total number of sparse levels with compound affine
|
|
/// expressions, summed over all operands of the `GenericOp`.
|
|
static unsigned getNumNonTrivialIdxExpOnSparseLvls(linalg::GenericOp op) {
|
|
unsigned num = 0;
|
|
for (OpOperand &t : op->getOpOperands())
|
|
num += getNumNonTrivialIdxExpOnSparseLvls(op.getMatchingIndexingMap(&t),
|
|
t.get());
|
|
return num;
|
|
}
|
|
|
|
// Returns true iff output has nontrivial affine indices.
|
|
static bool hasNonTrivialAffineOnSparseOut(linalg::GenericOp op) {
|
|
OpOperand *out = op.getDpsInitOperand(0);
|
|
if (getSparseTensorType(out->get()).isAllDense())
|
|
return false;
|
|
return getNumNonTrivialIdxExpOnSparseLvls(op.getMatchingIndexingMap(out),
|
|
out->get());
|
|
}
|
|
|
|
/// Helper method to inspect sparse encodings in the tensor types.
|
|
/// Fills the per-dimension sparsity information for all tensors.
|
|
/// Returns true if the sparse annotations and affine subscript
|
|
/// expressions of all tensors are admissible. Returns false if
|
|
/// no annotations are found or inadmissible constructs occur.
|
|
/// We currently support two different ways to handle non-trivial index
|
|
/// expression on sparse tensors, and they accept different affine expressions.
|
|
/// When using dependent index reducton-based approach, it currently only
|
|
/// supports affine addition index expression.
|
|
static bool findSparseAnnotations(CodegenEnv &env, bool idxReducBased) {
|
|
bool annotated = false;
|
|
for (OpOperand &t : env.op()->getOpOperands()) {
|
|
const TensorId tid = env.makeTensorId(t.getOperandNumber());
|
|
const auto map = env.op().getMatchingIndexingMap(&t);
|
|
const auto enc = getSparseTensorEncoding(t.get().getType());
|
|
if (enc)
|
|
annotated = true;
|
|
const Level lvlRank = map.getNumResults();
|
|
assert(!enc || lvlRank == enc.getLvlRank());
|
|
assert(static_cast<Level>(env.op().getRank(&t)) == lvlRank);
|
|
// We only need to do index reduction if there is at least one
|
|
// non-trivial index expression on sparse levels. If all non-trivial
|
|
// index expression is on dense levels, we can efficiently rely on
|
|
// the random access to locate the element.
|
|
bool needIdxReduc =
|
|
enc && getNumNonTrivialIdxExpOnSparseLvls(map, t.get()) != 0;
|
|
// If then current tensor being inspected requires affine index, it need
|
|
// to be sliced.
|
|
for (Level l = 0; l < lvlRank; l++) {
|
|
const AffineExpr a = map.getResult(l);
|
|
const LevelType lt = enc.getLvlType(l);
|
|
if (idxReducBased && needIdxReduc) {
|
|
if (!findDepIdxSet(env.merger(), tid, l, a, lt))
|
|
return false; // inadmissible affine expression
|
|
} else {
|
|
if (!findAffine(env.merger(), tid, l, a, lt))
|
|
return false; // inadmissible affine expression
|
|
}
|
|
}
|
|
}
|
|
return annotated;
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Sparsifier synthesis methods (statements and expressions).
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
/// Local bufferization of all dense and sparse data structures.
|
|
static void genBuffers(CodegenEnv &env, OpBuilder &builder) {
|
|
linalg::GenericOp op = env.op();
|
|
Location loc = op.getLoc();
|
|
assert(op.getNumOperands() == op.getNumDpsInputs() + 1);
|
|
|
|
SmallVector<Range, 4> loopRange =
|
|
llvm::cast<linalg::LinalgOp>(op.getOperation())
|
|
.createLoopRanges(builder, loc);
|
|
|
|
env.emitter().initializeLoopEmit(
|
|
builder, loc,
|
|
/// Generates buffer for the output tensor.
|
|
/// Note that all sparse kernels assume that when all elements are written
|
|
/// to (viz. x(i) = y(i) * z(i)), the output buffer is already initialized
|
|
/// to all zeroes and only nonzeroes values are computed and written out.
|
|
/// For updates (viz. x(i) += y(i) * z(i)), only nonzeroes values are used
|
|
/// for the updates and no assumption on the original contents of the
|
|
/// output buffer is necessary.
|
|
[&op](OpBuilder &builder, Location loc, Value memref,
|
|
Value tensor) -> Value {
|
|
// Must not be a sparse tensor.
|
|
assert(!getSparseTensorEncoding(tensor.getType()));
|
|
// Two output tensor references should point to the same object.
|
|
OpOperand *lhs = op.getDpsInitOperand(0);
|
|
assert(lhs->get() == tensor);
|
|
// An output tensor can simply materialize from the buffer of the tensor
|
|
// that appears in the outs() clause. For updates, this has the
|
|
// advantage that only the nonzero value are involved in the
|
|
// computation, keeping the operation O(nnz). In all other cases, we are
|
|
// forced to zero out the buffer to enforce the assumption above, which
|
|
// may negatively impact running complexity (viz. O(n^2 + nnz) vs.
|
|
// O(nnz) for matrices).
|
|
// TODO: use better analysis to avoid zeroing out the buffer?
|
|
bool isInit = op.isInitTensor(lhs);
|
|
Value init = memref;
|
|
if (!isInit) {
|
|
Value zero = constantZero(builder, loc,
|
|
getElementTypeOrSelf(tensor.getType()));
|
|
builder.create<linalg::FillOp>(loc, ValueRange{zero},
|
|
ValueRange{init});
|
|
}
|
|
return init;
|
|
},
|
|
[&loopRange](OpBuilder &b, Location loc, Level l) {
|
|
assert(l < loopRange.size());
|
|
return mlir::getValueOrCreateConstantIndexOp(b, loc, loopRange[l].size);
|
|
});
|
|
}
|
|
|
|
/// Generates index for load/store on sparse tensor.
|
|
static Value genIndex(CodegenEnv &env, OpOperand *t) {
|
|
const auto map = env.op().getMatchingIndexingMap(t);
|
|
const auto stt = getSparseTensorType(t->get());
|
|
const Level lvlRank = stt.getLvlRank();
|
|
assert(static_cast<Level>(map.getNumResults()) == lvlRank);
|
|
const AffineExpr a = map.getResult(lvlRank - 1);
|
|
assert(a.getKind() == AffineExprKind::DimId);
|
|
const LoopId idx = env.makeLoopId(cast<AffineDimExpr>(a).getPosition());
|
|
return env.getLoopVar(idx);
|
|
}
|
|
|
|
/// Generates subscript for load/store on a dense or sparse tensor.
|
|
static Value genSubscript(CodegenEnv &env, OpBuilder &builder, OpOperand *t,
|
|
SmallVectorImpl<Value> &args) {
|
|
const Location loc = env.op().getLoc();
|
|
const TensorId tid = env.makeTensorId(t->getOperandNumber());
|
|
const auto map = env.op().getMatchingIndexingMap(t);
|
|
const auto stt = getSparseTensorType(t->get());
|
|
if (stt.hasEncoding()) {
|
|
// For sparse tensors we only push the last-level's position onto `args`.
|
|
const auto pos = env.emitter().getPosits()[tid].back();
|
|
assert(pos);
|
|
args.push_back(pos);
|
|
} else {
|
|
// For dense tensors we push all level's coordinates onto `args`.
|
|
const Level lvlRank = stt.getLvlRank();
|
|
assert(static_cast<Level>(map.getNumResults()) == lvlRank);
|
|
for (Level l = 0; l < lvlRank; l++) {
|
|
const auto lvlExpr = map.getResult(l);
|
|
const auto lvlCrd = env.emitter().genAffine(builder, loc, lvlExpr);
|
|
args.push_back(lvlCrd);
|
|
}
|
|
}
|
|
return env.emitter().getValBuffer()[tid];
|
|
}
|
|
|
|
/// Generates insertion code to implement dynamic tensor load.
|
|
static Value genInsertionLoad(CodegenEnv &env, OpBuilder &builder,
|
|
OpOperand *t) {
|
|
linalg::GenericOp op = env.op();
|
|
Location loc = op.getLoc();
|
|
// Direct lexicographic coordinate order, tensor loads as zero.
|
|
if (!env.isExpand()) {
|
|
Type tp = getElementTypeOrSelf(t->get().getType());
|
|
return constantZero(builder, loc, tp);
|
|
}
|
|
// Load from expanded access pattern.
|
|
Value index = genIndex(env, t);
|
|
return builder.create<memref::LoadOp>(loc, env.getExpandValues(), index);
|
|
}
|
|
|
|
/// Generates insertion code to implement dynamic tensor load for reduction.
|
|
static Value genInsertionLoadReduce(CodegenEnv &env, OpBuilder &builder,
|
|
OpOperand *t) {
|
|
linalg::GenericOp op = env.op();
|
|
Location loc = op.getLoc();
|
|
Value identity = env.getCustomRedId();
|
|
// Direct lexicographic coordinate order, tensor loads as identity.
|
|
if (!env.isExpand())
|
|
return identity;
|
|
// Load from expanded access pattern if filled, identity otherwise.
|
|
Value values = env.getExpandValues();
|
|
Value filled = env.getExpandFilled();
|
|
Value index = genIndex(env, t);
|
|
Value isFilled = builder.create<memref::LoadOp>(loc, filled, index);
|
|
Value valAtIndex = builder.create<memref::LoadOp>(loc, values, index);
|
|
return builder.create<arith::SelectOp>(loc, isFilled, valAtIndex, identity);
|
|
}
|
|
|
|
/// Generates insertion code to implement dynamic tensor store.
|
|
static void genInsertionStore(CodegenEnv &env, OpBuilder &builder, OpOperand *t,
|
|
Value rhs) {
|
|
linalg::GenericOp op = env.op();
|
|
Location loc = op.getLoc();
|
|
// Direct insertion in lexicographic coordinate order.
|
|
if (!env.isExpand()) {
|
|
const LoopId numLoops = op.getRank(t);
|
|
// Retrieves the first `numLoop` induction variables.
|
|
SmallVector<Value> ivs = llvm::to_vector(llvm::drop_end(
|
|
env.emitter().getLoopIVsRange(), env.getCurrentDepth() - numLoops));
|
|
Value chain = env.getInsertionChain();
|
|
if (env.isValidLexInsert()) {
|
|
// Generates runtime check for a valid lex during reduction,
|
|
// to avoid inserting the identity value for empty reductions.
|
|
// if (validLexInsert) then
|
|
// insert(rhs) into chain
|
|
// return updated chain
|
|
// else
|
|
// return unmodified chain
|
|
scf::IfOp ifValidLexInsert = builder.create<scf::IfOp>(
|
|
loc, chain.getType(), env.getValidLexInsert(),
|
|
/*else=*/true);
|
|
// True branch.
|
|
builder.setInsertionPointToStart(ifValidLexInsert.thenBlock());
|
|
Value res = builder.create<InsertOp>(loc, rhs, chain, ivs);
|
|
builder.create<scf::YieldOp>(loc, res);
|
|
// False branch.
|
|
builder.setInsertionPointToStart(ifValidLexInsert.elseBlock());
|
|
builder.create<scf::YieldOp>(loc, chain);
|
|
// Value assignment.
|
|
builder.setInsertionPointAfter(ifValidLexInsert);
|
|
env.updateInsertionChain(ifValidLexInsert.getResult(0));
|
|
} else {
|
|
// Generates regular insertion chain.
|
|
env.updateInsertionChain(builder.create<InsertOp>(loc, rhs, chain, ivs));
|
|
}
|
|
return;
|
|
}
|
|
// Generates insertion code along expanded access pattern.
|
|
// if (!expFilled[i]) then
|
|
// expFilled[i] = true
|
|
// expAdded[inserts++] = i
|
|
// endif
|
|
// values[i] = rhs
|
|
Value values = env.getExpandValues();
|
|
Value filled = env.getExpandFilled();
|
|
Value added = env.getExpandAdded();
|
|
Value count = env.getExpandCount();
|
|
Value index = genIndex(env, t);
|
|
Value fval = constantI1(builder, loc, false);
|
|
Value tval = constantI1(builder, loc, true);
|
|
// If statement.
|
|
Value isFilled = builder.create<memref::LoadOp>(loc, filled, index);
|
|
Value cond = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
|
|
isFilled, fval);
|
|
scf::IfOp ifOp = builder.create<scf::IfOp>(loc, builder.getIndexType(), cond,
|
|
/*else=*/true);
|
|
// True branch.
|
|
builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
|
|
builder.create<memref::StoreOp>(loc, tval, filled, index);
|
|
builder.create<memref::StoreOp>(loc, index, added, count);
|
|
Value one = constantIndex(builder, loc, 1);
|
|
Value add = builder.create<arith::AddIOp>(loc, count, one);
|
|
builder.create<scf::YieldOp>(loc, add);
|
|
// False branch.
|
|
builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
|
|
builder.create<scf::YieldOp>(loc, count);
|
|
builder.setInsertionPointAfter(ifOp);
|
|
// Value assignment.
|
|
env.updateExpandCount(ifOp.getResult(0));
|
|
builder.create<memref::StoreOp>(loc, rhs, values, index);
|
|
}
|
|
|
|
/// Generates a load on a dense or sparse tensor.
|
|
static Value genTensorLoad(CodegenEnv &env, OpBuilder &builder, ExprId exp) {
|
|
// Test if the load was hoisted to a higher loop nest.
|
|
Value val = env.exp(exp).val;
|
|
if (val)
|
|
return val;
|
|
// Load during insertion.
|
|
linalg::GenericOp op = env.op();
|
|
OpOperand *t = &op->getOpOperand(env.exp(exp).tensor);
|
|
if (env.isSparseOutput(t)) {
|
|
if (env.isCustomReduc())
|
|
return genInsertionLoadReduce(env, builder, t);
|
|
return genInsertionLoad(env, builder, t);
|
|
}
|
|
// Actual load.
|
|
SmallVector<Value> args;
|
|
Value ptr = genSubscript(env, builder, t, args);
|
|
return builder.create<memref::LoadOp>(op.getLoc(), ptr, args);
|
|
}
|
|
|
|
/// Generates a store on a dense or sparse tensor.
|
|
static void genTensorStore(CodegenEnv &env, OpBuilder &builder, ExprId exp,
|
|
Value rhs) {
|
|
// Only unary and binary are allowed to return an uninitialized rhs
|
|
// to indicate missing output. Or otherwise a custom reduction that
|
|
// received no value to accumulate.
|
|
if (!rhs) {
|
|
assert(env.exp(exp).kind == TensorExp::Kind::kUnary ||
|
|
env.exp(exp).kind == TensorExp::Kind::kBinary ||
|
|
env.exp(exp).kind == TensorExp::Kind::kReduce);
|
|
return;
|
|
}
|
|
// Test if this is a scalarized reduction.
|
|
if (env.isReduc()) {
|
|
env.updateReduc(rhs);
|
|
return;
|
|
}
|
|
// Regular store.
|
|
linalg::GenericOp op = env.op();
|
|
Location loc = op.getLoc();
|
|
OpOperand *t = op.getDpsInitOperand(0);
|
|
if (!env.isSparseOutput(t)) {
|
|
SmallVector<Value> args;
|
|
Value ptr = genSubscript(env, builder, t, args);
|
|
builder.create<memref::StoreOp>(loc, rhs, ptr, args);
|
|
return;
|
|
}
|
|
// Store during sparse insertion.
|
|
if (env.exp(exp).kind != TensorExp::Kind::kSelect) {
|
|
genInsertionStore(env, builder, t, rhs);
|
|
return;
|
|
}
|
|
// Select operation insertion.
|
|
Value chain = env.getInsertionChain();
|
|
scf::IfOp ifOp =
|
|
builder.create<scf::IfOp>(loc, chain.getType(), rhs, /*else=*/true);
|
|
builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
|
|
// Existing value was preserved to be used here.
|
|
assert(env.exp(exp).val);
|
|
Value v0 = env.exp(exp).val;
|
|
genInsertionStore(env, builder, t, v0);
|
|
env.merger().clearExprValue(exp);
|
|
// Yield modified insertion chain along true branch.
|
|
Value mchain = env.getInsertionChain();
|
|
builder.create<scf::YieldOp>(op.getLoc(), mchain);
|
|
// Yield original insertion chain along false branch.
|
|
builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
|
|
builder.create<scf::YieldOp>(loc, chain);
|
|
// Done with if statement.
|
|
env.updateInsertionChain(ifOp->getResult(0));
|
|
builder.setInsertionPointAfter(ifOp);
|
|
}
|
|
|
|
/// Generates an invariant value.
|
|
inline static Value genInvariantValue(CodegenEnv &env, ExprId exp) {
|
|
return env.exp(exp).val;
|
|
}
|
|
|
|
/// Semi-ring branches are simply inlined by the sparsifier. Prior
|
|
/// analysis has verified that all computations are "local" to the inlined
|
|
/// branch or otherwise invariantly defined outside the loop nest, with the
|
|
/// exception of index computations, which need to be relinked to actual
|
|
/// inlined cloned code.
|
|
static Value relinkBranch(CodegenEnv &env, RewriterBase &rewriter, Block *block,
|
|
Value e) {
|
|
if (auto arg = dyn_cast<BlockArgument>(e)) {
|
|
// Direct arguments of the original linalg op must be converted
|
|
// into dense tensor loads. Note that we should not encounter
|
|
// anything else. This needs to be verified by semi-ring ops.
|
|
linalg::GenericOp op = env.op();
|
|
if (arg.getOwner()->getParentOp() == op) {
|
|
const TensorId tid = env.makeTensorId(arg.getArgNumber());
|
|
OpOperand *t = &op->getOpOperand(tid);
|
|
assert(!getSparseTensorType(t->get()).hasEncoding()); // dense!
|
|
SmallVector<Value> args;
|
|
Value ptr = genSubscript(env, rewriter, t, args);
|
|
return rewriter.create<memref::LoadOp>(op.getLoc(), ptr, args);
|
|
}
|
|
} else if (Operation *def = e.getDefiningOp()) {
|
|
// Handle index computation.
|
|
if (auto indexOp = dyn_cast<linalg::IndexOp>(def))
|
|
return env.getLoopVar(env.makeLoopId(indexOp.getDim()));
|
|
// When still defined in new body, recurse into operands.
|
|
if (def->getBlock() == block) {
|
|
rewriter.setInsertionPoint(def);
|
|
for (unsigned i = 0, n = def->getNumOperands(); i < n; i++) {
|
|
rewriter.updateRootInPlace(def, [&]() {
|
|
def->setOperand(
|
|
i, relinkBranch(env, rewriter, block, def->getOperand(i)));
|
|
});
|
|
}
|
|
}
|
|
}
|
|
return e;
|
|
}
|
|
|
|
/// Recursively generates tensor expression.
|
|
static Value genExp(CodegenEnv &env, RewriterBase &rewriter, ExprId e) {
|
|
if (e == ::mlir::sparse_tensor::detail::kInvalidId)
|
|
return Value();
|
|
|
|
linalg::GenericOp op = env.op();
|
|
Location loc = op.getLoc();
|
|
const TensorExp &exp = env.exp(e);
|
|
const auto kind = exp.kind;
|
|
if (kind == TensorExp::Kind::kTensor)
|
|
return genTensorLoad(env, rewriter, e);
|
|
if (kind == TensorExp::Kind::kInvariant)
|
|
return genInvariantValue(env, e);
|
|
if (kind == TensorExp::Kind::kLoopVar)
|
|
return env.getLoopVar(exp.loop);
|
|
|
|
if (kind == TensorExp::Kind::kReduce)
|
|
env.startCustomReduc(e); // enter custom
|
|
|
|
// If either lhs/rhs is a synthetic zero, we infer the type for the zero value
|
|
// based on the type of the other operand.
|
|
Value v0, v1;
|
|
if (exp.children.e0 != ::mlir::sparse_tensor::detail::kInvalidId &&
|
|
env.exp(exp.children.e0).kind == TensorExp::Kind::kSynZero) {
|
|
v1 = genExp(env, rewriter, exp.children.e1);
|
|
v0 = constantZero(rewriter, loc, v1.getType());
|
|
} else if (exp.children.e1 != ::mlir::sparse_tensor::detail::kInvalidId &&
|
|
env.exp(exp.children.e1).kind == TensorExp::Kind::kSynZero) {
|
|
v0 = genExp(env, rewriter, exp.children.e0);
|
|
v1 = constantZero(rewriter, loc, v0.getType());
|
|
} else {
|
|
v0 = genExp(env, rewriter, exp.children.e0);
|
|
v1 = genExp(env, rewriter, exp.children.e1);
|
|
}
|
|
|
|
Value ee;
|
|
if (kind == TensorExp::Kind::kReduce && (!v0 || !v1)) {
|
|
// custom reduce did not receive a value
|
|
} else {
|
|
ee = env.merger().buildExp(rewriter, loc, e, v0, v1);
|
|
if (ee &&
|
|
(kind == TensorExp::Kind::kUnary || kind == TensorExp::Kind::kBinary ||
|
|
kind == TensorExp::Kind::kBinaryBranch ||
|
|
kind == TensorExp::Kind::kReduce ||
|
|
kind == TensorExp::Kind::kSelect)) {
|
|
OpBuilder::InsertionGuard guard(rewriter);
|
|
ee = relinkBranch(env, rewriter, ee.getParentBlock(), ee);
|
|
}
|
|
}
|
|
|
|
if (kind == TensorExp::Kind::kReduce)
|
|
env.endCustomReduc(); // exit custom
|
|
|
|
if (kind == TensorExp::Kind::kSelect)
|
|
env.merger().setExprValue(e, v0); // Preserve value for later use.
|
|
|
|
return ee;
|
|
}
|
|
|
|
/// Hoists loop invariant tensor loads for which indices have been exhausted.
|
|
static void genInvariants(CodegenEnv &env, OpBuilder &builder, ExprId exp,
|
|
LoopId curr, bool isStart) {
|
|
if (exp == ::mlir::sparse_tensor::detail::kInvalidId)
|
|
return;
|
|
if (env.exp(exp).kind == TensorExp::Kind::kTensor) {
|
|
// Inspect tensor indices.
|
|
linalg::GenericOp op = env.op();
|
|
OpOperand &t = op->getOpOperand(env.exp(exp).tensor);
|
|
const auto map = op.getMatchingIndexingMap(&t);
|
|
const auto stt = getSparseTensorType(t.get());
|
|
const Level lvlRank = stt.getLvlRank();
|
|
assert(static_cast<Level>(map.getNumResults()) == lvlRank);
|
|
bool isCurrentLoop = curr == 0; // for scalar tensors
|
|
for (Level l = 0; l < lvlRank; l++) {
|
|
const AffineExpr a = map.getResult(l);
|
|
if (!isInvariantAffine(a, curr, /*out*/ isCurrentLoop))
|
|
return; // still in play
|
|
}
|
|
// All exhausted at current level.
|
|
if (!isCurrentLoop)
|
|
return;
|
|
// Generate code for a scalarized reduction or invariant. Note that
|
|
// because custom reduction lhs may occur several times in the IR,
|
|
// we have a built-in safety for only initializing and wrapping-up
|
|
// the scalarized reduction once.
|
|
OpOperand *lhs = op.getDpsInitOperand(0);
|
|
if (lhs == &t) {
|
|
// Start or end a scalarized reduction.
|
|
if (isStart) {
|
|
if (env.isCustomReduc()) {
|
|
if (!env.isReduc())
|
|
env.startReduc(exp, env.getCustomRedId());
|
|
} else {
|
|
env.startReduc(exp, genTensorLoad(env, builder, exp));
|
|
}
|
|
if (env.hasSparseOutput())
|
|
env.startValidLexInsert(
|
|
constantI1(builder, env.op().getLoc(), false));
|
|
} else {
|
|
if (!env.isCustomReduc() || env.isReduc())
|
|
genTensorStore(env, builder, exp, env.endReduc());
|
|
if (env.hasSparseOutput())
|
|
env.endValidLexInsert();
|
|
}
|
|
} else {
|
|
// Start or end loop invariant hoisting of a tensor load.
|
|
if (isStart) {
|
|
env.merger().setExprValue(exp, genTensorLoad(env, builder, exp));
|
|
} else {
|
|
env.merger().clearExprValue(exp);
|
|
}
|
|
}
|
|
} else if (env.exp(exp).kind != TensorExp::Kind::kInvariant &&
|
|
env.exp(exp).kind != TensorExp::Kind::kLoopVar &&
|
|
env.exp(exp).kind != TensorExp::Kind::kSynZero) {
|
|
// Traverse into the binary operations. Note that we only hoist
|
|
// tensor loads, since subsequent MLIR/LLVM passes know how to
|
|
// deal with all other kinds of derived loop invariants.
|
|
if (env.exp(exp).kind == TensorExp::Kind::kReduce)
|
|
env.startCustomReduc(exp); // enter custom
|
|
const ExprId e0 = env.exp(exp).children.e0;
|
|
const ExprId e1 = env.exp(exp).children.e1;
|
|
genInvariants(env, builder, e0, curr, isStart);
|
|
genInvariants(env, builder, e1, curr, isStart);
|
|
if (env.exp(exp).kind == TensorExp::Kind::kReduce)
|
|
env.endCustomReduc(); // exit custom
|
|
}
|
|
}
|
|
|
|
/// Generates an expanded access pattern in innermost dimension.
|
|
static void genExpand(CodegenEnv &env, OpBuilder &builder, LoopId curr,
|
|
bool isStart) {
|
|
linalg::GenericOp op = env.op();
|
|
OpOperand *lhs = op.getDpsInitOperand(0);
|
|
if (!env.atExpandLevel(lhs, op.getRank(lhs), curr))
|
|
return; // not needed at current level
|
|
assert(!env.isReduc());
|
|
// Generate start or end of an expanded access pattern. Note that because
|
|
// an expansion does not rely on the ongoing contents of the sparse storage
|
|
// scheme, we can use the original tensor as incoming SSA value (which
|
|
// simplifies codegen a bit). If expansion on the actual contents is ever
|
|
// needed, we will need to use the SSA value in the insertion chain instead.
|
|
Value tensor = lhs->get();
|
|
Location loc = op.getLoc();
|
|
if (isStart) {
|
|
auto dynShape = {ShapedType::kDynamic};
|
|
Type etp = cast<ShapedType>(tensor.getType()).getElementType();
|
|
Type t1 = MemRefType::get(dynShape, etp);
|
|
Type t2 = MemRefType::get(dynShape, builder.getI1Type());
|
|
Type t3 = MemRefType::get(dynShape, builder.getIndexType());
|
|
Type t4 = builder.getIndexType();
|
|
auto r = builder.create<ExpandOp>(loc, TypeRange({t1, t2, t3, t4}), tensor);
|
|
assert(r.getNumResults() == 4);
|
|
env.startExpand(r.getResult(0), r.getResult(1), r.getResult(2),
|
|
r.getResult(3));
|
|
} else {
|
|
SmallVector<Value> indices;
|
|
for (LoopId i = 0; i < curr; i++)
|
|
indices.push_back(env.emitter().getLoopIV(i));
|
|
Value values = env.getExpandValues();
|
|
Value filled = env.getExpandFilled();
|
|
Value added = env.getExpandAdded();
|
|
Value count = env.getExpandCount();
|
|
Value chain = env.getInsertionChain();
|
|
Value compress = builder.create<CompressOp>(loc, values, filled, added,
|
|
count, chain, indices);
|
|
env.updateInsertionChain(compress);
|
|
env.endExpand();
|
|
}
|
|
}
|
|
|
|
/// Returns parallelization strategy. Any implicit loop in the Linalg
|
|
/// operation that is marked "parallel" is a candidate. Whether it is actually
|
|
/// converted to a parallel operation depends on the requested strategy.
|
|
static bool isParallelFor(CodegenEnv &env, bool isOuter, bool isSparse) {
|
|
// Reject parallelization of sparse output.
|
|
if (env.hasSparseOutput())
|
|
return false;
|
|
// Parallel loops on tensor expansion can cause data races.
|
|
if (env.isExpand())
|
|
return false;
|
|
// Inspect strategy.
|
|
switch (env.options().parallelizationStrategy) {
|
|
case SparseParallelizationStrategy::kNone:
|
|
return false;
|
|
case SparseParallelizationStrategy::kDenseOuterLoop:
|
|
return isOuter && !isSparse;
|
|
case SparseParallelizationStrategy::kAnyStorageOuterLoop:
|
|
return isOuter;
|
|
case SparseParallelizationStrategy::kDenseAnyLoop:
|
|
return !isSparse;
|
|
case SparseParallelizationStrategy::kAnyStorageAnyLoop:
|
|
return true;
|
|
}
|
|
llvm_unreachable("unexpected parallelization strategy");
|
|
}
|
|
|
|
/// Whether or not the current loop being generated should be parallized (if
|
|
/// possible) according to the configuration.
|
|
static bool shouldTryParallize(CodegenEnv &env, LoopId curr,
|
|
ArrayRef<TensorLevel> tidLvls) {
|
|
linalg::GenericOp op = env.op();
|
|
auto iteratorTypes = op.getIteratorTypesArray();
|
|
bool isSparse = llvm::any_of(tidLvls, [curr, &env](TensorLevel tidLvl) {
|
|
// Queries the LT based on the tensor and loop id, as requested by
|
|
// `CodegenEnv::lt(TensorId, LoopId)`. The returned LT from CodegenEnv
|
|
// should be consistent with the LT indexed by <TensorId, Level>.
|
|
const auto lt = env.lt(env.unpackTensorLevel(tidLvl).first, curr);
|
|
return isCompressedLT(lt) || isSingletonLT(lt);
|
|
});
|
|
return isParallelFor(env, /*isOuter=*/curr == 0, isSparse);
|
|
}
|
|
|
|
/// Emit a loop to coiterate over the list of tensor levels. The generated loop
|
|
/// can either be a for loop or while loop depending on whether there is at most
|
|
/// one sparse level in the list.
|
|
static Operation *genCoIteration(CodegenEnv &env, OpBuilder &builder,
|
|
ArrayRef<TensorLevel> tidLvls,
|
|
bool tryParallel, bool needsUniv) {
|
|
Operation *loop = *env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
|
|
// Construct while-loop with a parameter for each index.
|
|
return env.emitter().enterCoIterationOverTensorsAtLvls(
|
|
builder, env.op().getLoc(), tidLvls, reduc, tryParallel,
|
|
/*genDedup=*/true, needsUniv);
|
|
});
|
|
assert(loop);
|
|
return loop;
|
|
}
|
|
|
|
/// Generates a for-loop or a while-loop, depending on whether it implements
|
|
/// singleton iteration or co-iteration over the given conjunction.
|
|
static Operation *genLoop(CodegenEnv &env, OpBuilder &builder, LoopId curr,
|
|
bool needsUniv, ArrayRef<TensorLevel> tidLvls) {
|
|
bool tryParallel = shouldTryParallize(env, curr, tidLvls);
|
|
return genCoIteration(env, builder, tidLvls, tryParallel, needsUniv);
|
|
}
|
|
|
|
/// Generates the induction structure for a while-loop.
|
|
static void finalizeWhileOp(CodegenEnv &env, OpBuilder &builder,
|
|
bool needsUniv) {
|
|
Location loc = env.op().getLoc();
|
|
// Finalize each else branch of all if statements.
|
|
if (env.isReduc() || env.isExpand() || env.getInsertionChain()) {
|
|
while (auto ifOp = dyn_cast_or_null<scf::IfOp>(
|
|
builder.getInsertionBlock()->getParentOp())) {
|
|
// Break on IfOp for slicing filtering.
|
|
if (ifOp->getAttr(LoopEmitter::getLoopEmitterLoopAttrName()) ==
|
|
StringAttr::get(ifOp->getContext(), "slice"))
|
|
break;
|
|
|
|
unsigned y = 0;
|
|
SmallVector<Value> yields;
|
|
if (env.isReduc()) {
|
|
yields.push_back(env.getReduc());
|
|
env.updateReduc(ifOp.getResult(y++));
|
|
if (env.isValidLexInsert()) {
|
|
yields.push_back(env.getValidLexInsert());
|
|
env.updateValidLexInsert(ifOp.getResult(y++));
|
|
}
|
|
}
|
|
if (env.isExpand()) {
|
|
yields.push_back(env.getExpandCount());
|
|
env.updateExpandCount(ifOp->getResult(y++));
|
|
}
|
|
if (env.getInsertionChain()) {
|
|
yields.push_back(env.getInsertionChain());
|
|
env.updateInsertionChain(ifOp->getResult(y++));
|
|
}
|
|
assert(y == yields.size());
|
|
builder.create<scf::YieldOp>(loc, yields);
|
|
builder.setInsertionPointAfter(ifOp);
|
|
}
|
|
}
|
|
// No need to set the insertion point here as LoopEmitter keeps track of the
|
|
// basic block where scf::Yield should be inserted.
|
|
}
|
|
|
|
/// Generates a single if-statement within a while-loop.
|
|
static scf::IfOp genIf(CodegenEnv &env, OpBuilder &builder, LoopId curr,
|
|
LatPointId p) {
|
|
Location loc = env.op().getLoc();
|
|
SmallVector<Type> types;
|
|
Value cond;
|
|
env.merger().foreachTensorLoopId(
|
|
p, /*simple=*/true,
|
|
[&](TensorLoopId b, TensorId tid, std::optional<Level> lvl, LevelType lt,
|
|
bool isIdxRed) {
|
|
if (isIdxRed) {
|
|
// Since there is no 1:1 mapping from loop to level (multiple loops
|
|
// are required to resolve one level with non-trivial index
|
|
// expression), we need to reconstruct the tensor level types if this
|
|
// loop requires index reduction condition.
|
|
assert(lvl.has_value() && isUndefLT(lt));
|
|
auto stt = getSparseTensorType(env.op().getInputs()[tid]);
|
|
lt = stt.getLvlType(*lvl);
|
|
}
|
|
assert(curr == env.merger().loop(b));
|
|
Value clause;
|
|
if (isCompressedLT(lt) || isSingletonLT(lt) ||
|
|
isLooseCompressedLT(lt) || is2OutOf4LT(lt)) {
|
|
assert(lvl.has_value());
|
|
const Value crd = env.emitter().getCoords()[tid][*lvl];
|
|
const Value lvar = env.getLoopVar(curr);
|
|
clause = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
|
|
crd, lvar);
|
|
} else {
|
|
assert(isDenseLT(lt) || isUndefLT(lt));
|
|
clause = constantI1(builder, loc, true);
|
|
}
|
|
cond = cond ? builder.create<arith::AndIOp>(loc, cond, clause) : clause;
|
|
});
|
|
if (env.isReduc()) {
|
|
types.push_back(env.getReduc().getType());
|
|
if (env.isValidLexInsert())
|
|
types.push_back(env.getValidLexInsert().getType());
|
|
}
|
|
if (env.isExpand())
|
|
types.push_back(builder.getIndexType());
|
|
if (env.getInsertionChain())
|
|
types.push_back(env.getInsertionChain().getType());
|
|
scf::IfOp ifOp = builder.create<scf::IfOp>(loc, types, cond, /*else=*/true);
|
|
builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
|
|
return ifOp;
|
|
}
|
|
|
|
/// Generates end of true branch of if-statement within a while-loop.
|
|
static void endIf(CodegenEnv &env, OpBuilder &builder, scf::IfOp ifOp,
|
|
Value redInput, Value cntInput, Value insInput,
|
|
Value validIns) {
|
|
SmallVector<Value> operands;
|
|
if (env.isReduc()) {
|
|
operands.push_back(env.getReduc());
|
|
env.updateReduc(redInput);
|
|
if (env.isValidLexInsert()) {
|
|
// Any overlapping indices during a reduction creates a valid lex insert.
|
|
operands.push_back(constantI1(builder, env.op().getLoc(), true));
|
|
env.updateValidLexInsert(validIns);
|
|
}
|
|
}
|
|
if (env.isExpand()) {
|
|
operands.push_back(env.getExpandCount());
|
|
env.updateExpandCount(cntInput);
|
|
}
|
|
if (env.getInsertionChain()) {
|
|
operands.push_back(env.getInsertionChain());
|
|
env.updateInsertionChain(insInput);
|
|
}
|
|
if (!operands.empty())
|
|
builder.create<scf::YieldOp>(env.op().getLoc(), operands);
|
|
builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Sparsifier synthesis methods (loop sequence).
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
static bool getAllTidLvlsInLatPoints(
|
|
CodegenEnv &env, LatPointId li, LoopId curr,
|
|
llvm::function_ref<void(TensorLevel, AffineExpr)> callback) {
|
|
const BitVector &simple = env.lat(li).simple;
|
|
const TensorId outTid = env.merger().getOutTensorID();
|
|
const std::optional<Level> outLvl = env.merger().getLvl(outTid, curr);
|
|
|
|
unsigned numloopCond = 0;
|
|
bool hasNonUnique = false;
|
|
env.merger().foreachTensorLoopId(
|
|
li, [&, curr](TensorLoopId b, TensorId tid, std::optional<Level> lvl,
|
|
LevelType lt, bool isIdxReduc) {
|
|
if (simple[b]) {
|
|
if (isIdxReduc) {
|
|
callback(env.makeTensorLevel(tid, *lvl), nullptr);
|
|
numloopCond++;
|
|
return;
|
|
}
|
|
if (isUndefLT(lt)) {
|
|
// An undefined lt in the lattices, we probably mean to
|
|
// generate a dense loop according to the synthetic tensor (for
|
|
// invariants and sparse output tensor).
|
|
if (env.merger().getSynTensorID() == tid) {
|
|
// Coiterating with an invariant
|
|
// e.g., out = prod(in[i][j] op invariant);
|
|
// or a broadcast
|
|
// e.g., out[i][j] = in[i] (j is undef for input)
|
|
//
|
|
// The level of the synthetic tensor is the current loop depth;
|
|
// the rank of the synthetic tensor equals to number of loops.
|
|
assert(curr == env.getCurrentDepth());
|
|
lvl = curr;
|
|
} else if (!lvl) {
|
|
// Skips invalid lvl (e.g., when this is a zero ranked tensor).
|
|
return;
|
|
}
|
|
}
|
|
hasNonUnique = !isUniqueLT(lt) || hasNonUnique;
|
|
callback(env.makeTensorLevel(tid, *lvl), nullptr);
|
|
numloopCond++;
|
|
} else if (isDenseLT(lt) || isIdxReduc) {
|
|
callback(env.makeTensorLevel(tid, *lvl), nullptr);
|
|
} else {
|
|
assert(isUndefLT(lt));
|
|
linalg::GenericOp op = env.op();
|
|
if (tid >= op.getNumDpsInputs())
|
|
// We only handle affine expression on input tensors (for now).
|
|
return;
|
|
OpOperand *operand = &op->getOpOperand(tid);
|
|
const auto stt = getSparseTensorType(operand->get());
|
|
// Non-annotated dense tensors requires no special handling.
|
|
if (!stt.hasEncoding())
|
|
return;
|
|
|
|
ArrayRef<AffineExpr> affines =
|
|
op.getMatchingIndexingMap(operand).getResults();
|
|
const Level lvlRank = stt.getLvlRank();
|
|
assert(affines.size() == static_cast<size_t>(lvlRank));
|
|
for (Level l = 0; l < lvlRank; l++) {
|
|
AffineExpr exp = affines[l];
|
|
// Skip simple affine expression and non-dense levels (which
|
|
// have their own filter loop).
|
|
if (isa<AffineDimExpr>(exp) || !stt.isDenseLvl(l))
|
|
continue;
|
|
|
|
// Constant affine expression are handled in genLoop.
|
|
if (!isa<AffineConstantExpr>(exp)) {
|
|
bool isCurrentLoop = false;
|
|
assert(curr == env.getCurrentDepth());
|
|
if (isInvariantAffine(exp, curr + 1, /*out*/ isCurrentLoop) &&
|
|
isCurrentLoop) {
|
|
// If the compound affine is invariant and we are right at the
|
|
// level. We need to generate the address according to the
|
|
// affine expression. This is also the best place we can do it
|
|
// to avoid putting it inside inner loops.
|
|
callback(env.makeTensorLevel(tid, l), exp);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
if (isDenseLT(env.lt(outTid, curr))) {
|
|
// Note that we generate dense indices of the output tensor
|
|
// unconditionally, since they may not appear in the lattice, but may be
|
|
// needed for linearized env.
|
|
callback(env.makeTensorLevel(outTid, *outLvl), nullptr);
|
|
}
|
|
|
|
if (numloopCond == 0) {
|
|
// Corner cases where the loop bound is defined by a *unused* operand, in
|
|
// this case, we just generate a dense "fake" loop by iterating over the
|
|
// synthetic tensor.
|
|
callback(env.makeTensorLevel(env.merger().getSynTensorID(), curr), nullptr);
|
|
numloopCond++;
|
|
}
|
|
// If we just need to one loop conditions and the conditions is not imposed on
|
|
// non-unique level, the loop can be generated by a for loop.
|
|
return numloopCond == 1 && !hasNonUnique;
|
|
}
|
|
|
|
/// Starts a loop sequence at given level. Returns true if
|
|
/// the universal loop index must be maintained at this level.
|
|
static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp,
|
|
LoopId curr, LatSetId lts) {
|
|
assert(!env.getLoopVar(curr));
|
|
// Emit invariants at this loop sequence level.
|
|
genInvariants(env, builder, exp, curr, /*isStart=*/true);
|
|
// Emit access pattern expansion for sparse tensor output.
|
|
genExpand(env, builder, curr, /*isStart=*/true);
|
|
// Emit further initialization at this loop sequence level.
|
|
const LatPointId l0 = env.set(lts)[0];
|
|
|
|
SmallVector<TensorLevel> tidLvls;
|
|
getAllTidLvlsInLatPoints(env, l0, curr, [&](TensorLevel tl, AffineExpr) {
|
|
tidLvls.emplace_back(tl);
|
|
});
|
|
|
|
env.emitter().enterNewLoopSeq(builder, env.op().getLoc(), tidLvls);
|
|
|
|
// Maintain the universal index only if it is actually
|
|
// consumed by a subsequent lattice point.
|
|
for (const LatPointId li : env.set(lts).drop_front())
|
|
if (!env.merger().hasAnySparse(env.lat(li).simple))
|
|
return true;
|
|
|
|
return false;
|
|
}
|
|
|
|
// Generates dense affine address for encoding.
|
|
static void genConstantDenseAddressFromLevel(CodegenEnv &env,
|
|
OpBuilder &builder, TensorId tid,
|
|
Level startLvl) {
|
|
// TODO: Handle affine expression on output tensor.
|
|
linalg::GenericOp op = env.op();
|
|
assert(tid < op.getNumDpsInputs());
|
|
OpOperand *input = op.getDpsInputOperands()[tid];
|
|
const auto lvlExprs = op.getMatchingIndexingMap(input).getResults();
|
|
const auto enc = getSparseTensorEncoding(input->get().getType());
|
|
if (enc) {
|
|
const Location loc = op.getLoc();
|
|
const TensorId tid = env.makeTensorId(input->getOperandNumber());
|
|
const Level lvlRank = enc.getLvlRank();
|
|
assert(lvlExprs.size() == static_cast<size_t>(lvlRank));
|
|
for (Level l = startLvl; l < lvlRank; l++) {
|
|
AffineExpr lvlExpr = lvlExprs[l];
|
|
if (enc.isDenseLvl(l) && isa<AffineConstantExpr>(lvlExpr))
|
|
env.emitter().genDenseAffineAddress(
|
|
builder, loc, env.makeTensorLevel(tid, l), lvlExpr);
|
|
else
|
|
return; // break on first non-dense non-constant level
|
|
}
|
|
}
|
|
}
|
|
|
|
// We can generate address for constant affine expression before any loops
|
|
// starting from the first level as they do not depend on anything.
|
|
// E.g., [Dense, Dense, Sparse] -> (1, 2, d0), the addresses for the first two
|
|
// levels can be determined before loops.
|
|
static void genInitConstantDenseAddress(CodegenEnv &env,
|
|
RewriterBase &rewriter) {
|
|
for (TensorId tid = 0, e = env.op().getNumDpsInputs(); tid < e; tid++)
|
|
genConstantDenseAddressFromLevel(env, rewriter, tid, 0);
|
|
}
|
|
|
|
/// Returns true if the lattice bit can be iterated by a for loop.
|
|
static bool translateBitsToTidLvlPairs(
|
|
CodegenEnv &env, LatPointId li, LoopId curr,
|
|
SmallVectorImpl<TensorLevel> &tidLvls,
|
|
SmallVectorImpl<std::pair<TensorLevel, AffineExpr>> &affineTidLvls) {
|
|
return getAllTidLvlsInLatPoints(env, li, curr,
|
|
[&](TensorLevel tl, AffineExpr exp) {
|
|
if (exp)
|
|
affineTidLvls.emplace_back(tl, exp);
|
|
else
|
|
tidLvls.emplace_back(tl);
|
|
});
|
|
}
|
|
|
|
/// Starts a single loop in current sequence.
|
|
static std::pair<Operation *, bool> startLoop(CodegenEnv &env,
|
|
OpBuilder &builder, LoopId curr,
|
|
LatPointId li, bool needsUniv) {
|
|
// The set of tensors + lvls to generate loops on
|
|
SmallVector<TensorLevel> tidLvls;
|
|
|
|
// The set of dense tensors with non-trivial affine expression that just
|
|
// becomes invariant and the address are generated at the current level.
|
|
SmallVector<std::pair<TensorLevel, AffineExpr>> affineTidLvls;
|
|
bool isSingleCond =
|
|
translateBitsToTidLvlPairs(env, li, curr, tidLvls, affineTidLvls);
|
|
|
|
// Emit the for/while-loop control.
|
|
Operation *loop = genLoop(env, builder, curr, needsUniv, tidLvls);
|
|
Location loc = env.op().getLoc();
|
|
for (auto [tidLvl, exp] : affineTidLvls) {
|
|
env.emitter().genDenseAffineAddress(builder, loc, tidLvl, exp);
|
|
}
|
|
|
|
// Until now, we have entered every <tid, lvl> pair in {cond, extra,
|
|
// affine}Tids/Lvls. The addresses of the upcoming levels which are dependent
|
|
// on constant affines expression may now be determined.
|
|
auto allTidLvls =
|
|
llvm::concat<TensorLevel>(tidLvls, llvm::make_first_range(affineTidLvls));
|
|
for (auto [tid, lvl] : env.unpackTensorLevelRange(allTidLvls)) {
|
|
if (tid != env.merger().getOutTensorID() &&
|
|
tid != env.merger().getSynTensorID())
|
|
genConstantDenseAddressFromLevel(env, builder, tid, lvl + 1);
|
|
}
|
|
|
|
return std::make_pair(loop, isSingleCond);
|
|
}
|
|
|
|
/// Ends a single loop in current sequence. Returns new values for needsUniv.
|
|
static bool endLoop(CodegenEnv &env, RewriterBase &rewriter, Operation *loop,
|
|
LatPointId li, bool needsUniv, bool isSingleCond) {
|
|
// Either a for-loop or a while-loop that iterates over a slice.
|
|
if (isSingleCond) {
|
|
// Any iteration creates a valid lex insert.
|
|
if (env.isReduc() && env.isValidLexInsert())
|
|
env.updateValidLexInsert(constantI1(rewriter, env.op().getLoc(), true));
|
|
} else if (auto whileOp = dyn_cast<scf::WhileOp>(loop)) {
|
|
// End a while-loop.
|
|
finalizeWhileOp(env, rewriter, needsUniv);
|
|
} else {
|
|
needsUniv = false;
|
|
}
|
|
env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
|
|
env.emitter().exitCurrentLoop(rewriter, env.op().getLoc(), reduc);
|
|
return std::nullopt;
|
|
});
|
|
return needsUniv;
|
|
}
|
|
|
|
/// Ends a loop sequence at given level.
|
|
static void endLoopSeq(CodegenEnv &env, OpBuilder &builder, unsigned exp,
|
|
unsigned at) {
|
|
assert(!env.getLoopVar(at));
|
|
env.emitter().exitCurrentLoopSeq(builder, env.op().getLoc());
|
|
// Unmark bookkeeping of invariants and loop index.
|
|
genInvariants(env, builder, exp, at, /*isStart=*/false);
|
|
// Finalize access pattern expansion for sparse tensor output.
|
|
genExpand(env, builder, at, /*isStart=*/false);
|
|
}
|
|
|
|
/// Recursively generates code while computing iteration lattices in order
|
|
/// to manage the complexity of implementing co-iteration over unions
|
|
/// and intersections of sparse iterations spaces.
|
|
static void genStmt(CodegenEnv &env, RewriterBase &rewriter, ExprId exp,
|
|
LoopId curr) {
|
|
assert(curr == env.getCurrentDepth());
|
|
|
|
// At each leaf, assign remaining tensor (sub)expression to output tensor.
|
|
if (curr == env.getLoopNum()) {
|
|
Value rhs = genExp(env, rewriter, exp);
|
|
genTensorStore(env, rewriter, exp, rhs);
|
|
return;
|
|
}
|
|
|
|
// Construct iteration lattices for current loop index.
|
|
const LatSetId lts =
|
|
env.merger().optimizeSet(env.merger().buildLattices(exp, curr));
|
|
|
|
// Start a loop sequence.
|
|
bool needsUniv = startLoopSeq(env, rewriter, exp, curr, lts);
|
|
|
|
// Emit a loop for every lattice point L0 >= Li in this loop sequence.
|
|
// We cannot change this to `for (const LatPointId li : env.set(lts))`
|
|
// because the loop body causes data-movement which invalidates
|
|
// the iterator.
|
|
const unsigned lsize = env.set(lts).size();
|
|
for (unsigned i = 0; i < lsize; i++) {
|
|
const LatPointId li = env.set(lts)[i];
|
|
// Start a loop.
|
|
auto [loop, isSingleCond] = startLoop(env, rewriter, curr, li, needsUniv);
|
|
|
|
// Visit all lattices points with Li >= Lj to generate the
|
|
// loop-body, possibly with if statements for coiteration.
|
|
Value redInput = env.getReduc();
|
|
Value cntInput = env.getExpandCount();
|
|
Value insInput = env.getInsertionChain();
|
|
Value validIns = env.getValidLexInsert();
|
|
// We cannot change this to `for (const LatPointId lj : env.set(lts))`
|
|
// because the loop body causes data-movement which invalidates the
|
|
// iterator.
|
|
for (unsigned j = 0; j < lsize; j++) {
|
|
const LatPointId lj = env.set(lts)[j];
|
|
const ExprId ej = env.lat(lj).exp;
|
|
if (li == lj || env.merger().latGT(li, lj)) {
|
|
// Recurse into body of each branch.
|
|
if (!isSingleCond) {
|
|
scf::IfOp ifOp = genIf(env, rewriter, curr, lj);
|
|
genStmt(env, rewriter, ej, curr + 1);
|
|
endIf(env, rewriter, ifOp, redInput, cntInput, insInput, validIns);
|
|
} else {
|
|
genStmt(env, rewriter, ej, curr + 1);
|
|
}
|
|
}
|
|
}
|
|
|
|
// End a loop.
|
|
needsUniv = endLoop(env, rewriter, loop, curr, needsUniv, isSingleCond);
|
|
}
|
|
|
|
// End a loop sequence.
|
|
endLoopSeq(env, rewriter, exp, curr);
|
|
assert(curr == env.getCurrentDepth());
|
|
}
|
|
|
|
/// Converts the result computed by the sparse kernel into the required form.
|
|
static void genResult(CodegenEnv &env, RewriterBase &rewriter) {
|
|
linalg::GenericOp op = env.op();
|
|
OpOperand *lhs = op.getDpsInitOperand(0);
|
|
Value tensor = lhs->get();
|
|
Type resType = tensor.getType();
|
|
if (getSparseTensorEncoding(resType)) {
|
|
// The sparse tensor rematerializes from the original sparse tensor's
|
|
// underlying sparse storage format. For an insertion chain, the
|
|
// tensor materializes from the chain with 'hasInserts' enabled.
|
|
bool hasInserts = false;
|
|
if (Value chain = env.getInsertionChain()) {
|
|
hasInserts = true;
|
|
tensor = chain;
|
|
}
|
|
rewriter.replaceOpWithNewOp<LoadOp>(op, resType, tensor, hasInserts);
|
|
} else {
|
|
// To rematerialize an non-annotated tensor, simply load it
|
|
// from the bufferized value.
|
|
Value val = env.emitter().getValBuffer()[env.merger().getOutTensorID()];
|
|
rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, resType, val);
|
|
}
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Sparsifier rewriting methods.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace {
|
|
|
|
/// Sparse rewriting rule for generic Lingalg operation.
|
|
struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
|
|
public:
|
|
GenericOpSparsifier(MLIRContext *context, SparsificationOptions o)
|
|
: OpRewritePattern<linalg::GenericOp>(context), options(o) {}
|
|
|
|
LogicalResult matchAndRewrite(linalg::GenericOp op,
|
|
PatternRewriter &rewriter) const override {
|
|
// Only accept single output operations with pure tensor semantics.
|
|
if (op.getNumDpsInits() != 1 || !op.hasPureTensorSemantics())
|
|
return failure();
|
|
|
|
// Only accept trivial affine indices.
|
|
if (hasNonTrivialAffineOnSparseOut(op))
|
|
return failure();
|
|
|
|
// Only accept scheduled loops.
|
|
if (!op->hasAttr("sorted")) {
|
|
return rewriter.notifyMatchFailure(
|
|
op, "Loops not yet scheduled, try run --sparse-reinterpret-map "
|
|
"before sparsification.");
|
|
}
|
|
|
|
// Must have been demapped as well if the generic op is sorted.
|
|
assert(!hasAnyNonIdentityOperandsOrResults(op));
|
|
|
|
// Sets up a code generation environment.
|
|
const unsigned numTensors = op->getNumOperands();
|
|
const unsigned numLoops = op.getNumLoops();
|
|
bool needIdxRed = getNumNonTrivialIdxExpOnSparseLvls(op) != 0;
|
|
// If we have indexing map like (d0) -> (0, d0), there might be more
|
|
// levels then loops because of the constant index, that means we can not
|
|
// use numLoops as the upper bound for ranks of all tensors.
|
|
// TODO: Constant indices are currently not support on sparse tensor, but
|
|
// are allowed in non-annotated dense tensor. Support it, it would be
|
|
// required for sparse tensor slice rank reducing too.
|
|
Level maxLvlRank = 0;
|
|
for (auto operand : op.getOperands()) {
|
|
if (auto rtp = dyn_cast<RankedTensorType>(operand.getType())) {
|
|
maxLvlRank = std::max(maxLvlRank, SparseTensorType(rtp).getLvlRank());
|
|
}
|
|
}
|
|
|
|
// Detects sparse annotations and translates the per-level sparsity
|
|
// information for all tensors to loop indices in the kernel.
|
|
CodegenEnv env(op, options, numTensors, numLoops, maxLvlRank);
|
|
if (!findSparseAnnotations(env, needIdxRed))
|
|
return failure();
|
|
|
|
// Only standard reduction operations (add, sub, or, xor) that can be
|
|
// sparsified by merely reducing the stored values are admissible. More
|
|
// elaborate reduction operations (such as mul, and, min, max) would need
|
|
// to know whether implicit zeros occur as well. They can still be
|
|
// implemented with a custom reduction operation, accepted here as well.
|
|
if (op.getNumReductionLoops() > 0) {
|
|
Operation *yield = op.getRegion().front().getTerminator();
|
|
assert(isa<linalg::YieldOp>(yield));
|
|
Operation *redop = yield->getOperand(0).getDefiningOp();
|
|
if (!isa<arith::AddFOp>(redop) && !isa<complex::AddOp>(redop) &&
|
|
!isa<arith::AddIOp>(redop) && !isa<arith::SubFOp>(redop) &&
|
|
!isa<complex::SubOp>(redop) && !isa<arith::SubIOp>(redop) &&
|
|
!isa<arith::OrIOp>(redop) && !isa<arith::XOrIOp>(redop) &&
|
|
!isa<ReduceOp>(redop)) {
|
|
return failure();
|
|
}
|
|
}
|
|
|
|
// Constructs the tensor expressions tree from `op`, returns failure if the
|
|
// tree can not be built or the tensor expression is inadmissible.
|
|
if (failed(env.initTensorExp()))
|
|
return failure();
|
|
|
|
// Recursively generates code if admissible.
|
|
env.startEmit();
|
|
genBuffers(env, rewriter);
|
|
// TODO: Constant affine expression should be handled differently when using
|
|
// slice-based codegen, it does not matter now because we already reject the
|
|
// constant expression at an earlier stage.
|
|
genInitConstantDenseAddress(env, rewriter);
|
|
genStmt(env, rewriter, env.getExprId(), 0);
|
|
genResult(env, rewriter);
|
|
return success();
|
|
}
|
|
|
|
private:
|
|
/// Options to control sparse code generation.
|
|
SparsificationOptions options;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
/// Populates the given patterns list with rewriting rules required for
|
|
/// the sparsification of linear algebra operations.
|
|
void mlir::populateSparsificationPatterns(
|
|
RewritePatternSet &patterns, const SparsificationOptions &options) {
|
|
patterns.add<GenericOpSparsifier>(patterns.getContext(), options);
|
|
}
|