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

[mlir][sparse] finalize sparse output in the presence of reductions 

This revision implements sparse outputs (from scratch) in all cases where
the loops can be reordered with all but one parallel loops outer. If the
inner parallel loop appears inside one or more reductions loops, then an
access pattern expansion is required (aka. workspaces in TACO speak).

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D115091
This commit is contained in:
Aart Bik 2021-12-03 16:55:43 -08:00
parent d9aa5f6158
commit 4f2ec7f983
9 changed files with 786 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

75
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
View File

@ -219,6 +219,81 @@ def SparseTensor_LexInsertOp : SparseTensor_Op<"lex_insert", []>,
" type($tensor) `,` type($indices) `,` type($value)";
}
def SparseTensor_ExpandOp : SparseTensor_Op<"expand", []>,
Arguments<(ins AnyTensor:$tensor)>,
Results<(outs AnyStridedMemRefOfRank<1>:$values,
StridedMemRefRankOf<[I1],[1]>:$filled,
StridedMemRefRankOf<[Index],[1]>:$added,
Index:$count)> {
string summary = "Expands an access pattern for insertion";
string description = [{
Performs an access pattern expansion for the innermost dimensions of the
given tensor. This operation is useful to implement kernels in which a
sparse tensor appears as output. This technique is known under several
different names and using several alternative implementations,
for example, phase counter [Gustavson71], expanded or switch array
[Pissanetzky84], in phase scan [Duff90], access pattern expansion [Bik96],
and workspaces [Kjolstad2018].
The values and filled array have a size the suffices for a *dense* innermost
dimension (e.g. a full row for matrices). The added array and count are used
to store new indices when a false value is encountered in the filled array.
All arrays should be allocated before the loop (possibly even shared between
loops in a future optimization) so that their *dense* intitialization can be
amortized over many iterations. Setting and resetting the dense arrays in
the loop nest itself is kept *sparse* by only iterating over set elements
through an indirection using the added array, so that the operations are
kept proportional to the number of nonzeros.
Note that this operation is "impure" in the sense that its behavior
is solely defined by side-effects and not SSA values. The semantics
may be refined over time as our sparse abstractions evolve.
Example:
```mlir
%values, %filled, %added, %count = sparse_tensor.expand %0
: tensor<4x4xf64, #CSR> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
```
}];
let assemblyFormat = "$tensor attr-dict `:` type($tensor) `to` type($values)"
" `,` type($filled) `,` type($added) `,` type($count)";
}
def SparseTensor_CompressOp : SparseTensor_Op<"compress", []>,
Arguments<(ins AnyTensor:$tensor,
StridedMemRefRankOf<[Index],[1]>:$indices,
AnyStridedMemRefOfRank<1>:$values,
StridedMemRefRankOf<[I1],[1]>:$filled,
StridedMemRefRankOf<[Index],[1]>:$added,
Index:$count)> {
string summary = "Compressed an access pattern for insertion";
string description = [{
Finishes a single access pattern by moving the inserted elements
into the sparse storage scheme. The values and filled array are reset
in a *sparse* fashion by only iterating over set elements through an
indirection using the added array, so that the operations are kept
proportional to the number of nonzeros. See the 'expand' operation
for more details.
Note that this operation is "impure" in the sense that its behavior
is solely defined by side-effects and not SSA values. The semantics
may be refined over time as our sparse abstractions evolve.
Example:
```mlir
sparse_tensor.compress %0, %1, %values, %filled, %added, %2
: tensor<4x4xf64, #CSR>, memref<?xindex>, memref<?xf64>,
memref<?xi1>, memref<?xindex>, index
```
}];
let assemblyFormat = "$tensor `,` $indices `,` $values `,` $filled `,`"
" $added `,` $count attr-dict `:` type($tensor) `,`"
" type($indices) `,` type($values) `,` type($filled) `,`"
" type($added) `,` type($count)";
}
def SparseTensor_LoadOp : SparseTensor_Op<"load", [SameOperandsAndResultType]>,
Arguments<(ins AnyTensor:$tensor, UnitAttr:$hasInserts)>,
Results<(outs AnyTensor:$result)> {

12
mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
View File

@ -305,6 +305,18 @@ static LogicalResult verify(LexInsertOp op) {
return success();
}
static LogicalResult verify(ExpandOp op) {
if (!getSparseTensorEncoding(op.tensor().getType()))
return op.emitError("expected a sparse tensor for expansion");
return success();
}
static LogicalResult verify(CompressOp op) {
if (!getSparseTensorEncoding(op.tensor().getType()))
return op.emitError("expected a sparse tensor for compression");
return success();
}
static LogicalResult verify(LoadOp op) {
if (!getSparseTensorEncoding(op.tensor().getType()))
return op.emitError("expected a sparse tensor to materialize");

88
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
View File

@ -257,10 +257,17 @@ static void sizesFromPtr(ConversionPatternRewriter &rewriter,
/// type, but returns it as type `memref<? x $tp>` (rather than as type
/// `memref<$sz x $tp>`).
static Value genAlloca(ConversionPatternRewriter &rewriter, Location loc,
unsigned sz, Type tp) {
Value sz, Type tp) {
auto memTp = MemRefType::get({ShapedType::kDynamicSize}, tp);
Value a = constantIndex(rewriter, loc, sz);
return rewriter.create<memref::AllocaOp>(loc, memTp, ValueRange{a});
return rewriter.create<memref::AllocaOp>(loc, memTp, ValueRange{sz});
}
/// Generates an uninitialized temporary buffer of the given size and
/// type, but returns it as type `memref<? x $tp>` (rather than as type
/// `memref<$sz x $tp>`).
static Value genAlloca(ConversionPatternRewriter &rewriter, Location loc,
unsigned sz, Type tp) {
return genAlloca(rewriter, loc, constantIndex(rewriter, loc, sz), tp);
}
/// Generates an uninitialized temporary buffer with room for one value
@ -911,6 +918,78 @@ public:
}
};
class SparseTensorExpandConverter : public OpConversionPattern<ExpandOp> {
public:
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(ExpandOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
Location loc = op->getLoc();
ShapedType srcType = op.tensor().getType().cast<ShapedType>();
Type eltType = srcType.getElementType();
Type boolType = rewriter.getIntegerType(1);
Type idxType = rewriter.getIndexType();
// All initialization should be done on entry of the loop nest.
rewriter.setInsertionPointAfter(op.tensor().getDefiningOp());
// Determine the size for access expansion.
auto enc = getSparseTensorEncoding(srcType);
Value src = adaptor.getOperands()[0];
Value sz = genDimSizeCall(rewriter, op, enc, src, srcType.getRank() - 1);
// Allocate temporary stack buffers for values, filled-switch, and indices.
Value values = genAlloca(rewriter, loc, sz, eltType);
Value filled = genAlloca(rewriter, loc, sz, boolType);
Value indices = genAlloca(rewriter, loc, sz, idxType);
Value zero = constantZero(rewriter, loc, idxType);
// Reset the values/filled-switch to all-zero/false. Note that this
// introduces an O(N) operation into the computation, but this reset
// operation is amortized over the innermost loops for the access
// pattern expansion.
rewriter.create<linalg::FillOp>(loc, constantZero(rewriter, loc, eltType),
values);
rewriter.create<linalg::FillOp>(loc, constantZero(rewriter, loc, boolType),
filled);
// Replace expansion op with these buffers and initial index.
assert(op.getNumResults() == 4);
rewriter.replaceOp(op, {values, filled, indices, zero});
return success();
}
};
class SparseTensorCompressConverter : public OpConversionPattern<CompressOp> {
public:
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(CompressOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
// Note that this method call resets the values/filled-switch back to
// all-zero/false by only iterating over the set elements, so the
// complexity remains proportional to the sparsity of the expanded
// access pattern.
Type srcType = op.tensor().getType();
Type eltType = srcType.cast<ShapedType>().getElementType();
StringRef name;
if (eltType.isF64())
name = "expInsertF64";
else if (eltType.isF32())
name = "expInsertF32";
else if (eltType.isInteger(64))
name = "expInsertI64";
else if (eltType.isInteger(32))
name = "expInsertI32";
else if (eltType.isInteger(16))
name = "expInsertI16";
else if (eltType.isInteger(8))
name = "expInsertI8";
else
return failure();
TypeRange noTp;
auto fn =
getFunc(op, name, noTp, adaptor.getOperands(), /*emitCInterface=*/true);
rewriter.replaceOpWithNewOp<CallOp>(op, noTp, fn, adaptor.getOperands());
return success();
}
};
} // namespace
//===----------------------------------------------------------------------===//
@ -926,6 +1005,7 @@ void mlir::populateSparseTensorConversionPatterns(TypeConverter &typeConverter,
SparseTensorInitConverter, SparseTensorConvertConverter,
SparseTensorReleaseConverter, SparseTensorToPointersConverter,
SparseTensorToIndicesConverter, SparseTensorToValuesConverter,
SparseTensorLoadConverter, SparseTensorLexInsertConverter>(
SparseTensorLoadConverter, SparseTensorLexInsertConverter,
SparseTensorExpandConverter, SparseTensorCompressConverter>(
typeConverter, patterns.getContext());
}

197
mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
View File

@ -54,7 +54,8 @@ struct CodeGen {
pidxs(numTensors, std::vector<Value>(numLoops)),
idxs(numTensors, std::vector<Value>(numLoops)), redExp(-1u), redVal(),
redKind(kNoReduc), sparseOut(op), outerParNest(nest), lexIdx(),
curVecLength(1), curVecMask() {}
expValues(), expFilled(), expAdded(), expCount(), curVecLength(1),
curVecMask() {}
/// Sparsification options.
SparsificationOptions options;
/// Universal dense indices and upper bounds (by index). The loops array
@ -81,10 +82,15 @@ struct CodeGen {
Reduction redKind;
// Sparse tensor as output. Implemented either through direct injective
// insertion in lexicographic index order (where indices are updated
// in the temporary array `lexIdx`) or TODO: access pattern expansion
// in the temporary array `lexIdx`) or through access pattern expansion
// in the innermost loop nest (`expValues` through `expCount`).
OpOperand *sparseOut;
unsigned outerParNest;
Value lexIdx;
Value expValues;
Value expFilled;
Value expAdded;
Value expCount;
// Current vector length and mask.
unsigned curVecLength;
Value curVecMask;
@ -334,8 +340,8 @@ static bool isAdmissableTensorExp(Merger &merger, linalg::GenericOp op,
}
// Determine admissable dynamic insertion situations:
// (1) fully injective, since there are no reductions,
// (2) admissable 1-d expansion in innermost dimension. TODO: accept
if (nest == op.getRank(lhs)) {
// (2) admissable 1-d expansion in innermost dimension.
if (nest >= op.getRank(lhs) - 1) {
*sparseOut = lhs;
outerParNest = nest;
return true;
@ -680,6 +686,16 @@ static Value genAffine(CodeGen &codegen, PatternRewriter &rewriter,
}
}
/// Generates index for load/store on sparse tensor.
static Value genIndex(CodeGen &codegen, linalg::GenericOp op, OpOperand *t) {
auto map = op.getTiedIndexingMap(t);
auto enc = getSparseTensorEncoding(t->get().getType());
AffineExpr a = map.getResult(perm(enc, map.getNumResults() - 1));
assert(a.getKind() == AffineExprKind::DimId);
unsigned idx = a.cast<AffineDimExpr>().getPosition();
return codegen.loops[idx];
}
/// Generates subscript for load/store on a dense or sparse tensor.
static Value genSubscript(CodeGen &codegen, PatternRewriter &rewriter,
linalg::GenericOp op, OpOperand *t,
@ -705,6 +721,62 @@ static Value genSubscript(CodeGen &codegen, PatternRewriter &rewriter,
return codegen.buffers[tensor];
}
/// Generates insertion code to implement dynamic tensor load.
static Value genInsertionLoad(CodeGen &codegen, PatternRewriter &rewriter,
linalg::GenericOp op, OpOperand *t) {
Location loc = op.getLoc();
// Direct lexicographic index order, tensor loads as zero.
if (!codegen.expValues) {
Type tp = getElementTypeOrSelf(t->get().getType());
return rewriter.create<arith::ConstantOp>(loc, tp,
rewriter.getZeroAttr(tp));
}
// Load from expanded access pattern.
Value index = genIndex(codegen, op, t);
return rewriter.create<memref::LoadOp>(loc, codegen.expValues, index);
}
/// Generates insertion code to implement dynamic tensor store.
static void genInsertionStore(CodeGen &codegen, PatternRewriter &rewriter,
linalg::GenericOp op, OpOperand *t, Value rhs) {
Location loc = op.getLoc();
// Direct insertion in lexicographic index order.
if (!codegen.expValues) {
rewriter.create<LexInsertOp>(loc, t->get(), codegen.lexIdx, rhs);
return;
}
// Generates insertion code along expanded access pattern.
// if (!expFilled[i]) then
// expFilled[i] = true
// expAdded[inserts++] = i
// endif
// values[i] = rhs
Value index = genIndex(codegen, op, t);
Value fval = rewriter.create<arith::ConstantIntOp>(loc, 0, 1); // false
Value tval = rewriter.create<arith::ConstantIntOp>(loc, 1, 1); // true
// If statement.
Value filled = rewriter.create<memref::LoadOp>(loc, codegen.expFilled, index);
Value cond = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
filled, fval);
scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, rewriter.getIndexType(),
cond, /*else=*/true);
// True branch.
rewriter.setInsertionPointToStart(&ifOp.thenRegion().front());
rewriter.create<memref::StoreOp>(loc, tval, codegen.expFilled, index);
rewriter.create<memref::StoreOp>(loc, index, codegen.expAdded,
codegen.expCount);
Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
Value add = rewriter.create<arith::AddIOp>(loc, codegen.expCount, one);
rewriter.create<scf::YieldOp>(loc, add);
// False branch.
rewriter.setInsertionPointToStart(&ifOp.elseRegion().front());
rewriter.create<scf::YieldOp>(loc, codegen.expCount);
rewriter.setInsertionPointAfter(ifOp);
// Value assignment.
codegen.expCount = ifOp.getResult(0);
rewriter.create<memref::StoreOp>(loc, rhs, codegen.expValues, index);
}
/// Generates a load on a dense or sparse tensor.
static Value genTensorLoad(Merger &merger, CodeGen &codegen,
PatternRewriter &rewriter, linalg::GenericOp op,
@ -716,13 +788,10 @@ static Value genTensorLoad(Merger &merger, CodeGen &codegen,
return genVectorInvariantValue(codegen, rewriter, val);
return val;
}
// Insertion (a sparse tensor output "loads" as zero).
// Load during insertion.
OpOperand *t = op.getInputAndOutputOperands()[merger.exp(exp).tensor];
if (t == codegen.sparseOut) {
Type tp = getElementTypeOrSelf(t->get().getType());
return rewriter.create<arith::ConstantOp>(op.getLoc(), tp,
rewriter.getZeroAttr(tp));
}
if (t == codegen.sparseOut)
return genInsertionLoad(codegen, rewriter, op, t);
// Actual load.
SmallVector<Value, 4> args;
Value ptr = genSubscript(codegen, rewriter, op, t, args);
@ -744,10 +813,10 @@ static void genTensorStore(Merger &merger, CodeGen &codegen,
updateReduc(merger, codegen, rhs);
return;
}
// Insertion.
// Store during insertion.
OpOperand *t = op.getOutputOperand(0);
if (t == codegen.sparseOut) {
rewriter.create<LexInsertOp>(loc, t->get(), codegen.lexIdx, rhs);
genInsertionStore(codegen, rewriter, op, t, rhs);
return;
}
// Actual store.
@ -916,6 +985,42 @@ static void genInvariants(Merger &merger, CodeGen &codegen,
}
}
/// Generates an expanded access pattern in innermost dimension.
static void genExpansion(Merger &merger, CodeGen &codegen,
PatternRewriter &rewriter, linalg::GenericOp op,
unsigned at, bool atStart) {
OpOperand *lhs = codegen.sparseOut;
if (!lhs || codegen.outerParNest != op.getRank(lhs) - 1 ||
at != codegen.outerParNest)
return; // not needed at this level
// Generate start or end of an expanded access pattern.
Value tensor = lhs->get();
Location loc = op.getLoc();
if (atStart) {
auto dynShape = {ShapedType::kDynamicSize};
Type etp = tensor.getType().cast<ShapedType>().getElementType();
Type t1 = MemRefType::get(dynShape, etp);
Type t2 = MemRefType::get(dynShape, genIntType(rewriter, 1));
Type t3 = MemRefType::get(dynShape, genIntType(rewriter, 0));
Type t4 = rewriter.getIndexType();
auto res =
rewriter.create<ExpandOp>(loc, TypeRange({t1, t2, t3, t4}), tensor);
assert(res.getNumResults() == 4);
assert(!codegen.expValues);
codegen.expValues = res.getResult(0);
codegen.expFilled = res.getResult(1);
codegen.expAdded = res.getResult(2);
codegen.expCount = res.getResult(3);
} else {
assert(codegen.expValues);
rewriter.create<CompressOp>(loc, tensor, codegen.lexIdx, codegen.expValues,
codegen.expFilled, codegen.expAdded,
codegen.expCount);
codegen.expValues = codegen.expFilled = codegen.expAdded =
codegen.expCount = Value();
}
}
/// Generates initialization code for the subsequent loop sequence at
/// current index level. Returns true if the loop sequence needs to
/// maintain the universal index.
@ -1069,9 +1174,13 @@ static Operation *genFor(Merger &merger, CodeGen &codegen,
}
operands.push_back(codegen.redVal);
}
if (codegen.expValues)
operands.push_back(codegen.expCount);
scf::ForOp forOp = rewriter.create<scf::ForOp>(loc, lo, hi, step, operands);
if (codegen.redVal)
updateReduc(merger, codegen, forOp.getRegionIterArgs().front());
if (codegen.expValues)
codegen.expCount = forOp.getRegionIterArgs().back();
// Assign induction variable to sparse or dense index.
Value iv = forOp.getInductionVar();
if (isSparse)
@ -1106,6 +1215,10 @@ static Operation *genWhile(Merger &merger, CodeGen &codegen,
types.push_back(codegen.redVal.getType());
operands.push_back(codegen.redVal);
}
if (codegen.expValues) {
types.push_back(indexType);
operands.push_back(codegen.expCount);
}
if (needsUniv) {
types.push_back(indexType);
operands.push_back(codegen.loops[idx]);
@ -1135,6 +1248,8 @@ static Operation *genWhile(Merger &merger, CodeGen &codegen,
}
if (codegen.redVal)
updateReduc(merger, codegen, after->getArgument(o++));
if (codegen.expValues)
codegen.expCount = after->getArgument(o++);
if (needsUniv)
codegen.loops[idx] = after->getArgument(o++);
assert(o == operands.size());
@ -1215,8 +1330,9 @@ static void genLocals(Merger &merger, CodeGen &codegen,
}
}
// Move the insertion indices in lexicographic index order.
if (codegen.sparseOut) {
// Move the insertion indices in lexicographic index order. During access
// pattern expansion, we can skip setting the innermost dimension.
if (codegen.sparseOut && !codegen.expValues) {
Value pos = rewriter.create<arith::ConstantIndexOp>(loc, at);
rewriter.create<memref::StoreOp>(loc, codegen.loops[idx], codegen.lexIdx,
pos);
@ -1231,11 +1347,21 @@ static void genWhileInduction(Merger &merger, CodeGen &codegen,
scf::WhileOp whileOp) {
Location loc = op.getLoc();
// Finalize each else branch of all if statements.
if (codegen.redVal) {
if (codegen.redVal || codegen.expValues) {
while (auto ifOp = dyn_cast_or_null<scf::IfOp>(
rewriter.getInsertionBlock()->getParentOp())) {
rewriter.create<scf::YieldOp>(loc, codegen.redVal);
updateReduc(merger, codegen, ifOp.getResult(0));
unsigned y = 0;
SmallVector<Value, 4> yields;
if (codegen.redVal) {
yields.push_back(codegen.redVal);
updateReduc(merger, codegen, ifOp.getResult(y++));
}
if (codegen.expValues) {
yields.push_back(codegen.expCount);
codegen.expCount = ifOp->getResult(y++);
}
assert(y == yields.size());
rewriter.create<scf::YieldOp>(loc, yields);
rewriter.setInsertionPointAfter(ifOp);
}
}
@ -1266,6 +1392,10 @@ static void genWhileInduction(Merger &merger, CodeGen &codegen,
operands.push_back(codegen.redVal);
updateReduc(merger, codegen, whileOp->getResult(o++));
}
if (codegen.expValues) {
operands.push_back(codegen.expCount);
codegen.expCount = whileOp->getResult(o++);
}
if (needsUniv) {
operands.push_back(
rewriter.create<arith::AddIOp>(loc, codegen.loops[idx], one));
@ -1287,6 +1417,10 @@ static void genForInduction(Merger &merger, CodeGen &codegen,
operands.push_back(codegen.redVal);
updateReduc(merger, codegen, loop->getResult(o++));
}
if (codegen.expValues) {
operands.push_back(codegen.expCount);
codegen.expCount = loop->getResult(o++);
}
assert(o == operands.size());
if (o > 0)
rewriter.create<scf::YieldOp>(loc, operands);
@ -1318,6 +1452,8 @@ static scf::IfOp genIf(Merger &merger, CodeGen &codegen,
}
if (codegen.redVal)
types.push_back(codegen.redVal.getType());
if (codegen.expValues)
types.push_back(rewriter.getIndexType());
scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, types, cond, /*else=*/true);
rewriter.setInsertionPointToStart(&ifOp.thenRegion().front());
return ifOp;
@ -1325,11 +1461,19 @@ static scf::IfOp genIf(Merger &merger, CodeGen &codegen,
/// Generates end of true branch of if-statement within a while-loop.
static void endIf(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
linalg::GenericOp op, scf::IfOp ifOp, Value ifInput) {
linalg::GenericOp op, scf::IfOp ifOp, Operation *loop,
Value redInput, Value cntInput) {
SmallVector<Value, 4> operands;
if (codegen.redVal) {
rewriter.create<scf::YieldOp>(op.getLoc(), codegen.redVal);
updateReduc(merger, codegen, ifInput);
operands.push_back(codegen.redVal);
updateReduc(merger, codegen, redInput);
}
if (codegen.expValues) {
operands.push_back(codegen.expCount);
codegen.expCount = cntInput;
}
if (!operands.empty())
rewriter.create<scf::YieldOp>(op.getLoc(), operands);
rewriter.setInsertionPointToStart(&ifOp.elseRegion().front());
}
@ -1348,6 +1492,8 @@ static bool startLoopSeq(Merger &merger, CodeGen &codegen,
assert(!codegen.loops[idx]);
// Emit invariants at this loop sequence level.
genInvariants(merger, codegen, rewriter, op, exp, ldx, /*atStart=*/true);
// Emit access pattern expansion for sparse tensor output.
genExpansion(merger, codegen, rewriter, op, at, /*atStart=*/true);
// Emit further intitialization at this loop sequence level.
unsigned l0 = merger.set(lts)[0];
bool needsUniv =
@ -1399,7 +1545,7 @@ static bool endLoop(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
/// Ends a loop sequence at given level.
static void endLoopSeq(Merger &merger, CodeGen &codegen,
PatternRewriter &rewriter, linalg::GenericOp op,
unsigned exp, unsigned idx, unsigned ldx) {
unsigned exp, unsigned at, unsigned idx, unsigned ldx) {
assert(codegen.curVecLength == 1);
codegen.loops[idx] = Value();
// Bring a pending reduction back from SIMD form when sequence ends.
@ -1409,6 +1555,8 @@ static void endLoopSeq(Merger &merger, CodeGen &codegen,
genVectorReducEnd(codegen, rewriter, op.getLoc(), vtp));
// Unmark bookkeeping of invariants and loop index.
genInvariants(merger, codegen, rewriter, op, exp, ldx, /*atStart=*/false);
// Finalize access pattern expansion for sparse tensor output.
genExpansion(merger, codegen, rewriter, op, at, /*atStart=*/false);
}
/// Recursively generates code while computing iteration lattices in order
@ -1443,7 +1591,8 @@ static void genStmt(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
// Visit all lattices points with Li >= Lj to generate the
// loop-body, possibly with if statements for coiteration.
Value ifInput = codegen.redVal;
Value redInput = codegen.redVal;
Value cntInput = codegen.expCount;
bool isWhile = dyn_cast<scf::WhileOp>(loop) != nullptr;
for (unsigned j = 0; j < lsize; j++) {
unsigned lj = merger.set(lts)[j];
@ -1454,7 +1603,7 @@ static void genStmt(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
scf::IfOp ifOp =
genIf(merger, codegen, rewriter, op, idx, merger.lat(lj).simple);
genStmt(merger, codegen, rewriter, op, topSort, ej, at + 1);
endIf(merger, codegen, rewriter, op, ifOp, ifInput);
endIf(merger, codegen, rewriter, op, ifOp, loop, redInput, cntInput);
} else {
genStmt(merger, codegen, rewriter, op, topSort, ej, at + 1);
}
@ -1467,7 +1616,7 @@ static void genStmt(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
}
// End a loop sequence.
endLoopSeq(merger, codegen, rewriter, op, exp, idx, ldx);
endLoopSeq(merger, codegen, rewriter, op, exp, at, idx, ldx);
}
/// Converts the result computed by the sparse kernel into the required form.

107
mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
View File

@ -163,10 +163,10 @@ private:
/// function overloading to implement "partial" method specialization.
class SparseTensorStorageBase {
public:
// Dimension size query.
/// Dimension size query.
virtual uint64_t getDimSize(uint64_t) = 0;
// Overhead storage.
/// Overhead storage.
virtual void getPointers(std::vector<uint64_t> **, uint64_t) { fatal("p64"); }
virtual void getPointers(std::vector<uint32_t> **, uint64_t) { fatal("p32"); }
virtual void getPointers(std::vector<uint16_t> **, uint64_t) { fatal("p16"); }
@ -176,7 +176,7 @@ public:
virtual void getIndices(std::vector<uint16_t> **, uint64_t) { fatal("i16"); }
virtual void getIndices(std::vector<uint8_t> **, uint64_t) { fatal("i8"); }
// Primary storage.
/// Primary storage.
virtual void getValues(std::vector<double> **) { fatal("valf64"); }
virtual void getValues(std::vector<float> **) { fatal("valf32"); }
virtual void getValues(std::vector<int64_t> **) { fatal("vali64"); }
@ -184,13 +184,35 @@ public:
virtual void getValues(std::vector<int16_t> **) { fatal("vali16"); }
virtual void getValues(std::vector<int8_t> **) { fatal("vali8"); }
// Element-wise insertion in lexicographic index order.
/// Element-wise insertion in lexicographic index order.
virtual void lexInsert(uint64_t *, double) { fatal("insf64"); }
virtual void lexInsert(uint64_t *, float) { fatal("insf32"); }
virtual void lexInsert(uint64_t *, int64_t) { fatal("insi64"); }
virtual void lexInsert(uint64_t *, int32_t) { fatal("insi32"); }
virtual void lexInsert(uint64_t *, int16_t) { fatal("ins16"); }
virtual void lexInsert(uint64_t *, int8_t) { fatal("insi8"); }
/// Expanded insertion.
virtual void expInsert(uint64_t *, double *, bool *, uint64_t *, uint64_t) {
fatal("expf64");
}
virtual void expInsert(uint64_t *, float *, bool *, uint64_t *, uint64_t) {
fatal("expf32");
}
virtual void expInsert(uint64_t *, int64_t *, bool *, uint64_t *, uint64_t) {
fatal("expi64");
}
virtual void expInsert(uint64_t *, int32_t *, bool *, uint64_t *, uint64_t) {
fatal("expi32");
}
virtual void expInsert(uint64_t *, int16_t *, bool *, uint64_t *, uint64_t) {
fatal("expi16");
}
virtual void expInsert(uint64_t *, int8_t *, bool *, uint64_t *, uint64_t) {
fatal("expi8");
}
/// Finishes insertion.
virtual void endInsert() = 0;
virtual ~SparseTensorStorageBase() {}
@ -289,6 +311,35 @@ public:
insPath(cursor, diff, top, val);
}
/// Partially specialize expanded insertions based on template types.
/// Note that this method resets the values/filled-switch array back
/// to all-zero/false while only iterating over the nonzero elements.
void expInsert(uint64_t *cursor, V *values, bool *filled, uint64_t *added,
uint64_t count) override {
if (count == 0)
return;
// Sort.
std::sort(added, added + count);
// Restore insertion path for first insert.
uint64_t rank = getRank();
uint64_t index = added[0];
cursor[rank - 1] = index;
lexInsert(cursor, values[index]);
assert(filled[index]);
values[index] = 0;
filled[index] = false;
// Subsequent insertions are quick.
for (uint64_t i = 1; i < count; i++) {
assert(index < added[i] && "non-lexicographic insertion");
index = added[i];
cursor[rank - 1] = index;
insPath(cursor, rank - 1, added[i - 1] + 1, values[index]);
assert(filled[index]);
values[index] = 0.0;
filled[index] = false;
}
}
/// Finalizes lexicographic insertions.
void endInsert() override {
if (values.empty())
@ -683,8 +734,7 @@ typedef uint64_t index_t;
#define IMPL_SPARSEVALUES(NAME, TYPE, LIB) \
void _mlir_ciface_##NAME(StridedMemRefType<TYPE, 1> *ref, void *tensor) { \
assert(ref); \
assert(tensor); \
assert(ref &&tensor); \
std::vector<TYPE> *v; \
static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v); \
ref->basePtr = ref->data = v->data(); \
@ -696,8 +746,7 @@ typedef uint64_t index_t;
#define IMPL_GETOVERHEAD(NAME, TYPE, LIB) \
void _mlir_ciface_##NAME(StridedMemRefType<TYPE, 1> *ref, void *tensor, \
index_t d) { \
assert(ref); \
assert(tensor); \
assert(ref &&tensor); \
std::vector<TYPE> *v; \
static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v, d); \
ref->basePtr = ref->data = v->data(); \
@ -710,9 +759,7 @@ typedef uint64_t index_t;
void *_mlir_ciface_##NAME(void *tensor, TYPE value, \
StridedMemRefType<index_t, 1> *iref, \
StridedMemRefType<index_t, 1> *pref) { \
assert(tensor); \
assert(iref); \
assert(pref); \
assert(tensor &&iref &&pref); \
assert(iref->strides[0] == 1 && pref->strides[0] == 1); \
assert(iref->sizes[0] == pref->sizes[0]); \
const index_t *indx = iref->data + iref->offset; \
@ -726,10 +773,11 @@ typedef uint64_t index_t;
}
#define IMPL_GETNEXT(NAME, V) \
bool _mlir_ciface_##NAME(void *tensor, StridedMemRefType<uint64_t, 1> *iref, \
bool _mlir_ciface_##NAME(void *tensor, StridedMemRefType<index_t, 1> *iref, \
StridedMemRefType<V, 0> *vref) { \
assert(tensor &&iref &&vref); \
assert(iref->strides[0] == 1); \
uint64_t *indx = iref->data + iref->offset; \
index_t *indx = iref->data + iref->offset; \
V *value = vref->data + vref->offset; \
const uint64_t isize = iref->sizes[0]; \
auto iter = static_cast<SparseTensorCOO<V> *>(tensor); \
@ -747,12 +795,32 @@ typedef uint64_t index_t;
#define IMPL_LEXINSERT(NAME, V) \
void _mlir_ciface_##NAME(void *tensor, StridedMemRefType<index_t, 1> *cref, \
V val) { \
assert(tensor &&cref); \
assert(cref->strides[0] == 1); \
uint64_t *cursor = cref->data + cref->offset; \
index_t *cursor = cref->data + cref->offset; \
assert(cursor); \
static_cast<SparseTensorStorageBase *>(tensor)->lexInsert(cursor, val); \
}
#define IMPL_EXPINSERT(NAME, V) \
void _mlir_ciface_##NAME( \
void *tensor, StridedMemRefType<index_t, 1> *cref, \
StridedMemRefType<V, 1> *vref, StridedMemRefType<bool, 1> *fref, \
StridedMemRefType<index_t, 1> *aref, index_t count) { \
assert(tensor &&cref &&vref &&fref &&aref); \
assert(cref->strides[0] == 1); \
assert(vref->strides[0] == 1); \
assert(fref->strides[0] == 1); \
assert(aref->strides[0] == 1); \
assert(vref->sizes[0] == fref->sizes[0]); \
index_t *cursor = cref->data + cref->offset; \
V *values = vref->data + vref->offset; \
bool *filled = fref->data + fref->offset; \
index_t *added = aref->data + aref->offset; \
static_cast<SparseTensorStorageBase *>(tensor)->expInsert( \
cursor, values, filled, added, count); \
}
/// Constructs a new sparse tensor. This is the "swiss army knife"
/// method for materializing sparse tensors into the computation.
///
@ -912,12 +980,21 @@ IMPL_LEXINSERT(lexInsertI32, int32_t)
IMPL_LEXINSERT(lexInsertI16, int16_t)
IMPL_LEXINSERT(lexInsertI8, int8_t)
/// Helper to insert using expansion, one per value type.
IMPL_EXPINSERT(expInsertF64, double)
IMPL_EXPINSERT(expInsertF32, float)
IMPL_EXPINSERT(expInsertI64, int64_t)
IMPL_EXPINSERT(expInsertI32, int32_t)
IMPL_EXPINSERT(expInsertI16, int16_t)
IMPL_EXPINSERT(expInsertI8, int8_t)
#undef CASE
#undef IMPL_SPARSEVALUES
#undef IMPL_GETOVERHEAD
#undef IMPL_ADDELT
#undef IMPL_GETNEXT
#undef IMPL_INSERTLEX
#undef IMPL_LEXINSERT
#undef IMPL_EXPINSERT
//===----------------------------------------------------------------------===//
//

27
mlir/test/Dialect/SparseTensor/conversion.mlir
View File

@ -441,3 +441,30 @@ func @sparse_insert(%arg0: tensor<128xf32, #SparseVector>,
return
}
// CHECK-LABEL: func @sparse_expansion()
// %[[S:.*]] = call @sparseDimSize
// %[[V:.*]] = memref.alloca(%[[S]]) : memref<?xf64>
// %[[F:.*]] = memref.alloca(%[[S]]) : memref<?xi1>
// %[[A:.*]] = memref.alloca(%[[S]]) : memref<?xindex>
// linalg.fill(%{{.*}}, %[[V]]) : f64, memref<?xf64>
// linalg.fill(%{{.*}}, %[[F]]) : i1, memref<?xi1>
// CHECK: return
func @sparse_expansion() {
%c = arith.constant 8 : index
%0 = sparse_tensor.init [%c, %c] : tensor<8x8xf64, #SparseMatrix>
%values, %filled, %added, %count = sparse_tensor.expand %0
: tensor<8x8xf64, #SparseMatrix> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
return
}
// CHECK-LABEL: func @sparse_compression(
// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>,
// CHECK: call @expInsertF64(%[[A]],
// CHECK: return
func @sparse_compression(%arg0: tensor<8x8xf64, #SparseMatrix>,
%arg1: memref<?xindex>, %arg2: memref<?xf64>, %arg3: memref<?xi1>,
%arg4: memref<?xindex>, %arg5: index) {
sparse_tensor.compress %arg0, %arg1, %arg2, %arg3, %arg4, %arg5
: tensor<8x8xf64, #SparseMatrix>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index
return
}

19
mlir/test/Dialect/SparseTensor/invalid.mlir
View File

@ -160,6 +160,25 @@ func @sparse_unannotated_insert(%arg0: tensor<128xf64>, %arg1: memref<?xindex>,
// -----
func @sparse_unannotated_expansion(%arg0: tensor<128xf64>) {
// expected-error@+1 {{expected a sparse tensor for expansion}}
%values, %filled, %added, %count = sparse_tensor.expand %arg0
: tensor<128xf64> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
return
}
// -----
func @sparse_unannotated_compression(%arg0: tensor<128xf64>, %arg1: memref<?xindex>,
%arg2: memref<?xf64>, %arg3: memref<?xi1>,
%arg4: memref<?xindex>, %arg5: index) {
// expected-error@+1 {{expected a sparse tensor for compression}}
sparse_tensor.compress %arg0, %arg1, %arg2, %arg3, %arg4, %arg5
: tensor<128xf64>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index
}
// -----
func @sparse_convert_unranked(%arg0: tensor<*xf32>) -> tensor<10xf32> {
// expected-error@+1 {{unexpected type in convert}}
%0 = sparse_tensor.convert %arg0 : tensor<*xf32> to tensor<10xf32>

30
mlir/test/Dialect/SparseTensor/roundtrip.mlir
View File

@ -149,3 +149,33 @@ func @sparse_insert(%arg0: tensor<128xf64, #SparseVector>, %arg1: memref<?xindex
sparse_tensor.lex_insert %arg0, %arg1, %arg2 : tensor<128xf64, #SparseVector>, memref<?xindex>, f64
return
}
// -----
#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
// CHECK-LABEL: func @sparse_expansion(
// CHECK-SAME: %[[A:.*]]: tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>)
// CHECK: sparse_tensor.expand %[[A]]
// CHECK: return
func @sparse_expansion(%arg0: tensor<8x8xf64, #SparseMatrix>) {
%values, %filled, %added, %count = sparse_tensor.expand %arg0
: tensor<8x8xf64, #SparseMatrix> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
return
}
// -----
#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
// CHECK-LABEL: func @sparse_compression(
// CHECK-SAME: %[[A:.*]]: tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>,
// CHECK: sparse_tensor.compress %[[A]]
// CHECK: return
func @sparse_compression(%arg0: tensor<8x8xf64, #SparseMatrix>,
%arg1: memref<?xindex>, %arg2: memref<?xf64>, %arg3: memref<?xi1>,
%arg4: memref<?xindex>, %arg5: index) {
sparse_tensor.compress %arg0, %arg1, %arg2, %arg3, %arg4, %arg5
: tensor<8x8xf64, #SparseMatrix>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index
return
}

274
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir Normal file
View File

@ -0,0 +1,274 @@
// RUN: mlir-opt %s \
// RUN: --linalg-generalize-named-ops --linalg-fuse-elementwise-ops \
// RUN: --sparsification --sparse-tensor-conversion \
// RUN: --linalg-bufferize --convert-linalg-to-loops \
// RUN: --convert-vector-to-scf --convert-scf-to-std \
// RUN: --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
// RUN: --std-bufferize --finalizing-bufferize --lower-affine \
// RUN: --convert-vector-to-llvm --convert-memref-to-llvm \
// RUN: --convert-std-to-llvm --reconcile-unrealized-casts | \
// RUN: mlir-cpu-runner \
// RUN: -e entry -entry-point-result=void \
// RUN: -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
// RUN: FileCheck %s
//
#CSR = #sparse_tensor.encoding<{
dimLevelType = [ "dense", "compressed" ],
dimOrdering = affine_map<(i,j) -> (i,j)>
}>
#DCSR = #sparse_tensor.encoding<{
dimLevelType = [ "compressed", "compressed" ],
dimOrdering = affine_map<(i,j) -> (i,j)>
}>
module {
//
// Computes C = A x B with all matrices dense.
//
func @matmul1(%A: tensor<4x8xf64>,
%B: tensor<8x4xf64>) -> tensor<4x4xf64> {
%C = arith.constant dense<0.0> : tensor<4x4xf64>
%D = linalg.matmul
ins(%A, %B: tensor<4x8xf64>, tensor<8x4xf64>)
outs(%C: tensor<4x4xf64>) -> tensor<4x4xf64>
return %D: tensor<4x4xf64>
}
//
// Computes C = A x B with all matrices sparse (SpMSpM) in CSR.
//
func @matmul2(%A: tensor<4x8xf64, #CSR>,
%B: tensor<8x4xf64, #CSR>) -> tensor<4x4xf64, #CSR> {
%c4 = arith.constant 4 : index
%C = sparse_tensor.init [%c4, %c4] : tensor<4x4xf64, #CSR>
%D = linalg.matmul
ins(%A, %B: tensor<4x8xf64, #CSR>, tensor<8x4xf64, #CSR>)
outs(%C: tensor<4x4xf64, #CSR>) -> tensor<4x4xf64, #CSR>
return %D: tensor<4x4xf64, #CSR>
}
//
// Computes C = A x B with all matrices sparse (SpMSpM) in DCSR.
//
func @matmul3(%A: tensor<4x8xf64, #DCSR>,
%B: tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> {
%c4 = arith.constant 4 : index
%C = sparse_tensor.init [%c4, %c4] : tensor<4x4xf64, #DCSR>
%D = linalg.matmul
ins(%A, %B: tensor<4x8xf64, #DCSR>, tensor<8x4xf64, #DCSR>)
outs(%C: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR>
return %D: tensor<4x4xf64, #DCSR>
}
//
// Main driver.
//
func @entry() {
%c0 = arith.constant 0 : index
%d1 = arith.constant -1.0 : f64
// Initialize various matrices, dense for stress testing,
// and sparse to verify correct nonzero structure.
%da = arith.constant dense<[
[ 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1 ],
[ 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2 ],
[ 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3, 8.3 ],
[ 1.4, 2.4, 3.4, 4.4, 5.4, 6.4, 7.4, 8.4 ]
]> : tensor<4x8xf64>
%db = arith.constant dense<[
[ 10.1, 11.1, 12.1, 13.1 ],
[ 10.2, 11.2, 12.2, 13.2 ],
[ 10.3, 11.3, 12.3, 13.3 ],
[ 10.4, 11.4, 12.4, 13.4 ],
[ 10.5, 11.5, 12.5, 13.5 ],
[ 10.6, 11.6, 12.6, 13.6 ],
[ 10.7, 11.7, 12.7, 13.7 ],
[ 10.8, 11.8, 12.8, 13.8 ]
]> : tensor<8x4xf64>
%sa = arith.constant dense<[
[ 0.0, 2.1, 0.0, 0.0, 0.0, 6.1, 0.0, 0.0 ],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 ],
[ 0.0, 2.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 ],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0 ]
]> : tensor<4x8xf64>
%sb = arith.constant dense<[
[ 0.0, 0.0, 0.0, 1.0 ],
[ 0.0, 0.0, 2.0, 0.0 ],
[ 0.0, 3.0, 0.0, 0.0 ],
[ 4.0, 0.0, 0.0, 0.0 ],
[ 0.0, 0.0, 0.0, 0.0 ],
[ 0.0, 5.0, 0.0, 0.0 ],
[ 0.0, 0.0, 6.0, 0.0 ],
[ 0.0, 0.0, 7.0, 8.0 ]
]> : tensor<8x4xf64>
// Convert all these matrices to sparse format.
%a1 = sparse_tensor.convert %da : tensor<4x8xf64> to tensor<4x8xf64, #CSR>
%a2 = sparse_tensor.convert %da : tensor<4x8xf64> to tensor<4x8xf64, #DCSR>
%a3 = sparse_tensor.convert %sa : tensor<4x8xf64> to tensor<4x8xf64, #CSR>
%a4 = sparse_tensor.convert %sa : tensor<4x8xf64> to tensor<4x8xf64, #DCSR>
%b1 = sparse_tensor.convert %db : tensor<8x4xf64> to tensor<8x4xf64, #CSR>
%b2 = sparse_tensor.convert %db : tensor<8x4xf64> to tensor<8x4xf64, #DCSR>
%b3 = sparse_tensor.convert %sb : tensor<8x4xf64> to tensor<8x4xf64, #CSR>
%b4 = sparse_tensor.convert %sb : tensor<8x4xf64> to tensor<8x4xf64, #DCSR>
// Call kernels with dense.
%0 = call @matmul1(%da, %db)
: (tensor<4x8xf64>, tensor<8x4xf64>) -> tensor<4x4xf64>
%1 = call @matmul2(%a1, %b1)
: (tensor<4x8xf64, #CSR>,
tensor<8x4xf64, #CSR>) -> tensor<4x4xf64, #CSR>
%2 = call @matmul3(%a2, %b2)
: (tensor<4x8xf64, #DCSR>,
tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR>
// Call kernels with one sparse.
%3 = call @matmul1(%sa, %db)
: (tensor<4x8xf64>, tensor<8x4xf64>) -> tensor<4x4xf64>
%4 = call @matmul2(%a3, %b1)
: (tensor<4x8xf64, #CSR>,
tensor<8x4xf64, #CSR>) -> tensor<4x4xf64, #CSR>
%5 = call @matmul3(%a4, %b2)
: (tensor<4x8xf64, #DCSR>,
tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR>
// Call kernels with sparse.
%6 = call @matmul1(%sa, %sb)
: (tensor<4x8xf64>, tensor<8x4xf64>) -> tensor<4x4xf64>
%7 = call @matmul2(%a3, %b3)
: (tensor<4x8xf64, #CSR>,
tensor<8x4xf64, #CSR>) -> tensor<4x4xf64, #CSR>
%8 = call @matmul3(%a4, %b4)
: (tensor<4x8xf64, #DCSR>,
tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR>
//
// CHECK: ( ( 388.76, 425.56, 462.36, 499.16 ),
// CHECK-SAME: ( 397.12, 434.72, 472.32, 509.92 ),
// CHECK-SAME: ( 405.48, 443.88, 482.28, 520.68 ),
// CHECK-SAME: ( 413.84, 453.04, 492.24, 531.44 ) )
//
%m0 = bufferization.to_memref %0 : memref<4x4xf64>
%v0 = vector.transfer_read %m0[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
vector.print %v0 : vector<4x4xf64>
//
// CHECK: ( ( 388.76, 425.56, 462.36, 499.16 ),
// CHECK-SAME: ( 397.12, 434.72, 472.32, 509.92 ),
// CHECK-SAME: ( 405.48, 443.88, 482.28, 520.68 ),
// CHECK-SAME: ( 413.84, 453.04, 492.24, 531.44 ) )
//
%c1 = sparse_tensor.convert %1 : tensor<4x4xf64, #CSR> to tensor<4x4xf64>
%m1 = bufferization.to_memref %c1 : memref<4x4xf64>
%v1 = vector.transfer_read %m1[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
vector.print %v1 : vector<4x4xf64>
//
// CHECK: ( ( 388.76, 425.56, 462.36, 499.16 ),
// CHECK-SAME: ( 397.12, 434.72, 472.32, 509.92 ),
// CHECK-SAME: ( 405.48, 443.88, 482.28, 520.68 ),
// CHECK-SAME: ( 413.84, 453.04, 492.24, 531.44 ) )
//
%c2 = sparse_tensor.convert %2 : tensor<4x4xf64, #DCSR> to tensor<4x4xf64>
%m2 = bufferization.to_memref %c2 : memref<4x4xf64>
%v2 = vector.transfer_read %m2[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
vector.print %v2 : vector<4x4xf64>
//
// CHECK: ( ( 86.08, 94.28, 102.48, 110.68 ),
// CHECK-SAME: ( 0, 0, 0, 0 ),
// CHECK-SAME: ( 23.46, 25.76, 28.06, 30.36 ),
// CHECK-SAME: ( 10.8, 11.8, 12.8, 13.8 ) )
//
%m3 = bufferization.to_memref %3 : memref<4x4xf64>
%v3 = vector.transfer_read %m3[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
vector.print %v3 : vector<4x4xf64>
//
// CHECK: ( ( 86.08, 94.28, 102.48, 110.68 ),
// CHECK-SAME: ( 0, 0, 0, 0 ),
// CHECK-SAME: ( 23.46, 25.76, 28.06, 30.36 ),
// CHECK-SAME: ( 10.8, 11.8, 12.8, 13.8 ) )
//
%c4 = sparse_tensor.convert %4 : tensor<4x4xf64, #CSR> to tensor<4x4xf64>
%m4 = bufferization.to_memref %c4 : memref<4x4xf64>
%v4 = vector.transfer_read %m4[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
vector.print %v4 : vector<4x4xf64>
//
// CHECK: ( ( 86.08, 94.28, 102.48, 110.68 ),
// CHECK-SAME: ( 0, 0, 0, 0 ),
// CHECK-SAME: ( 23.46, 25.76, 28.06, 30.36 ),
// CHECK-SAME: ( 10.8, 11.8, 12.8, 13.8 ) )
//
%c5 = sparse_tensor.convert %5 : tensor<4x4xf64, #DCSR> to tensor<4x4xf64>
%m5 = bufferization.to_memref %c5 : memref<4x4xf64>
%v5 = vector.transfer_read %m5[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
vector.print %v5 : vector<4x4xf64>
//
// CHECK: ( ( 0, 30.5, 4.2, 0 ), ( 0, 0, 0, 0 ), ( 0, 0, 4.6, 0 ), ( 0, 0, 7, 8 ) )
//
%m6 = bufferization.to_memref %6 : memref<4x4xf64>
%v6 = vector.transfer_read %m6[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
vector.print %v6 : vector<4x4xf64>
//
// CHECK: ( ( 0, 30.5, 4.2, 0 ), ( 0, 0, 0, 0 ), ( 0, 0, 4.6, 0 ), ( 0, 0, 7, 8 ) )
//
%c7 = sparse_tensor.convert %7 : tensor<4x4xf64, #CSR> to tensor<4x4xf64>
%m7 = bufferization.to_memref %c7 : memref<4x4xf64>
%v7 = vector.transfer_read %m7[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
vector.print %v7 : vector<4x4xf64>
//
// CHECK: ( ( 0, 30.5, 4.2, 0 ), ( 0, 0, 0, 0 ), ( 0, 0, 4.6, 0 ), ( 0, 0, 7, 8 ) )
//
%c8 = sparse_tensor.convert %8 : tensor<4x4xf64, #DCSR> to tensor<4x4xf64>
%m8 = bufferization.to_memref %c8 : memref<4x4xf64>
%v8 = vector.transfer_read %m8[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
vector.print %v8 : vector<4x4xf64>
//
// Sanity check on nonzeros.
//
// CHECK: ( 30.5, 4.2, 4.6, 7, 8, -1, -1, -1 )
// CHECK: ( 30.5, 4.2, 4.6, 7, 8, -1, -1, -1 )
//
%val7 = sparse_tensor.values %7 : tensor<4x4xf64, #CSR> to memref<?xf64>
%val8 = sparse_tensor.values %8 : tensor<4x4xf64, #DCSR> to memref<?xf64>
%nz7 = vector.transfer_read %val7[%c0], %d1 : memref<?xf64>, vector<8xf64>
%nz8 = vector.transfer_read %val8[%c0], %d1 : memref<?xf64>, vector<8xf64>
vector.print %nz7 : vector<8xf64>
vector.print %nz8 : vector<8xf64>
// Release the resources.
sparse_tensor.release %a1 : tensor<4x8xf64, #CSR>
sparse_tensor.release %a2 : tensor<4x8xf64, #DCSR>
sparse_tensor.release %a3 : tensor<4x8xf64, #CSR>
sparse_tensor.release %a4 : tensor<4x8xf64, #DCSR>
sparse_tensor.release %b1 : tensor<8x4xf64, #CSR>
sparse_tensor.release %b2 : tensor<8x4xf64, #DCSR>
sparse_tensor.release %b3 : tensor<8x4xf64, #CSR>
sparse_tensor.release %b4 : tensor<8x4xf64, #DCSR>
sparse_tensor.release %1 : tensor<4x4xf64, #CSR>
sparse_tensor.release %2 : tensor<4x4xf64, #DCSR>
sparse_tensor.release %4 : tensor<4x4xf64, #CSR>
sparse_tensor.release %5 : tensor<4x4xf64, #DCSR>
sparse_tensor.release %7 : tensor<4x4xf64, #CSR>
sparse_tensor.release %8 : tensor<4x4xf64, #DCSR>
memref.dealloc %m0 : memref<4x4xf64>
memref.dealloc %m1 : memref<4x4xf64>
memref.dealloc %m2 : memref<4x4xf64>
memref.dealloc %m3 : memref<4x4xf64>
memref.dealloc %m4 : memref<4x4xf64>
memref.dealloc %m5 : memref<4x4xf64>
memref.dealloc %m6 : memref<4x4xf64>
memref.dealloc %m7 : memref<4x4xf64>
memref.dealloc %m8 : memref<4x4xf64>
return
}
}