address comments.

Author: Peiming Liu

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

@@ -1769,18 +1769,18 @@ def SparseTensor_CoIterateOp : SparseTensor_Op<"coiterate",
       });
     }
 
-    // The block arguments starts with referenced coordinates, follows by
-    // user-provided iteration arguments and ends with iterators.
+    // The block arguments starts with user-provided iteration arguments,
+    // follows by referenced coordinates and ends with iterators.
     Block::BlockArgListType getCrds(unsigned regionIdx) {
       return getRegion(regionIdx).getArguments()
-          .take_front(getCrdUsedLvls().count());
+          .slice(getNumRegionIterArgs(), getCrdUsedLvls().count());
     }
-    unsigned getNumRegionIterArgs(unsigned regionIdx) {
+    unsigned getNumRegionIterArgs() {
       return getInitArgs().size();
     }
     Block::BlockArgListType getRegionIterArgs(unsigned regionIdx) {
       return getRegion(regionIdx).getArguments()
-          .slice(getCrdUsedLvls().count(), getNumRegionIterArgs(regionIdx));
+          .take_front(getNumRegionIterArgs());
     }
     Block::BlockArgListType getRegionIterators(unsigned regionIdx) {
       return getRegion(regionIdx).getArguments()

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

@@ -2293,16 +2293,18 @@ parseSparseCoIterateLoop(OpAsmParser &parser, OperationState &state,
   if (parser.parseOperandList(spaces, OpAsmParser::Delimiter::Paren))
     return failure();
 
-  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, ...)"
   SmallVector<OpAsmParser::UnresolvedOperand> initArgs;
   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"
@@ -2326,8 +2328,8 @@ parseSparseCoIterateLoop(OpAsmParser &parser, OperationState &state,
   state.operands.append(spacesVals);
 
   if (hasIterArgs) {
-    // Strip off leading args that used for coordinates.
-    MutableArrayRef args = MutableArrayRef(blockArgs).drop_front(numCrds);
+    // Strip off trailing args that used for coordinates.
+    MutableArrayRef args = MutableArrayRef(blockArgs).drop_back(numCrds);
     if (args.size() != initArgs.size() || args.size() != state.types.size()) {
       return parser.emitError(
           parser.getNameLoc(),
@@ -2609,8 +2611,10 @@ void CoIterateOp::build(OpBuilder &builder, OperationState &odsState,
       cast<IterSpaceType>(iterSpaces.front().getType()).getSpaceDim();
   // All ones.
   I64BitSet set((1 << rank) - 1);
-  // Fake cases bits. We need to preallocate all the regions as Region can not
-  // be dynamically added later after the operation is created.
+  // Generates all-zero case bits (they only serve as placeholders), which are
+  // supposed to be overriden later. We need to preallocate all the regions as
+  // mlir::Region cannot be dynamically added later after the operation is
+  // created.
   SmallVector<int64_t> caseBits(numCases, 0);
   ArrayAttr cases = builder.getI64ArrayAttr(caseBits);
   return CoIterateOp::build(builder, odsState, initArgs.getTypes(), iterSpaces,
@@ -2701,7 +2705,7 @@ ValueRange CoIterateOp::getYieldedValues(unsigned regionIdx) {
 
 LogicalResult CoIterateOp::verifyRegions() {
   for (unsigned r = 0, e = getNumRegions(); r < e; r++) {
-    if (getNumRegionIterArgs(r) != getNumResults())
+    if (getNumRegionIterArgs() != getNumResults())
       return emitOpError(
           "mismatch in number of basic block args and defined values");
 

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

@@ -904,7 +904,7 @@ static void finalizeWhileOp(CodegenEnv &env, OpBuilder &builder,
   // basic block where scf::Yield should be inserted.
 }
 
-/// Generate a case region in the coiterate operation.
+/// Generates a case region in the coiterate operation.
 static void genCoIterationCase(CodegenEnv &env, OpBuilder &builder,
                                unsigned caseIdx, LatPointId allCase,
                                LatPointId curCase,
@@ -920,8 +920,8 @@ static void genCoIterationCase(CodegenEnv &env, OpBuilder &builder,
     if (curCaseBits.test(set))
       caseBit.set(idx);
 
-  env.emitter().enterCurCoIterationCase(builder, env.op().getLoc(), caseBit,
-                                        caseIdx, reduc);
+  env.emitter().enterCurrentCoIterationCase(builder, env.op().getLoc(), caseBit,
+                                            caseIdx, reduc);
 }
 
 /// Generates a single if-statement within a while-loop.

mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.cpp

@@ -615,10 +615,11 @@ bool LoopEmitter::shouldIteratedByForLoop(ArrayRef<SparseIterator *> spIters) {
   return true;
 }
 
-Region *LoopEmitter::enterCurCoIterationCase(OpBuilder &builder, Location loc,
-                                             I64BitSet caseBit,
-                                             unsigned caseIdx,
-                                             MutableArrayRef<Value> reduc) {
+Region *LoopEmitter::enterCurrentCoIterationCase(OpBuilder &builder,
+                                                 Location loc,
+                                                 I64BitSet caseBit,
+                                                 unsigned caseIdx,
+                                                 MutableArrayRef<Value> reduc) {
   auto coIterOp = cast<CoIterateOp>(loopStack.back().loop);
   SmallVector<Attribute> cases(coIterOp.getCases().getAsRange<Attribute>());
   cases[caseIdx] = builder.getI64IntegerAttr(caseBit);
@@ -628,11 +629,12 @@ Region *LoopEmitter::enterCurCoIterationCase(OpBuilder &builder, Location loc,
   assert(caseRegion.getBlocks().empty() &&
          "re-initialize the same coiteration case region.");
 
-  // Each block starts with a list of used coordinates of index type.
+  // Each block starts with by a list of user-provided iteration arguments.
+  TypeRange iterArgsTps = coIterOp.getInitArgs().getTypes();
+  // Followed by a list of used coordinates of index type.
   SmallVector<Type> blockArgTps(coIterOp.getCrdUsedLvls().count(),
                                 builder.getIndexType());
-  // Follows by a list of user-provided iteration arguments.
-  TypeRange iterArgsTps = coIterOp.getInitArgs().getTypes();
+
   blockArgTps.append(iterArgsTps.begin(), iterArgsTps.end());
   // Ends with a set of iterators that defines the actually iteration space.
   for (auto i : caseBit.bits()) {

mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h

@@ -148,9 +148,9 @@ class LoopEmitter {
       unsigned numCases, MutableArrayRef<Value> reduc = {},
       bool isParallel = false, bool needsUniv = false);
 
-  Region *enterCurCoIterationCase(OpBuilder &builder, Location loc,
-                                  I64BitSet caseBit, unsigned caseIdx,
-                                  MutableArrayRef<Value> reduc);
+  Region *enterCurrentCoIterationCase(OpBuilder &builder, Location loc,
+                                      I64BitSet caseBit, unsigned caseIdx,
+                                      MutableArrayRef<Value> reduc);
 
   /// Generates code to exit the current loop (e.g., generates yields, forwards
   /// loop induction variables, etc).

mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir

@@ -1,6 +1,6 @@
 // RUN: mlir-opt %s --sparse-reinterpret-map -sparsification="sparse-emit-strategy=sparse-iterator" --cse | FileCheck %s --check-prefix="ITER"
 
-// TODO: temporially disabled since there is no lowering rules from `coiterate` to `scf`.
+// TODO: temporarilly disabled since there is no lowering rules from `coiterate` to `scf`.
 // R_U_N: mlir-opt %s --sparse-reinterpret-map -sparsification="sparse-emit-strategy=sparse-iterator" --cse --sparse-space-collapse --lower-sparse-iteration-to-scf --loop-invariant-code-motion | FileCheck %s