[mlir][sparse] support SoA COO in codegen path. (#82439)

*NOTE*: the `SoA` property only makes a difference on codegen path, and
is ignored in libgen path at the moment (only SoA COO is supported).
This commit is contained in:
Peiming Liu
2024-02-20 19:06:21 -06:00
committed by GitHub
parent 4ca0480a4f
commit 5248a98724
9 changed files with 43 additions and 30 deletions

View File

@@ -194,7 +194,7 @@ static void createAllocFields(OpBuilder &builder, Location loc,
valHeuristic =
builder.create<arith::MulIOp>(loc, valHeuristic, lvlSizesValues[lvl]);
} else if (sizeHint) {
if (stt.getCOOStart() == 0) {
if (stt.getAoSCOOStart() == 0) {
posHeuristic = constantIndex(builder, loc, 2);
crdHeuristic = builder.create<arith::MulIOp>(
loc, constantIndex(builder, loc, lvlRank), sizeHint); // AOS
@@ -1316,7 +1316,7 @@ struct SparseAssembleOpConverter : public OpConversionPattern<AssembleOp> {
Value posBack = c0; // index to the last value in the position array
Value memSize = c1; // memory size for current array
Level trailCOOStart = stt.getCOOStart();
Level trailCOOStart = stt.getAoSCOOStart();
Level trailCOORank = stt.getLvlRank() - trailCOOStart;
// Sets up SparseTensorSpecifier.
for (Level lvl = 0, lvlRank = stt.getLvlRank(); lvl < lvlRank; lvl++) {
@@ -1453,7 +1453,7 @@ struct SparseNewConverter : public OpConversionPattern<NewOp> {
const auto dstTp = getSparseTensorType(op.getResult());
// Creating COO with NewOp is handled by direct IR codegen. All other cases
// are handled by rewriting.
if (!dstTp.hasEncoding() || dstTp.getCOOStart() != 0)
if (!dstTp.hasEncoding() || dstTp.getAoSCOOStart() != 0)
return failure();
// Implement as follows: