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

[PeimingLiu]

Stacked PRs:

• ->[mlir][sparse] unify block arguments order between iterate/coiterate operations. #105567
• [mlir][sparse] refactoring sparse_tensor.iterate lowering pattern implementation. #105566
• [mlir][sparse] partially support lowering sparse coiteration loops to scf.while/for. #105565

---

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

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-sparse

Author: Peiming Liu (PeimingLiu)

Changes

Stacked PRs:

• ->#105567
• #105566
• #105565

---

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

---

Full diff: https://github.com/llvm/llvm-project/pull/105567.diff

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td (+3-4)
• (modified) mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp (+17-14)
• (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp (+11-25)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
index 20512f972e67cd..96a61419a541f7 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
@@ -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]) {
@@ -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();
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 16856b958d4f13..b21bc1a93036c4 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -2228,9 +2228,10 @@ 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"));
@@ -2238,6 +2239,8 @@ parseSparseIterateLoop(OpAsmParser &parser, OperationState &state,
     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))
@@ -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(),
@@ -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) {
@@ -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);
@@ -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() {
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp
index f7fcabb0220b50..71a229bea990c0 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp
@@ -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)>
@@ -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);
@@ -164,12 +158,8 @@ 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);
@@ -177,12 +167,8 @@ static ValueRange genLoopWithIterator(
 
     // 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());
@@ -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(),
@@ -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,