1925 lines
77 KiB
C++
1925 lines
77 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 "CodegenEnv.h"
|
|
#include "CodegenUtils.h"
|
|
#include "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;
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Declarations
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace {
|
|
|
|
/// Iteration graph sorting.
|
|
enum class SortMask : unsigned {
|
|
// The individual mask bits.
|
|
kIncludeDenseOutput = 0x1, // b001
|
|
kIncludeDenseInput = 0x2, // b010
|
|
kIncludeUndef = 0x4, // b100
|
|
// The subsets of mask bits.
|
|
kIncludeAll = 0x7, // b111
|
|
kIncludeDense = 0x3, // b011
|
|
kSparseOnly = 0x0, // b000
|
|
};
|
|
|
|
inline static bool includesAny(SortMask mask1, SortMask mask2) {
|
|
return static_cast<unsigned>(mask1) & static_cast<unsigned>(mask2);
|
|
}
|
|
|
|
inline static bool includesDenseInput(SortMask mask) {
|
|
return includesAny(mask, SortMask::kIncludeDenseInput);
|
|
}
|
|
|
|
inline static bool includesDenseOutput(SortMask mask) {
|
|
return includesAny(mask, SortMask::kIncludeDenseOutput);
|
|
}
|
|
|
|
inline static bool includesDense(SortMask mask) {
|
|
return includesAny(mask, SortMask::kIncludeDense);
|
|
}
|
|
|
|
inline static bool includesUndef(SortMask mask) {
|
|
return includesAny(mask, SortMask::kIncludeUndef);
|
|
}
|
|
|
|
/// A helper class that visits an affine expression and tries to find an
|
|
/// AffineDimExpr to which the corresponding iterator from a GenericOp matches
|
|
/// the desired iterator type.
|
|
class AffineDimFinder : public AffineExprVisitor<AffineDimFinder> {
|
|
public:
|
|
explicit AffineDimFinder(linalg::GenericOp op)
|
|
: iterTypes(op.getIteratorTypesArray()) {}
|
|
void visitDimExpr(AffineDimExpr expr) {
|
|
if (pickedDim == nullptr || pickIterType == iterTypes[expr.getPosition()]) {
|
|
pickedDim = expr;
|
|
}
|
|
}
|
|
|
|
/// Set the desired iterator type that we want to pick.
|
|
void setPickedIterType(utils::IteratorType iterType) {
|
|
pickIterType = iterType;
|
|
}
|
|
|
|
/// Get the desired AffineDimExpr.
|
|
AffineDimExpr getDimExpr() const { return pickedDim.cast<AffineDimExpr>(); }
|
|
|
|
private:
|
|
/// The picked AffineDimExpr after visit. This must be stored as
|
|
/// `AffineExpr` rather than `AffineDimExpr`, because the latter
|
|
/// doesn't have a default ctor.
|
|
AffineExpr pickedDim;
|
|
/// The iterator type that we want.
|
|
utils::IteratorType pickIterType;
|
|
/// The mapping between dim=>iterator type.
|
|
SmallVector<utils::IteratorType> iterTypes;
|
|
};
|
|
|
|
// Flattens an affine expression into a list of AffineDimExprs.
|
|
struct AffineDimCollector : public AffineExprVisitor<AffineDimCollector> {
|
|
void visitDimExpr(AffineDimExpr expr) { dims.push_back(expr); }
|
|
SmallVector<AffineDimExpr> dims;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Sparse compiler analysis methods.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
// TODO: the "idx"-vs-"ldx" naming convention is not self-explanatory,
|
|
// and those letters are too easy to confuse visually. We should switch
|
|
// to a more self-explanatory naming convention like "curLoop"-vs-"prevLoop"
|
|
// (assuming that's the actual meaning behind the "idx"-vs-"ldx" convention).
|
|
|
|
/// Determines if affine expression is invariant.
|
|
static bool isInvariantAffine(AffineExpr a, ArrayRef<LoopId> loopStack,
|
|
LoopId ldx, bool &isAtLoop) {
|
|
switch (a.getKind()) {
|
|
case AffineExprKind::DimId: {
|
|
const LoopId i = a.cast<AffineDimExpr>().getPosition();
|
|
if (i == ldx) {
|
|
isAtLoop = true;
|
|
// Must be invariant if we are at the given loop.
|
|
return true;
|
|
}
|
|
bool isInvariant = false;
|
|
for (LoopId l : loopStack) {
|
|
isInvariant = (l == i);
|
|
if (isInvariant)
|
|
break;
|
|
}
|
|
return isInvariant;
|
|
}
|
|
case AffineExprKind::Add:
|
|
case AffineExprKind::Mul: {
|
|
auto binOp = a.cast<AffineBinaryOpExpr>();
|
|
return isInvariantAffine(binOp.getLHS(), loopStack, ldx, isAtLoop) &&
|
|
isInvariantAffine(binOp.getRHS(), loopStack, ldx, isAtLoop);
|
|
}
|
|
default: {
|
|
assert(a.isa<AffineConstantExpr>());
|
|
return true;
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Determines if affine expression is invariant.
|
|
static bool isInvariantAffine(CodegenEnv &env, AffineExpr a, LoopId ldx,
|
|
bool &isAtLoop) {
|
|
return isInvariantAffine(a, env.getCurrentLoopStack(), ldx, isAtLoop);
|
|
}
|
|
|
|
/// Helper method to construct a permuted dimension ordering
|
|
/// that adheres to the given topological sort.
|
|
//
|
|
// FIXME: does the above actually mean "dimensions", or should it say
|
|
// "level ordering"? The same dim/lvl confusion applies to all the code
|
|
// and comments in the definition below.
|
|
static AffineMap permute(CodegenEnv &env, AffineMap m) {
|
|
assert(m.getNumDims() + env.merger().getNumFilterLoops() ==
|
|
env.topSortSize() &&
|
|
"size mismatch");
|
|
// Construct the inverse of `m`; to avoid the asymptotic complexity
|
|
// of calling `m.getPermutedPosition` repeatedly.
|
|
//
|
|
// The variable `perm` must use `unsigned` rather than `Dimension`/`Level`,
|
|
// because that's what `AffineMap::getPermutationMap` requires.
|
|
// TODO: however, `perm` should be renamed to make clear what exactly
|
|
// it's storing a permutation of.
|
|
SmallVector<unsigned> perm;
|
|
const unsigned numResults = m.getNumResults();
|
|
BitVector worklist(numResults, true);
|
|
LoopOrd loopDepth = 1;
|
|
|
|
// Construct the permutation.
|
|
while (worklist.any() && loopDepth <= env.topSortSize()) {
|
|
const unsigned preSize = perm.size();
|
|
for (unsigned dim : worklist.set_bits()) {
|
|
bool isAtLoop = false;
|
|
if (m.getResult(dim).isa<AffineConstantExpr>() ||
|
|
(isInvariantAffine(m.getResult(dim), env.getLoopStackUpTo(loopDepth),
|
|
env.topSortAt(loopDepth - 1), isAtLoop) &&
|
|
isAtLoop)) {
|
|
// If the matching affine is constant expression or just become
|
|
// invariant. We can visit the dimension now without breaking the
|
|
// topSort constraint.
|
|
perm.push_back(dim);
|
|
}
|
|
}
|
|
|
|
// Removes resolved dimension.
|
|
for (unsigned i = preSize, e = perm.size(); i < e; i++)
|
|
worklist.reset(perm[i]);
|
|
|
|
// Try entering the next loop in the stack.
|
|
loopDepth++;
|
|
}
|
|
|
|
assert(perm.size() == numResults);
|
|
return AffineMap::getPermutationMap(perm, env.op().getContext());
|
|
}
|
|
|
|
/// 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.
|
|
/// filterIdx stores the current filter loop idx should be used for the next
|
|
/// compound affine sparse level, and it will be incremented by one when
|
|
/// used.
|
|
static bool findAffine(Merger &merger, TensorId tid, Level lvl, AffineExpr a,
|
|
DimLevelType dlt, LoopId &filterLdx,
|
|
bool setLvlFormat = true) {
|
|
switch (a.getKind()) {
|
|
case AffineExprKind::DimId: {
|
|
const LoopId idx = a.cast<AffineDimExpr>().getPosition();
|
|
if (!isUndefDLT(merger.getDimLevelType(tid, idx)))
|
|
return false; // used more than once
|
|
|
|
if (setLvlFormat)
|
|
merger.setLevelAndType(tid, idx, lvl, dlt);
|
|
return true;
|
|
}
|
|
case AffineExprKind::Add:
|
|
case AffineExprKind::Mul:
|
|
case AffineExprKind::Constant: {
|
|
if (!isDenseDLT(dlt) && setLvlFormat) {
|
|
assert(isUndefDLT(merger.getDimLevelType(tid, filterLdx)));
|
|
// Use a filter loop for sparse affine expression.
|
|
merger.setLevelAndType(tid, filterLdx++, lvl, dlt);
|
|
}
|
|
|
|
if (auto binOp = a.dyn_cast<AffineBinaryOpExpr>()) {
|
|
// We do not set dim level format for affine expresssion 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(), dlt, filterLdx,
|
|
false) &&
|
|
findAffine(merger, tid, lvl, binOp.getRHS(), dlt, filterLdx,
|
|
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, DimLevelType dlt,
|
|
bool isSubExp = false) {
|
|
switch (a.getKind()) {
|
|
case AffineExprKind::DimId: {
|
|
LoopId ldx = a.cast<AffineDimExpr>().getPosition();
|
|
if (!isUndefDLT(merger.getDimLevelType(tensor, ldx)))
|
|
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)
|
|
merger.setLevelAndType(tensor, ldx, lvl, dlt);
|
|
|
|
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(ldx, 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
|
|
// mulitple 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(ldx, tensor, lvl);
|
|
}
|
|
return true;
|
|
}
|
|
case AffineExprKind::Constant:
|
|
case AffineExprKind::Mul:
|
|
// TODO: Support Mul and Constant AffineExp for slice-based codegen
|
|
return false;
|
|
case AffineExprKind::Add: {
|
|
auto binOp = a.cast<AffineBinaryOpExpr>();
|
|
return findDepIdxSet(merger, tensor, lvl, binOp.getLHS(), dlt, true) &&
|
|
findDepIdxSet(merger, tensor, lvl, binOp.getRHS(), dlt, true);
|
|
}
|
|
default:
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/// Get the total number of compound affine expressions in the
|
|
/// `getMatchingIndexingMap` for the given tensor. For the following inputs:
|
|
///
|
|
/// map = (d0, d1, d2) => (d0 + d1, d2)
|
|
/// lvlTypes = ["compressed", "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 guaranted 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 = tensor.getType().dyn_cast<RankedTensorType>();
|
|
if (!rtp)
|
|
return 0;
|
|
const SparseTensorType stt(rtp);
|
|
|
|
// FIXME: There's some dim/lvl confusion here. The previous version of
|
|
// the code asserted that there are `lvlRank`-many expressions, but then
|
|
// the `exprs[d]` expression assumes there are in fact `dimRank`-many
|
|
// expressions. Even though `ArrayRef::operator[]` will check for OOB,
|
|
// the mismatch between the assertion and the usage belies that this code
|
|
// cannot support non-permutations.
|
|
//
|
|
// Elsewhere in this file the maps returned by
|
|
// `linalg::GenericOp::getMatchingIndexingMap` are inconsistent about
|
|
// whether they're expected to have `lvlRank`-many or `dimRank`-many
|
|
// expressions (cf., `genSubscript` vs `findSparseAnnotations`);
|
|
// so those are no help in determining which is actually intended.
|
|
//
|
|
// For now we work around this problem by asserting the two ranks agree.
|
|
const Dimension dimRank = stt.getDimRank();
|
|
const Level lvlRank = stt.getLvlRank();
|
|
assert(dimRank == lvlRank && "Non-permutations not currently supported");
|
|
const auto exprs = map.getResults();
|
|
assert(static_cast<Dimension>(exprs.size()) == dimRank &&
|
|
"AffineMap does not have dimension-rank many results");
|
|
(void)dimRank;
|
|
unsigned num = 0;
|
|
for (Level l = 0; l < lvlRank; l++) {
|
|
// FIXME: `toOrigDim` is deprecated.
|
|
const Dimension d = toOrigDim(stt.getEncoding(), l);
|
|
if (!exprs[d].isa<AffineDimExpr>() && !stt.isDenseLvl(l))
|
|
num++;
|
|
}
|
|
return num;
|
|
}
|
|
|
|
/// Get 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;
|
|
}
|
|
|
|
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 filter-loop-based approach, it accept (almost) arbitrary affine
|
|
/// index expression on sparse tensor but it is much less efficient, and will be
|
|
/// gradually removed from the codebase.
|
|
/// 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;
|
|
// `filterLdx` may be mutated by `findAffine`.
|
|
LoopId filterLdx = env.merger().getStartingFilterLoopId();
|
|
for (OpOperand &t : env.op()->getOpOperands()) {
|
|
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 TensorId tid = t.getOperandNumber();
|
|
AffineExpr a = map.getResult(toOrigDim(enc, l));
|
|
DimLevelType dlt = enc.getLvlType(l);
|
|
if (idxReducBased && needIdxReduc) {
|
|
if (!findDepIdxSet(env.merger(), tid, l, a, dlt))
|
|
return false; // inadmissible affine expression
|
|
} else {
|
|
if (!findAffine(env.merger(), tid, l, a, dlt, filterLdx))
|
|
return false; // inadmissible affine expression
|
|
}
|
|
}
|
|
}
|
|
assert(filterLdx == env.merger().getNumLoops());
|
|
return annotated;
|
|
}
|
|
|
|
/// A helper to compute a topological sort. O(n^2) time complexity
|
|
/// as we use adj matrix for the graph.
|
|
/// The sorted result will put the first Reduction iterator to the
|
|
/// latest possible index.
|
|
/// FIXME(wrengr): correct the above "index"
|
|
///
|
|
/// The `inDegree` is indexed by `LoopId`, and the `adjM` is indexed by
|
|
/// `(LoopId,LoopId)`.
|
|
static bool topSortOptimal(CodegenEnv &env, LoopId n,
|
|
ArrayRef<utils::IteratorType> iteratorTypes,
|
|
std::vector<unsigned> &inDegree,
|
|
std::vector<std::vector<bool>> &adjM) {
|
|
std::vector<LoopId> redIt; // reduce iterator with 0 degree
|
|
std::vector<LoopId> parIt; // parallel iterator with 0 degree
|
|
std::vector<LoopId> filterIt; // filter loop with 0 degree
|
|
for (LoopId i = 0; i < n; i++) {
|
|
if (inDegree[i] == 0) {
|
|
if (env.merger().isFilterLoop(i))
|
|
filterIt.push_back(i);
|
|
else if (linalg::isReductionIterator(iteratorTypes[i]))
|
|
redIt.push_back(i);
|
|
else
|
|
parIt.push_back(i);
|
|
}
|
|
}
|
|
|
|
while (!redIt.empty() || !parIt.empty() || !filterIt.empty()) {
|
|
// We always choose in order of filter loop -> parallel loop -> reduction
|
|
// loop because
|
|
// 1. Putting reduction loop early might make the loop sequence
|
|
// inadmissible.
|
|
// 2. Filter loops should be put as early as possible for better
|
|
// performance, since only one (if any) iteration will carry the
|
|
// computation. E.g., for (1 to N)
|
|
// for (1 to M)
|
|
// for (1 to K)
|
|
// if (xxx)
|
|
// O(X) computation => O(NMK+NMX) time complexity
|
|
//
|
|
// By putting the filter loop one level up, we got
|
|
//
|
|
// for (1 to N)
|
|
// for (1 to K)
|
|
// if (xxx)
|
|
// for (1 to M)
|
|
// O(X) computation => O(NK+NMX) time complexity
|
|
auto &it = !filterIt.empty() ? filterIt : (!parIt.empty() ? parIt : redIt);
|
|
auto src = it.back();
|
|
env.topSortPushBack(src);
|
|
it.pop_back();
|
|
// Update in-degree, and push 0-degree node into worklist.
|
|
for (LoopId dst = 0; dst < n; dst++) {
|
|
if (adjM[src][dst] && --inDegree[dst] == 0) {
|
|
if (env.merger().isFilterLoop(dst))
|
|
filterIt.push_back(dst);
|
|
else if (linalg::isReductionIterator(iteratorTypes[dst]))
|
|
redIt.push_back(dst);
|
|
else
|
|
parIt.push_back(dst);
|
|
}
|
|
}
|
|
}
|
|
return env.topSortSize() == n;
|
|
}
|
|
|
|
/// Helper method to add all constraints from the indices in one affine
|
|
/// expression before all indices in the other affine expression. For
|
|
/// example i0+i1 < i2+i3+1 yields i0<i2, i0<i3, i1<i2, and i1<i3.
|
|
/// The affine expression `a` is empty iff `fidx` have a value, leading to
|
|
/// b = (i0 + i1) < fidx => i0 < fidx, i1 < fidx.
|
|
/// The affine expression `b` is empty iff `tidx` have a value, leading to
|
|
/// tidx < a = (i0 + i1) => tidx < i0, tidx < i1.
|
|
///
|
|
/// The `inDegree` is indexed by `LoopId`, and the `adjM` is indexed by
|
|
/// `(LoopId,LoopId)`.
|
|
static void addAffineOrderings(std::vector<std::vector<bool>> &adjM,
|
|
std::vector<unsigned> &inDegree, AffineExpr a,
|
|
AffineExpr b, std::optional<LoopId> fidx,
|
|
std::optional<LoopId> tidx) {
|
|
if (!a && !b) {
|
|
// Recursion leaf.
|
|
assert(fidx && tidx);
|
|
const LoopId f = *fidx, t = *tidx;
|
|
if (!adjM[f][t]) {
|
|
adjM[f][t] = true;
|
|
inDegree[t]++;
|
|
}
|
|
return;
|
|
}
|
|
// Picks an affine expression and expand (recurse into) it.
|
|
const auto toExpand = a ? a : b;
|
|
switch (toExpand.getKind()) {
|
|
case AffineExprKind::DimId: {
|
|
std::optional<LoopId> idx = toExpand.cast<AffineDimExpr>().getPosition();
|
|
if (toExpand == a)
|
|
addAffineOrderings(adjM, inDegree, AffineExpr(), b, idx, tidx);
|
|
else // toExpand == b
|
|
addAffineOrderings(adjM, inDegree, a, AffineExpr(), fidx, idx);
|
|
break;
|
|
}
|
|
case AffineExprKind::Add:
|
|
case AffineExprKind::Mul: {
|
|
auto binOp = toExpand.cast<AffineBinaryOpExpr>();
|
|
if (toExpand == a) {
|
|
addAffineOrderings(adjM, inDegree, binOp.getLHS(), b, fidx, tidx);
|
|
addAffineOrderings(adjM, inDegree, binOp.getRHS(), b, fidx, tidx);
|
|
} else {
|
|
addAffineOrderings(adjM, inDegree, a, binOp.getLHS(), fidx, tidx);
|
|
addAffineOrderings(adjM, inDegree, a, binOp.getRHS(), fidx, tidx);
|
|
}
|
|
break;
|
|
}
|
|
default:
|
|
break;
|
|
}
|
|
}
|
|
|
|
static void tryRelaxAffineConstraints(linalg::GenericOp op,
|
|
std::optional<LoopId> &fldx,
|
|
AffineExpr &fa,
|
|
std::optional<LoopId> &tldx,
|
|
AffineExpr &ta) {
|
|
// We use a heuristic here to only pick one dim expression from each
|
|
// compound affine expression to establish the order between two dense
|
|
// dimensions.
|
|
if (!tldx) {
|
|
AffineDimFinder finder(op);
|
|
// NOTE: The ordering can only be loosen when the destination level is
|
|
// dense (when !tldx), for [dense, sparse] -> (d0 + d1, d2), we still
|
|
// require both d0 < d2 and d1 < d2 to ensure correct ordering (i.e.,
|
|
// no ordering like d0->d2->d1).
|
|
// TODO: this is obviously a sub optimal solution.
|
|
if (!fldx && !fa.isa<AffineConstantExpr>()) {
|
|
// Heuristic: we prefer parallel loop for lhs to reduce the chance
|
|
// we add reduce < parallel ordering.
|
|
finder.setPickedIterType(utils::IteratorType::parallel);
|
|
finder.walkPostOrder(fa);
|
|
fa = finder.getDimExpr();
|
|
fldx = finder.getDimExpr().getPosition();
|
|
}
|
|
if (!ta.isa<AffineConstantExpr>()) {
|
|
// Heuristic: we prefer reduction loop for rhs to reduce the chance
|
|
// adding reduce < parallel ordering.
|
|
finder.setPickedIterType(utils::IteratorType::reduction);
|
|
finder.walkPostOrder(ta);
|
|
ta = finder.getDimExpr();
|
|
tldx = finder.getDimExpr().getPosition();
|
|
}
|
|
}
|
|
}
|
|
|
|
static void addFilterLoopBasedConstraints(CodegenEnv &env, OpOperand &t,
|
|
OpOperand *skip, SortMask mask,
|
|
std::vector<std::vector<bool>> &adjM,
|
|
std::vector<unsigned> &inDegree) {
|
|
// Get map and encoding.
|
|
auto map = env.op().getMatchingIndexingMap(&t);
|
|
auto enc = getSparseTensorEncoding(t.get().getType());
|
|
|
|
// Each tensor expression and optional dimension ordering (row-major
|
|
// by default) puts an ordering constraint on the loop indices. For
|
|
// example, the tensor expresion A_ijk forces the ordering i < j < k
|
|
// on the loop indices if no explicit dimension ordering is given.
|
|
for (Level l = 0, rank = map.getNumResults(); l < rank; l++) {
|
|
AffineExpr ta = map.getResult(toOrigDim(enc, l));
|
|
std::optional<LoopId> tldx =
|
|
env.merger().getLoopId(t.getOperandNumber(), l);
|
|
// Filter loops should be constructed after all the dependent loops,
|
|
// i.e., d0 + d1 < filter_loop(d0 + d1)
|
|
if (tldx && env.merger().isFilterLoop(*tldx)) {
|
|
assert(!ta.isa<AffineDimExpr>() && !isDenseDLT(enc.getDimLevelType()[l]));
|
|
addAffineOrderings(adjM, inDegree, ta, AffineExpr(), std::nullopt, tldx);
|
|
// Now that the ordering of affine expression is captured by filter
|
|
// loop idx, we only need to ensure the affine ordering against filter
|
|
// loop. Thus, we reset the affine express to nil here to mark it as
|
|
// resolved.
|
|
ta = AffineExpr();
|
|
}
|
|
|
|
// Skip tensor during cycle resolution, though order between filter loop
|
|
// and dependent loops need to be guaranteed unconditionally.
|
|
if (&t == skip)
|
|
continue;
|
|
|
|
if (l > 0) {
|
|
AffineExpr fa = map.getResult(toOrigDim(enc, l - 1));
|
|
std::optional<LoopId> fldx =
|
|
env.merger().getLoopId(t.getOperandNumber(), l - 1);
|
|
|
|
// Applying order constraints on every pair of dimExpr between two
|
|
// compound affine expressions can sometime too strict:
|
|
// E.g, for [dense, dense] -> (d0 + d1, d2 + d3).
|
|
// It is totally fine to have loop sequence d0->d2->d1->d3 instead of
|
|
// requiring d0 < d2, d1 < d2, d0 < d3, d1 < d3.
|
|
// We also relax the affine constraint when use slice-based algorithm
|
|
// as there is no filter loop for affine index on sparse dimension.
|
|
// TODO: do we really need the condition?
|
|
if (!includesDense(mask))
|
|
tryRelaxAffineConstraints(env.op(), fldx, fa, tldx, ta);
|
|
|
|
// (d0 + d1) < (d2 + d3), or
|
|
// filter_loop_d-1 < (d2 + d3), or
|
|
// (d0 + d1) < filter_loop_d, or
|
|
// filter_loop_d-1 < filter_loop_d depending on whether fa/ta is reset
|
|
// above.
|
|
addAffineOrderings(adjM, inDegree, fa, ta, fldx, tldx);
|
|
}
|
|
}
|
|
}
|
|
|
|
static void addSliceBasedConstraints(CodegenEnv &env, OpOperand &t,
|
|
OpOperand *skip, SortMask mask,
|
|
std::vector<std::vector<bool>> &adjM,
|
|
std::vector<unsigned> &inDegree) {
|
|
// Get map and encoding.
|
|
auto map = env.op().getMatchingIndexingMap(&t);
|
|
auto enc = getSparseTensorEncoding(t.get().getType());
|
|
|
|
// No special treatment for simple indices.
|
|
if (getNumNonTrivialIdxExpOnSparseLvls(map, t.get()) == 0)
|
|
return addFilterLoopBasedConstraints(env, t, skip, mask, adjM, inDegree);
|
|
|
|
// Skip tensor during cycle resolution, though order between filter loop
|
|
// and dependent loops need to be guaranteed unconditionally.
|
|
if (&t == skip)
|
|
return;
|
|
|
|
AffineDimFinder finder(env.op());
|
|
finder.setPickedIterType(utils::IteratorType::reduction);
|
|
// To compute iteration graph for tensor[d0 + d1 + d3, d4 + d5 + d6],
|
|
// we requires there exist d_x \in {d0, d1, d3} and d_y \in {d4, d5, d6},
|
|
// and d_x > d_y && {d0, d1, d3} - d_x > {d4, d5, d6} - d_y
|
|
for (Level lvl = 1, rank = map.getNumResults(); lvl < rank; lvl++) {
|
|
AffineExpr fa = map.getResult(toOrigDim(enc, lvl - 1));
|
|
AffineExpr ta = map.getResult(toOrigDim(enc, lvl));
|
|
|
|
// This is a heuristic, we pick an abitrary reduction loop from lhs and
|
|
// rhs and use them as d_x and d_y.
|
|
finder.walkPostOrder(fa);
|
|
AffineDimExpr fexp = finder.getDimExpr();
|
|
LoopId fldx = fexp.getPosition();
|
|
|
|
finder.walkPostOrder(ta);
|
|
AffineDimExpr texp = finder.getDimExpr();
|
|
LoopId tldx = texp.getPosition();
|
|
|
|
// d_x > d_y
|
|
if (!adjM[fldx][tldx]) {
|
|
adjM[fldx][tldx] = true;
|
|
inDegree[tldx]++;
|
|
}
|
|
|
|
AffineDimCollector fCollector;
|
|
fCollector.walkPostOrder(fa);
|
|
AffineDimCollector tCollector;
|
|
tCollector.walkPostOrder(ta);
|
|
|
|
// make sure dx and dy is the last;
|
|
for (auto fd : fCollector.dims) {
|
|
LoopId f = fd.getPosition();
|
|
if (f == fldx)
|
|
continue;
|
|
if (!adjM[f][fldx]) {
|
|
adjM[f][fldx] = true;
|
|
inDegree[fldx]++;
|
|
}
|
|
}
|
|
for (auto td : tCollector.dims) {
|
|
LoopId t = td.getPosition();
|
|
if (t == tldx)
|
|
continue;
|
|
if (!adjM[t][tldx]) {
|
|
adjM[t][tldx] = true;
|
|
inDegree[tldx]++;
|
|
}
|
|
}
|
|
// Since we only support affine addition, the order between two dim
|
|
// expression does not really matters.
|
|
// {d0, d1, d3} - d_x > {d4, d5, d6} - d_y
|
|
// This is to ensure that the affine expressions are reduced in sparse
|
|
// tensor level ordering.
|
|
// TODO: this ordering could probably be loosen if we support out-of-order
|
|
// reduction.
|
|
// TODO: the evaluation order need to be ensure to
|
|
// support affine multiplication.
|
|
for (auto fd : fCollector.dims) {
|
|
LoopId f = fd.getPosition();
|
|
if (f == fldx) // skip d_x
|
|
continue;
|
|
|
|
for (auto td : tCollector.dims) {
|
|
LoopId t = td.getPosition();
|
|
if (t == tldx) // skip d_y
|
|
continue;
|
|
if (!adjM[f][t]) {
|
|
adjM[f][t] = true;
|
|
inDegree[t]++;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Computes a topologically sorted iteration graph for the linalg operation.
|
|
/// Ensures all tensors are visited in natural index order. This is
|
|
/// essential for sparse storage formats since these only support access
|
|
/// along fixed dimensions. Even for dense storage formats, however, the natural
|
|
/// index order yields innermost unit-stride access with better spatial
|
|
/// locality.
|
|
static bool computeIterationGraph(CodegenEnv &env, SortMask mask,
|
|
OpOperand *skip, bool idxReducBased = false) {
|
|
// Set up an n x n from/to adjacency matrix of the iteration graph
|
|
// for the implicit loop indices i_0 .. i_n-1.
|
|
const LoopId n = env.merger().getNumLoops();
|
|
std::vector<std::vector<bool>> adjM(n, std::vector<bool>(n, false));
|
|
std::vector<unsigned> inDegree(n, 0); // in-degree of each node.
|
|
const auto iteratorTypes = env.op().getIteratorTypesArray();
|
|
// Iterate over the indexing maps of every tensor in the tensor expression.
|
|
for (OpOperand &t : env.op()->getOpOperands()) {
|
|
// Get map and encoding.
|
|
const auto map = env.op().getMatchingIndexingMap(&t);
|
|
const auto enc = getSparseTensorEncoding(t.get().getType());
|
|
assert(map.getNumDims() + getNumNonTrivialIdxExpOnSparseLvls(env.op()) ==
|
|
n);
|
|
|
|
// Skips dense inputs/outputs when not requested.
|
|
const bool isDenseInput = !enc && env.op().isDpsInput(&t);
|
|
const bool isDenseOutput = !enc && !isDenseInput;
|
|
if ((isDenseInput && !includesDenseInput(mask)) ||
|
|
(isDenseOutput && !includesDenseOutput(mask)))
|
|
continue;
|
|
|
|
// Push unrelated loops into sparse iteration space, so these
|
|
// will be skipped more often.
|
|
// TODO: Do we really need this?
|
|
if (includesUndef(mask)) {
|
|
const TensorId tensor = t.getOperandNumber();
|
|
for (LoopId i = 0; i < n; i++) {
|
|
const auto dltI = env.dlt(tensor, i);
|
|
if (isCompressedDLT(dltI) || isSingletonDLT(dltI)) {
|
|
for (LoopId j = 0; j < n; j++)
|
|
if (isUndefDLT(env.dlt(tensor, j))) {
|
|
adjM[i][j] = true;
|
|
inDegree[j]++;
|
|
}
|
|
} else {
|
|
assert(isDenseDLT(dltI) || isUndefDLT(dltI));
|
|
}
|
|
}
|
|
}
|
|
// Push unrelated loops into sparse iteration space, so these
|
|
// will be skipped more often.
|
|
if (idxReducBased)
|
|
addSliceBasedConstraints(env, t, skip, mask, adjM, inDegree);
|
|
else
|
|
addFilterLoopBasedConstraints(env, t, skip, mask, adjM, inDegree);
|
|
}
|
|
// Topologically sort the iteration graph to determine loop order.
|
|
// Report failure for a cyclic iteration graph.
|
|
env.topSortClear(n);
|
|
return topSortOptimal(env, n, iteratorTypes, inDegree, adjM);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Sparse compiler 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);
|
|
|
|
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;
|
|
});
|
|
}
|
|
|
|
/// Generates index for load/store on sparse tensor.
|
|
// FIXME: It's not entirely clear what "index" means here (i.e., is it
|
|
// a "coordinate", or "Ldx", or what). So the function should be renamed
|
|
// and/or the documentation expanded in order to clarify.
|
|
static Value genIndex(CodegenEnv &env, OpOperand *t) {
|
|
auto map = env.op().getMatchingIndexingMap(t);
|
|
const auto stt = getSparseTensorType(t->get());
|
|
const Level lvlRank = stt.getLvlRank();
|
|
assert(static_cast<Level>(map.getNumResults()) == lvlRank);
|
|
// FIXME: `toOrigDim` is deprecated.
|
|
// FIXME: above we asserted that there are `lvlRank` many results,
|
|
// but this is assuming there are in fact `dimRank` many results instead.
|
|
AffineExpr a = map.getResult(toOrigDim(stt.getEncoding(), lvlRank - 1));
|
|
assert(a.getKind() == AffineExprKind::DimId);
|
|
const LoopId idx = a.cast<AffineDimExpr>().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 = 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 LoopOrd numLoops = op.getRank(t);
|
|
// TODO: rewrite this to use `env.emitter().getLoopIVs(ivs)`
|
|
// instead. We just need to either assert that `numLoops ==
|
|
// env.emitter().getCurrentDepth()`, or else update the `getLoopIVs`
|
|
// method to take an optional parameter to restrict to a smaller depth.
|
|
SmallVector<Value> ivs;
|
|
ivs.reserve(numLoops);
|
|
for (LoopOrd n = 0; n < numLoops; n++) {
|
|
const auto iv = env.emitter().getLoopIV(n);
|
|
assert(iv);
|
|
ivs.push_back(iv);
|
|
}
|
|
Value chain = env.getInsertionChain();
|
|
if (!env.getValidLexInsert()) {
|
|
env.updateInsertionChain(builder.create<InsertOp>(loc, rhs, chain, ivs));
|
|
} else {
|
|
// 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));
|
|
}
|
|
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) {
|
|
linalg::GenericOp op = env.op();
|
|
Location loc = op.getLoc();
|
|
// Test if this is a scalarized reduction.
|
|
if (env.isReduc()) {
|
|
env.updateReduc(rhs);
|
|
return;
|
|
}
|
|
// Store during insertion.
|
|
OpOperand *t = op.getDpsInitOperand(0);
|
|
if (env.isSparseOutput(t)) {
|
|
if (!rhs) {
|
|
// Only unary and binary are allowed to return uninitialized rhs
|
|
// to indicate missing output.
|
|
assert(env.exp(exp).kind == kUnary || env.exp(exp).kind == kBinary);
|
|
} else if (env.exp(exp).kind == kSelect) {
|
|
// 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.exp(exp).val = Value();
|
|
// 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);
|
|
} else {
|
|
genInsertionStore(env, builder, t, rhs);
|
|
}
|
|
return;
|
|
}
|
|
// Actual store.
|
|
SmallVector<Value> args;
|
|
Value ptr = genSubscript(env, builder, t, args);
|
|
builder.create<memref::StoreOp>(loc, rhs, ptr, args);
|
|
}
|
|
|
|
/// 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 sparse compiler. 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, LoopId ldx) {
|
|
if (Operation *def = e.getDefiningOp()) {
|
|
if (auto indexOp = dyn_cast<linalg::IndexOp>(def))
|
|
return env.getLoopVar(indexOp.getDim());
|
|
if (def->getBlock() == block) {
|
|
for (unsigned i = 0, n = def->getNumOperands(); i < n; i++) {
|
|
rewriter.updateRootInPlace(def, [&]() {
|
|
def->setOperand(
|
|
i, relinkBranch(env, rewriter, block, def->getOperand(i), ldx));
|
|
});
|
|
}
|
|
}
|
|
}
|
|
return e;
|
|
}
|
|
|
|
/// Recursively generates tensor expression.
|
|
static Value genExp(CodegenEnv &env, RewriterBase &rewriter, ExprId e,
|
|
LoopId ldx) {
|
|
linalg::GenericOp op = env.op();
|
|
Location loc = op.getLoc();
|
|
|
|
if (e == kInvalidId)
|
|
return Value();
|
|
const TensorExp &exp = env.exp(e);
|
|
const auto kind = exp.kind;
|
|
if (kind == Kind::kTensor)
|
|
return genTensorLoad(env, rewriter, e);
|
|
if (kind == Kind::kInvariant)
|
|
return genInvariantValue(env, e);
|
|
if (kind == Kind::kLoopVar)
|
|
return env.getLoopVar(exp.loop);
|
|
|
|
if (kind == Kind::kReduce)
|
|
env.startCustomReduc(e); // enter custom
|
|
|
|
Value v0 = genExp(env, rewriter, exp.children.e0, ldx);
|
|
Value v1 = genExp(env, rewriter, exp.children.e1, ldx);
|
|
Value ee = env.merger().buildExp(rewriter, loc, e, v0, v1);
|
|
if (ee && (kind == Kind::kUnary || kind == Kind::kBinary ||
|
|
kind == Kind::kBinaryBranch || kind == Kind::kReduce ||
|
|
kind == Kind::kSelect))
|
|
ee = relinkBranch(env, rewriter, ee.getParentBlock(), ee, ldx);
|
|
|
|
if (kind == Kind::kReduce)
|
|
env.endCustomReduc(); // exit custom
|
|
|
|
if (kind == kSelect) {
|
|
assert(!exp.val);
|
|
env.exp(e).val = 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 ldx, bool atStart) {
|
|
if (exp == kInvalidId)
|
|
return;
|
|
if (env.exp(exp).kind == Kind::kTensor) {
|
|
// Inspect tensor indices.
|
|
bool isAtLoop = ldx == kInvalidId;
|
|
linalg::GenericOp op = env.op();
|
|
OpOperand &t = op->getOpOperand(env.exp(exp).tensor);
|
|
auto map = op.getMatchingIndexingMap(&t);
|
|
const auto stt = getSparseTensorType(t.get());
|
|
const Level lvlRank = stt.getLvlRank();
|
|
assert(static_cast<Level>(map.getNumResults()) == lvlRank);
|
|
for (Level l = 0; l < lvlRank; l++) {
|
|
// FIXME: `toOrigDim` is deprecated.
|
|
// FIXME: above we asserted that there are `lvlRank` many results,
|
|
// but this is assuming there are in fact `dimRank` many results instead.
|
|
AffineExpr a = map.getResult(toOrigDim(stt.getEncoding(), l));
|
|
const auto sldx = env.merger().getLoopId(t.getOperandNumber(), l);
|
|
if (sldx && env.merger().isFilterLoop(*sldx)) {
|
|
if (!env.getLoopVar(*sldx))
|
|
// The filter loops has not been constructed.
|
|
return;
|
|
if (*sldx == ldx)
|
|
isAtLoop = true;
|
|
} else if (!isInvariantAffine(env, a, ldx, isAtLoop))
|
|
return; // still in play
|
|
}
|
|
// All exhausted at this level (isAtLoop denotes exactly at this LoopId).
|
|
if (!isAtLoop)
|
|
return;
|
|
OpOperand *lhs = op.getDpsInitOperand(0);
|
|
if (lhs == &t) {
|
|
// Start or end a scalarized reduction.
|
|
if (atStart) {
|
|
Value load = env.isCustomReduc() ? env.getCustomRedId()
|
|
: genTensorLoad(env, builder, exp);
|
|
env.startReduc(exp, load);
|
|
if (env.hasSparseOutput())
|
|
env.setValidLexInsert(constantI1(builder, env.op().getLoc(), false));
|
|
} else {
|
|
genTensorStore(env, builder, exp, env.endReduc());
|
|
env.clearValidLexInsert();
|
|
}
|
|
} else {
|
|
// Start or end loop invariant hoisting of a tensor load.
|
|
env.exp(exp).val = atStart ? genTensorLoad(env, builder, exp) : Value();
|
|
}
|
|
} else if (env.exp(exp).kind != Kind::kInvariant &&
|
|
env.exp(exp).kind != Kind::kLoopVar) {
|
|
// 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 == 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, ldx, atStart);
|
|
genInvariants(env, builder, e1, ldx, atStart);
|
|
if (env.exp(exp).kind == Kind::kReduce)
|
|
env.endCustomReduc(); // exit custom
|
|
}
|
|
}
|
|
|
|
/// Generates an expanded access pattern in innermost dimension.
|
|
static void genExpand(CodegenEnv &env, OpBuilder &builder, LoopOrd at,
|
|
bool atStart) {
|
|
linalg::GenericOp op = env.op();
|
|
OpOperand *lhs = op.getDpsInitOperand(0);
|
|
if (!env.atExpandLevel(lhs, op.getRank(lhs), at))
|
|
return; // not needed at this level
|
|
assert(!env.isReduc());
|
|
// Generate start or end of an expanded access pattern. Note that because
|
|
// an expension 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 (atStart) {
|
|
auto dynShape = {ShapedType::kDynamic};
|
|
Type etp = tensor.getType().cast<ShapedType>().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 (LoopOrd i = 0; i < at; 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");
|
|
}
|
|
|
|
/// Generates a for-loop on a single index.
|
|
static Operation *genFor(CodegenEnv &env, OpBuilder &builder, bool isOuter,
|
|
bool isInner, LoopId ldx, ArrayRef<TensorId> tids,
|
|
ArrayRef<Level> lvls) {
|
|
linalg::GenericOp op = env.op();
|
|
Location loc = op.getLoc();
|
|
auto iteratorTypes = op.getIteratorTypesArray();
|
|
bool isSparse = llvm::any_of(tids, [ldx, &env](TensorId tid) {
|
|
const auto dlt = env.dlt(tid, ldx);
|
|
return isCompressedDLT(dlt) || isSingletonDLT(dlt);
|
|
});
|
|
|
|
bool isParallel = isParallelFor(env, isOuter, isSparse);
|
|
|
|
Operation *loop = *env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
|
|
if (env.merger().isFilterLoop(ldx)) {
|
|
const TensorId tid = tids.front();
|
|
const Level lvl = lvls.front();
|
|
// tids/lvls must only have one value because filter loops only
|
|
// corresponding to the one and only sparse tensor level.
|
|
assert(isSparse && tids.size() == 1 && lvls.size() == 1);
|
|
OpOperand *t = &op->getOpOperand(tid);
|
|
auto enc = getSparseTensorEncoding(t->get().getType());
|
|
// Retrieves the affine expression for the filter loop.
|
|
// FIXME: `toOrigDim` is deprecated.
|
|
AffineExpr a =
|
|
op.getMatchingIndexingMap(t).getResult(toOrigDim(enc, lvl));
|
|
return env.emitter().enterFilterLoopOverTensorAtLvl(builder, loc, tid,
|
|
lvl, a, reduc);
|
|
}
|
|
return env.emitter().enterLoopOverTensorAtLvl(builder, loc, tids, lvls,
|
|
reduc, isParallel);
|
|
});
|
|
assert(loop);
|
|
return loop;
|
|
}
|
|
|
|
/// Emit a while-loop for co-iteration over multiple indices.
|
|
static Operation *genWhile(CodegenEnv &env, OpBuilder &builder, LoopId idx,
|
|
bool needsUniv, ArrayRef<TensorId> tids,
|
|
ArrayRef<Level> lvls) {
|
|
Operation *loop = *env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
|
|
// Construct the while-loop with a parameter for each
|
|
// index.
|
|
return env.emitter().enterCoIterationOverTensorsAtLvls(
|
|
builder, env.op().getLoc(), tids, lvls, needsUniv, reduc);
|
|
});
|
|
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, LoopOrd at,
|
|
bool needsUniv, ArrayRef<TensorId> tids,
|
|
ArrayRef<Level> lvls, bool isFor) {
|
|
assert(tids.size() == lvls.size());
|
|
const LoopId idx = env.topSortAt(at);
|
|
if (isFor) {
|
|
bool isOuter = at == 0;
|
|
bool isInner = at == env.topSortSize() - 1;
|
|
return genFor(env, builder, isOuter, isInner, idx, tids, lvls);
|
|
}
|
|
return genWhile(env, builder, idx, needsUniv, tids, lvls);
|
|
}
|
|
|
|
/// Generates the induction structure for a while-loop.
|
|
static void finalizeWhileOp(CodegenEnv &env, OpBuilder &builder, LoopId idx,
|
|
bool needsUniv, BitVector &induction,
|
|
scf::WhileOp whileOp) {
|
|
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.getValidLexInsert()) {
|
|
yields.push_back(env.getValidLexInsert());
|
|
env.setValidLexInsert(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);
|
|
}
|
|
}
|
|
builder.setInsertionPointToEnd(&whileOp.getAfter().front());
|
|
}
|
|
|
|
/// Generates a single if-statement within a while-loop.
|
|
static scf::IfOp genIf(CodegenEnv &env, OpBuilder &builder, LoopId ldx,
|
|
BitVector &conditions) {
|
|
Location loc = env.op().getLoc();
|
|
SmallVector<Type> types;
|
|
Value cond;
|
|
for (TensorLoopId b = 0, be = conditions.size(); b < be; b++) {
|
|
if (!conditions[b])
|
|
continue;
|
|
const TensorId tid = env.merger().tensor(b);
|
|
assert(ldx == env.merger().loop(b));
|
|
Value clause;
|
|
const auto dlt = env.dlt(b);
|
|
if (isCompressedDLT(dlt) || isSingletonDLT(dlt)) {
|
|
const Level lvl = *env.merger().getLvl(tid, ldx);
|
|
const Value crd = env.emitter().getCoords()[tid][lvl];
|
|
const Value lvar = env.getLoopVar(ldx);
|
|
clause = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, crd,
|
|
lvar);
|
|
} else {
|
|
assert(isDenseDLT(dlt) || isUndefDLT(dlt));
|
|
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.getValidLexInsert())
|
|
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,
|
|
Operation *loop, Value redInput, Value cntInput,
|
|
Value insInput) {
|
|
SmallVector<Value> operands;
|
|
if (env.isReduc()) {
|
|
operands.push_back(env.getReduc());
|
|
env.updateReduc(redInput);
|
|
if (env.getValidLexInsert())
|
|
// Any overlapping indices during a reduction creates a valid lex insert.
|
|
operands.push_back(constantI1(builder, env.op().getLoc(), true));
|
|
}
|
|
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());
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Sparse compiler synthesis methods (loop sequence).
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
/// 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,
|
|
LoopOrd at, LoopId idx, LoopId ldx, LatSetId lts) {
|
|
assert(!env.getLoopVar(idx));
|
|
// Emit invariants at this loop sequence level.
|
|
genInvariants(env, builder, exp, ldx, /*atStart=*/true);
|
|
// Emit access pattern expansion for sparse tensor output.
|
|
genExpand(env, builder, at, /*atStart=*/true);
|
|
// Emit further intitialization at this loop sequence level.
|
|
const LatPointId l0 = env.set(lts)[0];
|
|
bool needsUniv = false;
|
|
|
|
SmallVector<TensorId> tids;
|
|
SmallVector<Level> lvls;
|
|
env.merger().foreachTensorLoopId(l0, [&](TensorLoopId b, TensorId tid,
|
|
std::optional<Level> lvl,
|
|
DimLevelType dlt, bool isIdxReduc) {
|
|
assert(env.merger().loop(b) == idx);
|
|
// FIXME: Dense index reduction can reuse the universal index as well.
|
|
if (!isIdxReduc && (isDenseDLT(dlt) || isUndefDLT(dlt))) {
|
|
needsUniv = true;
|
|
} else {
|
|
// sparse/singleton levels.
|
|
tids.push_back(tid);
|
|
lvls.push_back(*lvl);
|
|
}
|
|
});
|
|
|
|
env.emitter().enterNewLoopSeq(builder, env.op().getLoc(), tids, lvls);
|
|
|
|
// Maintain the universal index only if it is actually
|
|
// consumed by a subsequent lattice point.
|
|
if (needsUniv) {
|
|
unsigned lsize = env.set(lts).size();
|
|
for (unsigned i = 1; i < lsize; i++) {
|
|
const LatPointId li = env.set(lts)[i];
|
|
if (!env.merger().hasAnySparse(env.lat(li).simple) &&
|
|
!env.merger().hasSparseIdxReduction(env.lat(li).simple))
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
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 = input->getOperandNumber();
|
|
const Level lvlRank = enc.getLvlRank();
|
|
assert(lvlExprs.size() == static_cast<size_t>(lvlRank));
|
|
// FIXME: there is dim/lvl confusion here
|
|
for (Level l = startLvl; l < lvlRank; l++) {
|
|
// FIXME: `toOrigDim` is deprecated.
|
|
AffineExpr lvlExpr = lvlExprs[toOrigDim(enc, l)];
|
|
if (enc.isDenseLvl(l) && lvlExpr.isa<AffineConstantExpr>())
|
|
env.emitter().genDenseAffineAddress(builder, loc, tid, l, lvlExpr);
|
|
else
|
|
return; // break on first non-dense non-constant level
|
|
}
|
|
}
|
|
}
|
|
|
|
static void genInitConstantDenseAddress(CodegenEnv &env,
|
|
RewriterBase &rewriter) {
|
|
// We can generate address for constant affine expression before any loops
|
|
// starting from the first level as they do not depend on any thing.
|
|
// E.g., [Dense, Dense, Sparse] -> (1, 2, d0), the addresses for the first two
|
|
// levels can be determined before loops.
|
|
for (TensorId tid = 0, e = env.op().getNumDpsInputs(); tid < e; tid++)
|
|
genConstantDenseAddressFromLevel(env, rewriter, tid, 0);
|
|
}
|
|
|
|
/// Return true if the lattices bit can be iterated by a for loop.
|
|
static bool translateBitsToTidLvlPairs(
|
|
CodegenEnv &env, LatPointId li, LoopId ldx, SmallVectorImpl<TensorId> &tids,
|
|
SmallVectorImpl<Level> &lvls, SmallVectorImpl<TensorId> &affineTids,
|
|
SmallVectorImpl<Level> &affineLvls, SmallVectorImpl<AffineExpr> &exps) {
|
|
const BitVector &simple = env.lat(li).simple;
|
|
const TensorId outTid = env.merger().getOutTensorID();
|
|
const std::optional<Level> outLvl = env.merger().getLvl(outTid, ldx);
|
|
|
|
unsigned numloopCond = 0;
|
|
bool hasNonUnique = false;
|
|
|
|
env.merger().foreachTensorLoopId(
|
|
li, [&, ldx](TensorLoopId b, TensorId tid, std::optional<Level> lvl,
|
|
DimLevelType dlt, bool isIdxReduc) {
|
|
if (simple.test(b)) {
|
|
if (isIdxReduc) {
|
|
tids.push_back(tid);
|
|
lvls.push_back(*lvl);
|
|
numloopCond++;
|
|
return;
|
|
}
|
|
if (isUndefDLT(dlt)) {
|
|
// An undefined dlt in the lattices, we probably mean to
|
|
// iterate based on the level of output tensor. E.g., this
|
|
// could be a synthetic tensor (for invariants and sparse
|
|
// output tensor).
|
|
// out[i][j] = invariant; or a broadcast
|
|
// out[i][j] = in[i] (j is undef for input)
|
|
tid = outTid;
|
|
lvl = outLvl;
|
|
// Skips invalid lvl (e.g., when this is a zero ranked tensor).
|
|
if (!lvl)
|
|
return;
|
|
}
|
|
hasNonUnique = !isUniqueDLT(dlt) || hasNonUnique;
|
|
tids.push_back(tid);
|
|
lvls.push_back(*lvl);
|
|
numloopCond++;
|
|
} else if (isDenseDLT(dlt)) {
|
|
tids.push_back(tid);
|
|
lvls.push_back(*lvl);
|
|
} else {
|
|
assert(isUndefDLT(dlt));
|
|
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++) {
|
|
// FIXME: `toOrigDim` is deprecated.
|
|
AffineExpr exp = affines[toOrigDim(stt.getEncoding(), l)];
|
|
// Skip simple affine expression and non-dense levels (which
|
|
// have their own filter loop).
|
|
if (exp.isa<AffineDimExpr>() || !stt.isDenseLvl(l))
|
|
continue;
|
|
|
|
// Constant affine expression are handled in genLoop
|
|
if (!exp.isa<AffineConstantExpr>()) {
|
|
bool isAtLoop = false;
|
|
if (isInvariantAffine(env, exp, ldx, isAtLoop) && isAtLoop) {
|
|
// 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.
|
|
// NOTE: It assumes that the levels of the input tensor are
|
|
// initialized in order (and it is also currently guaranteed by
|
|
// computeIterationGraph), another more admissible approach
|
|
// might be accepting out-of-order access between consecutive
|
|
// dense levels.
|
|
affineTids.push_back(tid);
|
|
affineLvls.push_back(l);
|
|
exps.push_back(exp);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
if (isDenseDLT(env.dlt(outTid, ldx))) {
|
|
// 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.
|
|
tids.push_back(outTid);
|
|
lvls.push_back(*outLvl);
|
|
}
|
|
|
|
assert(numloopCond > 0);
|
|
// 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 single loop in current sequence.
|
|
static std::pair<Operation *, bool> startLoop(CodegenEnv &env,
|
|
OpBuilder &builder, unsigned at,
|
|
unsigned li, bool needsUniv) {
|
|
// The set of tensors + lvls to generate loops on
|
|
SmallVector<TensorId> tids, affineTids;
|
|
SmallVector<Level> lvls, affineLvls;
|
|
// The set of dense tensors with non-trivial affine expression that just
|
|
// becomes invariant and the address shall now be generated at the current
|
|
// level.
|
|
SmallVector<AffineExpr> affines;
|
|
bool isSingleCond = translateBitsToTidLvlPairs(
|
|
env, li, env.topSortAt(at), tids, lvls, affineTids, affineLvls, affines);
|
|
|
|
// Emit the for/while-loop control.
|
|
Operation *loop =
|
|
genLoop(env, builder, at, needsUniv, tids, lvls, isSingleCond);
|
|
Location loc = env.op().getLoc();
|
|
for (auto [tid, lvl, exp] : llvm::zip(affineTids, affineLvls, affines)) {
|
|
env.emitter().genDenseAffineAddress(builder, loc, tid, lvl, 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 allTids = llvm::concat<TensorId>(tids, affineTids);
|
|
auto allLvls = llvm::concat<Level>(lvls, affineLvls);
|
|
for (auto [tid, lvl] : llvm::zip(allTids, allLvls)) {
|
|
if (tid != env.merger().getOutTensorID())
|
|
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,
|
|
LoopId idx, LatPointId li, bool needsUniv) {
|
|
// End a while-loop.
|
|
if (auto whileOp = dyn_cast<scf::WhileOp>(loop)) {
|
|
finalizeWhileOp(env, rewriter, idx, needsUniv, env.lat(li).bits, whileOp);
|
|
} else if (auto forOp = dyn_cast<scf::ForOp>(loop)) {
|
|
// Any iteration of a reduction for-loop creates a valid lex insert.
|
|
if (env.isReduc() && env.getValidLexInsert())
|
|
env.setValidLexInsert(constantI1(rewriter, env.op().getLoc(), true));
|
|
} 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, ExprId exp,
|
|
LoopOrd at, LoopId idx, LoopId ldx) {
|
|
assert(!env.getLoopVar(idx));
|
|
env.emitter().exitCurrentLoopSeq();
|
|
// Unmark bookkeeping of invariants and loop index.
|
|
genInvariants(env, builder, exp, ldx, /*atStart=*/false);
|
|
// Finalize access pattern expansion for sparse tensor output.
|
|
genExpand(env, builder, at, /*atStart=*/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,
|
|
LoopOrd at) {
|
|
// At each leaf, assign remaining tensor (sub)expression to output tensor.
|
|
if (at == env.topSortSize()) {
|
|
const LoopId ldx = env.topSortAt(at - 1);
|
|
Value rhs = genExp(env, rewriter, exp, ldx);
|
|
genTensorStore(env, rewriter, exp, rhs);
|
|
return;
|
|
}
|
|
|
|
// Construct iteration lattices for current loop index, with L0 at top.
|
|
const LoopId idx = env.topSortAt(at);
|
|
const LoopId ldx = at == 0 ? kInvalidId : env.topSortAt(at - 1);
|
|
const LatSetId lts =
|
|
env.merger().optimizeSet(env.merger().buildLattices(exp, idx));
|
|
|
|
// Start a loop sequence.
|
|
bool needsUniv = startLoopSeq(env, rewriter, exp, at, idx, ldx, lts);
|
|
|
|
// Emit a loop for every lattice point L0 >= Li in this loop sequence.
|
|
unsigned lsize = env.set(lts).size();
|
|
for (unsigned i = 0; i < lsize; i++) {
|
|
// Start a loop.
|
|
const LatPointId li = env.set(lts)[i];
|
|
auto [loop, isSingleCond] = startLoop(env, rewriter, at, 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();
|
|
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, idx, env.lat(lj).simple);
|
|
genStmt(env, rewriter, ej, at + 1);
|
|
endIf(env, rewriter, ifOp, loop, redInput, cntInput, insInput);
|
|
} else {
|
|
genStmt(env, rewriter, ej, at + 1);
|
|
}
|
|
}
|
|
}
|
|
|
|
// End a loop.
|
|
needsUniv = endLoop(env, rewriter, loop, idx, li, needsUniv);
|
|
}
|
|
|
|
// End a loop sequence.
|
|
endLoopSeq(env, rewriter, exp, at, idx, ldx);
|
|
}
|
|
|
|
/// 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().back(); // value array
|
|
rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, resType, val);
|
|
}
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Sparse compiler 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 without affine index on sparse
|
|
// output.
|
|
if (op.getNumDpsInits() != 1 || hasNonTrivialAffineOnSparseOut(op))
|
|
return failure();
|
|
|
|
// Sets up a code generation environment.
|
|
const unsigned numTensors = op->getNumOperands();
|
|
const unsigned numLoops = op.getNumLoops();
|
|
const unsigned numFilterLoops = getNumNonTrivialIdxExpOnSparseLvls(op);
|
|
// TODO: we should probably always use slice-based codegen whenever
|
|
// possible, we can even intermix slice-based and filter-loop based codegen.
|
|
bool idxReducBased = options.enableIndexReduction && numFilterLoops != 0;
|
|
|
|
// If we uses slice based algorithm for affine index, we do not need filter
|
|
// loop.
|
|
CodegenEnv env(op, options, numTensors, numLoops,
|
|
/*numFilterLoops=*/idxReducBased ? 0 : numFilterLoops);
|
|
|
|
// Detects sparse annotations and translates the per-level sparsity
|
|
// information for all tensors to loop indices in the kernel.
|
|
if (!findSparseAnnotations(env, idxReducBased))
|
|
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();
|
|
|
|
// Computes a topologically sorted iteration graph to ensure tensors
|
|
// are visited in natural index order. Gradually relaxes the considered
|
|
// constraints until an acyclic iteration graph results, such that sparse
|
|
// code generation can proceed. As a last resort, an attempt is made
|
|
// to resolve cycles by inserting a conversion.
|
|
bool isAdmissible = false;
|
|
bool hasCycle = true;
|
|
|
|
// An const list of all masks that we used for interation graph
|
|
// computation. Must be ordered from more strict to less strict.
|
|
// Ideally (though might not be guaranteed), the eariler a constraint mask
|
|
// can be satisfied, the faster the generated kernel will be.
|
|
const auto allMasks = {
|
|
SortMask::kIncludeAll, SortMask::kIncludeDense,
|
|
SortMask::kIncludeDenseInput, SortMask::kIncludeDenseOutput,
|
|
SortMask::kIncludeUndef, SortMask::kSparseOnly};
|
|
for (const SortMask mask : allMasks) {
|
|
if (computeIterationGraph(env, mask, nullptr, idxReducBased)) {
|
|
hasCycle = false;
|
|
if (env.isAdmissibleTopoOrder()) {
|
|
isAdmissible = true;
|
|
break;
|
|
}
|
|
// else try a set of less strict constraints.
|
|
}
|
|
}
|
|
if (hasCycle) {
|
|
return idxReducBased
|
|
? failure() // TODO: should cycle be resolved differently?
|
|
: resolveCycle(env, rewriter); // one last shot
|
|
}
|
|
|
|
if (!isAdmissible)
|
|
return failure(); // inadmissible expression, reject
|
|
|
|
// 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 becasue we already reject the
|
|
// constant expression at a earlier stage.
|
|
genInitConstantDenseAddress(env, rewriter);
|
|
genStmt(env, rewriter, env.getExprId(), 0);
|
|
genResult(env, rewriter);
|
|
return success();
|
|
}
|
|
|
|
private:
|
|
// Last resort cycle resolution.
|
|
LogicalResult resolveCycle(CodegenEnv &env, PatternRewriter &rewriter) const {
|
|
// Compute topological sort while leaving out every
|
|
// sparse input tensor in succession until an acylic
|
|
// iteration graph results.
|
|
for (OpOperand *t : env.op().getDpsInputOperands()) {
|
|
const TensorId tid = t->getOperandNumber();
|
|
Value tval = t->get();
|
|
auto srcEnc = getSparseTensorEncoding(tval.getType());
|
|
if (!srcEnc || !computeIterationGraph(env, SortMask::kSparseOnly, t))
|
|
continue;
|
|
// Found an input tensor that resolves the cycle by inserting a
|
|
// conversion into a sparse tensor that adheres to the iteration
|
|
// graph order. Also releases the temporary sparse tensor.
|
|
//
|
|
// TODO: investigate fusing the conversion with computation,
|
|
// especially if it is a direct yield!
|
|
//
|
|
auto srcTp = getRankedTensorType(tval);
|
|
auto dstEnc = SparseTensorEncodingAttr::get(
|
|
getContext(), srcEnc.getDimLevelType(),
|
|
permute(env, env.op().getMatchingIndexingMap(t)), // new order
|
|
srcEnc.getHigherOrdering(), srcEnc.getPosWidth(),
|
|
srcEnc.getCrdWidth());
|
|
auto dstTp = RankedTensorType::get(srcTp.getShape(),
|
|
srcTp.getElementType(), dstEnc);
|
|
auto convert = rewriter.create<ConvertOp>(tval.getLoc(), dstTp, tval);
|
|
rewriter.updateRootInPlace(env.op(),
|
|
[&]() { env.op()->setOperand(tid, convert); });
|
|
rewriter.setInsertionPointAfter(env.op());
|
|
rewriter.create<bufferization::DeallocTensorOp>(tval.getLoc(), convert);
|
|
return success();
|
|
}
|
|
// Cannot be resolved with a single conversion.
|
|
// TODO: convert more than one?
|
|
return failure();
|
|
}
|
|
|
|
/// 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);
|
|
}
|