Skip to content

[mlir][sparse] unify block arguments order between iterate/coiterate operations. #105567

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 23, 2024
Merged

[mlir][sparse] unify block arguments order between iterate/coiterate operations. #105567

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 3 additions & 4 deletions mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1644,7 +1644,7 @@ def IterateOp : SparseTensor_Op<"iterate",
return getIterSpace().getType().getSpaceDim();
}
BlockArgument getIterator() {
return getRegion().getArguments().front();
return getRegion().getArguments().back();
}
std::optional<BlockArgument> getLvlCrd(Level lvl) {
if (getCrdUsedLvls()[lvl]) {
Expand All @@ -1654,9 +1654,8 @@ def IterateOp : SparseTensor_Op<"iterate",
return std::nullopt;
}
Block::BlockArgListType getCrds() {
// The first block argument is iterator, the remaining arguments are
// referenced coordinates.
return getRegion().getArguments().slice(1, getCrdUsedLvls().count());
// User-provided iteration arguments -> coords -> iterator.
return getRegion().getArguments().slice(getNumRegionIterArgs(), getCrdUsedLvls().count());
}
unsigned getNumRegionIterArgs() {
return getRegion().getArguments().size() - 1 - getCrdUsedLvls().count();
Expand Down
31 changes: 17 additions & 14 deletions mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2228,16 +2228,19 @@ parseSparseIterateLoop(OpAsmParser &parser, OperationState &state,
parser.getNameLoc(),
"mismatch in number of sparse iterators and sparse spaces");

if (failed(parseUsedCoordList(parser, state, blockArgs)))
SmallVector<OpAsmParser::Argument> coords;
if (failed(parseUsedCoordList(parser, state, coords)))
return failure();
size_t numCrds = blockArgs.size();
size_t numCrds = coords.size();

// Parse "iter_args(%arg = %init, ...)"
bool hasIterArgs = succeeded(parser.parseOptionalKeyword("iter_args"));
if (hasIterArgs)
if (parser.parseAssignmentList(blockArgs, initArgs))
return failure();

blockArgs.append(coords);

SmallVector<Type> iterSpaceTps;
// parse ": sparse_tensor.iter_space -> ret"
if (parser.parseColon() || parser.parseTypeList(iterSpaceTps))
Expand Down Expand Up @@ -2267,7 +2270,7 @@ parseSparseIterateLoop(OpAsmParser &parser, OperationState &state,

if (hasIterArgs) {
// Strip off leading args that used for coordinates.
MutableArrayRef args = MutableArrayRef(blockArgs).drop_front(numCrds);
MutableArrayRef args = MutableArrayRef(blockArgs).drop_back(numCrds);
if (args.size() != initArgs.size() || args.size() != state.types.size()) {
return parser.emitError(
parser.getNameLoc(),
Expand Down Expand Up @@ -2448,18 +2451,18 @@ void IterateOp::build(OpBuilder &builder, OperationState &odsState,
odsState.addTypes(initArgs.getTypes());
Block *bodyBlock = builder.createBlock(bodyRegion);

// First argument, sparse iterator
bodyBlock->addArgument(
llvm::cast<IterSpaceType>(iterSpace.getType()).getIteratorType(),
odsState.location);
// Starts with a list of user-provided loop arguments.
for (Value v : initArgs)
bodyBlock->addArgument(v.getType(), v.getLoc());

// Followed by a list of used coordinates.
// Follows by a list of used coordinates.
for (unsigned i = 0, e = crdUsedLvls.count(); i < e; i++)
bodyBlock->addArgument(builder.getIndexType(), odsState.location);

// Followed by a list of user-provided loop arguments.
for (Value v : initArgs)
bodyBlock->addArgument(v.getType(), v.getLoc());
// Ends with sparse iterator
bodyBlock->addArgument(
llvm::cast<IterSpaceType>(iterSpace.getType()).getIteratorType(),
odsState.location);
}

ParseResult IterateOp::parse(OpAsmParser &parser, OperationState &result) {
Expand All @@ -2473,9 +2476,9 @@ ParseResult IterateOp::parse(OpAsmParser &parser, OperationState &result) {
return parser.emitError(parser.getNameLoc(),
"expected only one iterator/iteration space");

iters.append(iterArgs);
iterArgs.append(iters);
Region *body = result.addRegion();
if (parser.parseRegion(*body, iters))
if (parser.parseRegion(*body, iterArgs))
return failure();

IterateOp::ensureTerminator(*body, parser.getBuilder(), result.location);
Expand Down Expand Up @@ -2580,7 +2583,7 @@ MutableArrayRef<OpOperand> IterateOp::getInitsMutable() {
}

Block::BlockArgListType IterateOp::getRegionIterArgs() {
return getRegion().getArguments().take_back(getNumRegionIterArgs());
return getRegion().getArguments().take_front(getNumRegionIterArgs());
}

std::optional<MutableArrayRef<OpOperand>> IterateOp::getYieldedValuesMutable() {
Expand Down
36 changes: 11 additions & 25 deletions mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ genCoIterateBranchNest(PatternRewriter &rewriter, Location loc, CoIterateOp op,

static ValueRange genLoopWithIterator(
PatternRewriter &rewriter, Location loc, SparseIterator *it,
ValueRange reduc, bool iterFirst,
ValueRange reduc,
function_ref<SmallVector<Value>(PatternRewriter &rewriter, Location loc,
Region &loopBody, SparseIterator *it,
ValueRange reduc)>
Expand All @@ -138,15 +138,9 @@ static ValueRange genLoopWithIterator(
}
return forOp.getResults();
}
SmallVector<Value> ivs;
// TODO: always put iterator SSA values at the end of argument list to be
// consistent with coiterate operation.
if (!iterFirst)
llvm::append_range(ivs, it->getCursor());
// Appends the user-provided values.
llvm::append_range(ivs, reduc);
if (iterFirst)
llvm::append_range(ivs, it->getCursor());

SmallVector<Value> ivs(reduc);
llvm::append_range(ivs, it->getCursor());

TypeRange types = ValueRange(ivs).getTypes();
auto whileOp = rewriter.create<scf::WhileOp>(loc, types, ivs);
Expand All @@ -164,25 +158,17 @@ static ValueRange genLoopWithIterator(
Region &dstRegion = whileOp.getAfter();
Block *after = rewriter.createBlock(&dstRegion, {}, types, l);
ValueRange aArgs = whileOp.getAfterArguments();
if (iterFirst) {
aArgs = it->linkNewScope(aArgs);
} else {
aArgs = aArgs.take_front(reduc.size());
it->linkNewScope(aArgs.drop_front(reduc.size()));
}
it->linkNewScope(aArgs.drop_front(reduc.size()));
aArgs = aArgs.take_front(reduc.size());

rewriter.setInsertionPointToStart(after);
SmallVector<Value> ret = bodyBuilder(rewriter, loc, dstRegion, it, aArgs);
rewriter.setInsertionPointToEnd(after);

// Forward loops
SmallVector<Value> yields;
ValueRange nx = it->forward(rewriter, loc);
if (iterFirst)
llvm::append_range(yields, nx);
llvm::append_range(yields, ret);
if (!iterFirst)
llvm::append_range(yields, nx);
llvm::append_range(yields, it->forward(rewriter, loc));
rewriter.create<scf::YieldOp>(loc, yields);
}
return whileOp.getResults().drop_front(it->getCursor().size());
Expand Down Expand Up @@ -258,13 +244,13 @@ class SparseIterateOpConverter : public OneToNOpConversionPattern<IterateOp> {

Block *block = op.getBody();
ValueRange ret = genLoopWithIterator(
rewriter, loc, it.get(), ivs, /*iterFirst=*/true,
rewriter, loc, it.get(), ivs,
[block](PatternRewriter &rewriter, Location loc, Region &loopBody,
SparseIterator *it, ValueRange reduc) -> SmallVector<Value> {
SmallVector<Value> blockArgs(it->getCursor());
SmallVector<Value> blockArgs(reduc);
// TODO: Also appends coordinates if used.
// blockArgs.push_back(it->deref(rewriter, loc));
llvm::append_range(blockArgs, reduc);
llvm::append_range(blockArgs, it->getCursor());

Block *dstBlock = &loopBody.getBlocks().front();
rewriter.inlineBlockBefore(block, dstBlock, dstBlock->end(),
Expand Down Expand Up @@ -404,7 +390,7 @@ class SparseCoIterateOpConverter

Block *block = &r.getBlocks().front();
ValueRange curResult = genLoopWithIterator(
rewriter, loc, validIters.front(), userReduc, /*iterFirst=*/false,
rewriter, loc, validIters.front(), userReduc,
/*bodyBuilder=*/
[block](PatternRewriter &rewriter, Location loc, Region &dstRegion,
SparseIterator *it,
Expand Down
Loading