[mlir][sparse] Refactoring sparsification using loop emitter.

[Peiming]

fixed some check tests...

[wrengr]

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

173–179

This whole section is very confusing due to the variable names. Is dimType supposed to be a DimLevelType? If so, then why do you use is{Compressed,Singleton}Dim instead of is{Compressed,Singleton}DLT? and why is the field called dims rather than something like dimLevelTypes which would make it clear? If not, then what is dimType supposed to be?

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

297

I don't feel like this class really belongs in this file— since all the other things in CodegenUtils.{h,cpp} are small generic helpers, whereas this class is both rather large and pretty specific. Why not split it out into a separate LoopEmitter.{h,cpp} file?

324–335

Why have this as an inline definition rather than moving it to the cpp file like most of the other methods?

339–344

Ditto. While this one is pretty short, I'm don't imagine there's any benefit to having it be inline.

444

Rename ts to tensors and ds to dims.

447–450

Can any of these be marked const?

454–466

Move this documentation to just above the struct LoopSeqInfo { rather than having it inside the class definition

469

Constructors should be moved to the top of the class / fields should be at the bottom of the class

473–475

If this comment is for loopStack then it should be placed above the loopStack declaration. Whereas, if this comment is for loopSeqStack then it should have a blank line separating the comment from the loopStack declaration and shouldn't have any blank lines between the comment and the loopSeqStack declaration.

[Peiming]

update check tests..

[Peiming]

remove unsupported vector tests...

[Peiming]

rebase

[Peiming]

address comments from Wren.

[Peiming]

mark varaible as const.

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

173–179

I think isCompressedDim looks good to me, I will leave the decision to Aart.

I will change the dims though! Thanks!

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

297

Sure, we can have some discussion on this!

324–335

Will do

447–450

You are right!

469

Deleted...

[Peiming]

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

444

I changed ts -> tids ( because there is a tensors field in the loop emitter that causes conflicts) hope it looks good to you ;-)

[Peiming]

add comments.

[Peiming]

cleanup..

[Peiming]

rebase..

[aartbik]

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

251

something went wrong with the rebase? this does not seem related?

269

you moved the class and changed it. perhaps this is a better place, but that makes the diff either harder to read.

can we first modify, and then later (once it is in), move the class, that makes it easier to read

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

196

This seems a functional change not really required for this revision, right?
perhaps you can leave that for later, since this is already rather hard to review?

408

a lot of these changes are because we removed vectorization

as discussed, that is okay, but this is something that can be done in a separate revision, right before this one, just so that the amount of changes per revision goes down

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir

7

Here and below, for all integration tests (not the CHECK tests) that are run without and with vectorization, simply remove the L7-L13 part. I will be easy to add that back later, so nothing is lost by removing it (in fact, at one point, we may want to run *all* integration tests with and without vectorization by some other means)

[Peiming]

splitted into

https://reviews.llvm.org/D136185
https://reviews.llvm.org/D136183

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

public:
   unsigned tensor(unsigned b) const { return b % numTensors; }
   unsigned index(unsigned b) const { return b / numTensors; }
 
   /// Returns true if bit corresponds to index of output tensor.
   bool isOutTensor(unsigned b, unsigned i) const {
     return tensor(b) == outTensor && index(b) == i;
   }
 
+  unsigned getOutTensorID() const { return outTensor; }
+
+  unsigned getSynTensorID() const { return outTensor + 1; }
+
   /// Returns true if given tensor iterates *only* in the given tensor
   /// expression. For the output tensor, this defines a "simply dynamic"
   /// operation [Bik96]. For instance: a(i) *= 2.0 or a(i) += a(i) for
   /// sparse vector a.
   bool isSingleCondition(unsigned t, unsigned e) const;
 
   /// Returns true if bit corresponds to given dimension level type.
   bool isDimLevelType(unsigned b, DimLvlType tp) const {

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

 
 namespace sparse_tensor {
 
 /// Shorthand aliases for the `emitCInterface` argument to `getFunc()`,
 /// `createFuncCall()`, and `replaceOpWithFuncCall()`.
 enum class EmitCInterface : bool { Off = false, On = true };
 
 //===----------------------------------------------------------------------===//
-// SparseTensorLoopEmiter class, manages sparse tensors and helps to generate
-// loop structure to (co-iterate) sparse tensors.
-//
-// An example usage:
-// To generate following loops over T1<?x?> and T2<?x?>
-//
-// for i in T1[0] {
-//   for j : T2[0] {
-//     for k : T1[1] {}
-//     for k : T2[1] {}
-//   }
-// }
-//
-// One can use
-//
-// SparseTensorLoopEmiter loopEmiter({T1, T1});
-// loopEmiter.initializeLoopEmit();
-// loopEmiter.enterLoopOverTensorAtDim(T1, 0);
-// loopEmiter.enterLoopOverTensorAtDim(T2, 0);
-// loopEmiter.enterLoopOverTensorAtDim(T1, 1);
-// loopEmiter.exitCurrentLoop();
-// loopEmiter.enterLoopOverTensorAtDim(T2, 1);
-// for 0 -> 3:
-//    loopEmiter.exitCurrentLoop();
-//===----------------------------------------------------------------------===//
-
-// TODO: Sparsification should also rely on this class to generate loops.
-class SparseTensorLoopEmitter {
-public:
-  /// Constructor: take an array of tensors inputs, on which the generated loops
-  /// will iterate on. The index of the tensor in the array is also the
-  /// tensor id (tid) used in related functions.
-  explicit SparseTensorLoopEmitter(ValueRange tensors,
-                                   bool isLastOutput = false);
-
-  ///
-  /// Core functions.
-  ///
-
-  /// Starts a loop emitting session:
-  /// 1. Generates all the buffers needed to iterate tensors.
-  /// 2. Generates the lo/hi bounds to iterate tensors[0].
-  void initializeLoopEmit(OpBuilder &builder, Location loc);
-
-  // TODO: Gets rid of `dim` in the argument list? Track the dimension we
-  // are currently at internally. Then it would be enterNextDimForTensor.
-
-  /// Emits loop over tensor[dim], it assumes that loops between
-  /// tensor[0...dim - 1] have already been generated.
-  /// It also prepares to enter tensor[dim + 1].
-  Operation *enterLoopOverTensorAtDim(OpBuilder &builder, Location loc,
-                                      size_t tid, size_t dim,
-                                      ArrayRef<Value> reduc = {});
-
-  /// Emits a coiteration loop over a set of tensors.
-  // TODO: not yet implemented
-  void enterCoiterationOverTensorsAtDims(OpBuilder &builder, Location loc,
-                                         ArrayRef<size_t> ts,
-                                         ArrayRef<size_t> ds);
-
-  /// Emits extra locals, since the locals might not be in simplified lattices
-  /// point used to generate the loops, but are still required to generates
-  /// expressions.
-  Value emitExtraLocalsForTensorsAtDims(OpBuilder &builder, Location loc,
-                                        size_t tid, size_t dim);
-
-  void exitCurrentLoop();
-
-  /// Return the array of coordinate for all the loop generated till now.
-  void getCoordinateArray(SmallVectorImpl<Value> &coords) {
-    for (auto &l : loopStack)
-      coords.push_back(l.idx);
-  }
-
-  ///
-  /// Getters.
-  ///
-
-  Value getTensorValueBuffer(size_t tid) { return valBuffer[tid]; }
-  Value getLastLevelTensorPointerIndex(size_t tid) {
-    return pidxs[tid].back();
-  };
-
-private:
-  struct LoopLevelInfo {
-    LoopLevelInfo(ArrayRef<size_t> ts, ArrayRef<size_t> ds, Value idx)
-        : tensors(ts), dims(ds), idx(idx) {}
-    llvm::SmallVector<size_t, 4> tensors;
-    llvm::SmallVector<size_t, 4> dims;
-    Value idx;
-  };
-
-  /// Return false if tid[dim] is a dense dimension that does not need to be
-  /// prepared (to be used by sparsification for needUniv).
-  bool prepareLoopOverTensorAtDim(OpBuilder &builder, Location loc, size_t tid,
-                                  size_t dim);
-
-  /// Input (TODO: and output) tensors.
-  std::vector<Value> tensors;
-  /// The dim type array for each tensor.
-  std::vector<std::vector<SparseTensorEncodingAttr::DimLevelType>> dims;
-  /// Sparse iteration information (by tensor and dim). These arrays
-  /// are updated to remain current within the current loop.
-  std::vector<std::vector<Value>> pidxs;
-  std::vector<std::vector<Value>> coord;
-  std::vector<std::vector<Value>> highs;
-  /// Universal dense indices and upper bounds (by index). The sizes array is
-  /// set once with the inferred dimension sizes.
-  std::vector<std::vector<Value>> sizes;
-  std::vector<std::vector<Value>> ptrBuffer; // to_pointers
-  std::vector<std::vector<Value>> idxBuffer; // to_indices
-  std::vector<Value> valBuffer;              // to_value
-
-  bool isLastOutput; // Is the last tensor output tensor
-  std::vector<LoopLevelInfo> loopStack;
-  // TODO: not yet used, it should track the current level for each tensor
-  // to help eliminate `dim` paramters from above APIs.
-  std::vector<size_t> curLv;
-};
-
-//===----------------------------------------------------------------------===//
 // ExecutionEngine/SparseTensorUtils helper functions.
 //===----------------------------------------------------------------------===//
 
 /// Converts an overhead storage bitwidth to its internal type-encoding.
 OverheadType overheadTypeEncoding(unsigned width);
 
 /// Converts an overhead storage type to its internal type-encoding.
 OverheadType overheadTypeEncoding(Type tp);

 /// Generates a constant of the internal dimension level type encoding.
 inline Value
 constantDimLevelTypeEncoding(OpBuilder &builder, Location loc,
                              SparseTensorEncodingAttr::DimLevelType dlt) {
   return constantI8(builder, loc,
                     static_cast<uint8_t>(dimLevelTypeEncoding(dlt)));
 }
 
+/// Computes the shape of destination tensor of a reshape operator. This is only
+/// used when operands have dynamic shape. The shape of the destination is
+/// stored into dstShape.
+void genReshapeDstShape(Location loc, PatternRewriter &rewriter,
+                        SmallVector<Value, 4> &dstShape,
+                        ArrayRef<Value> srcShape,
+                        ArrayRef<int64_t> staticDstShape,
+                        ArrayRef<ReassociationIndices> reassociation);
+
+/// Helper method to translate indices during a reshaping operation.
+void translateIndicesArray(OpBuilder &builder, Location loc,
+                           ArrayRef<ReassociationIndices> reassociation,
+                           ValueRange srcIndices, ArrayRef<Value> srcShape,
+                           ArrayRef<Value> dstShape,
+                           SmallVectorImpl<Value> &dstIndices);
+
+inline bool isZeroRankedTensorOrScalar(Type type) {
+  auto rtp = type.dyn_cast<RankedTensorType>();
+  return !rtp || rtp.getRank() == 0;
+}
+
+//===----------------------------------------------------------------------===//
+// SparseTensorLoopEmiter class, manages sparse tensors and helps to generate
+// loop structure to (co-iterate) sparse tensors.
+//
+// An example usage:
+// To generate following loops over T1<?x?> and T2<?x?>
+//
+// for i in T1[0] {
+//   for j : T2[0] {
+//     for k : T1[1] {}
+//     for k : T2[1] {}
+//   }
+// }
+//
+// One can use
+//
+// SparseTensorLoopEmiter loopEmiter({T1, T1});
+// loopEmiter.initializeLoopEmit();
+// loopEmiter.enterLoopOverTensorAtDim(T1, 0);
+// loopEmiter.enterLoopOverTensorAtDim(T2, 0);
+// loopEmiter.enterLoopOverTensorAtDim(T1, 1);
+// loopEmiter.exitCurrentLoop();
+// loopEmiter.enterLoopOverTensorAtDim(T2, 1);
+// for 0 -> 3:
+//    loopEmiter.exitCurrentLoop();
+//===----------------------------------------------------------------------===//
+
+// TODO: Sparsification should also rely on this class to generate loops.
+class SparseTensorLoopEmitter {
+public:
+  /// Optional callback function to setup dense output tensors when initialize
+  /// the loop emitter (e.g.,to fill a dense output with zeros).
+  using OutputUpdater = function_ref<Value(OpBuilder &builder, Location loc,
+                                           Value memref, Value tensor)>;
+
+  /// Constructor: take an array of tensors inputs, on which the generated loops
+  /// will iterate on. The index of the tensor in the array is also the
+  /// tensor id (tid) used in related functions.
+  /// If isSparseOut is set, loop emitter assume that the sparse output tensor
+  /// is empty, and will always generate loops on it based on the dim sizes.
+  explicit SparseTensorLoopEmitter(ValueRange tensors,
+                                   bool isLastOutput = false,
+                                   bool isSparseOut = false);
+
+  ///
+  /// Core functions.
+  ///
+
+  /// Starts a loop emitting session by generating all the buffers needed to
+  /// iterate tensors.
+  void initializeLoopEmit(OpBuilder &builder, Location loc,
+                          OutputUpdater updater = nullptr);
+
+  /// Enters a new loop sequence, the loops within the same sequence starts from
+  /// the break points of previous loop instead of starting over from 0.
+  /// e.g.,
+  /// {
+  ///   // loop sequence start.
+  ///   p0 = while(xxx)
+  ///     ...
+  ///     break p0
+  ///
+  ///   for (p0 -> end) // see how the loop starts from p0 instead of 0.
+  ///     ...
+  ///   // loop sequence end.
+  /// }
+  void enterNewLoopSeq(OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
+                       ArrayRef<size_t> dims);
+
+  void exitCurrentLoopSeq() {
+    assert(loopSeqStack.size() == loopStack.size() + 1);
+    loopSeqStack.pop_back();
+  }
+
+  // TODO: Gets rid of `dim` in the argument list? Track the dimension we
+  // are currently at internally. Then it would be enterNextDimForTensor.
+  // Still need a way to specify the dim for non annoated dense tensor though,
+  // as it can be accessed out of order.
+  /// Emits loop over tensor[dim], it assumes that loops between
+  /// tensor[0...dim - 1] have already been generated.
+  /// It also prepares to enter tensor[dim + 1].
+  /// The function will also perform in-place update on the `reduc` vector to
+  /// return the reduction variable used inside the generated loop.
+  Operation *enterLoopOverTensorAtDim(OpBuilder &builder, Location loc,
+                                      size_t tid, size_t dim,
+                                      SmallVectorImpl<Value> &reduc,
+                                      bool isParallel = false,
+                                      ArrayRef<size_t> extraTids = {},
+                                      ArrayRef<size_t> extraDims = {});
+
+  /// Emits a coiteration loop over a set of tensors.
+  Operation *enterCoiterationOverTensorsAtDims(
+      OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
+      ArrayRef<size_t> dims, SmallVectorImpl<Value> &reduc, bool needsUniv,
+      ArrayRef<size_t> extraTids = {}, ArrayRef<size_t> extraDims = {});
+
+  SmallVector<Value, 2> exitCurrentLoop(OpBuilder &builder, Location loc,
+                                        ArrayRef<Value> reduc = {});
+
+  /// Return the array of coordinate for all the loop generated till now.
+  void getCoordinateArray(SmallVectorImpl<Value> &coords) const {
+    for (auto &l : loopStack)
+      coords.push_back(l.iv);
+  }
+
+  /// Get loop induction variable at the given level.
+  Value getLoopIV(size_t level) const {
+    if (level < loopStack.size())
+      return loopStack[level].iv;
+    return nullptr;
+  }
+
+  ///
+  /// Getters.
+  ///
+  const std::vector<std::vector<Value>> &getPidxs() const { return pidxs; };
+  const std::vector<std::vector<Value>> &getCoord() const { return coord; };
+  const std::vector<std::vector<Value>> &getHighs() const { return highs; };
+  const std::vector<std::vector<Value>> &getPtrBuffer() const {
+    return ptrBuffer;
+  };
+  const std::vector<std::vector<Value>> &getIdxBuffer() const {
+    return idxBuffer;
+  };
+  const std::vector<Value> &getValBuffer() const { return valBuffer; };
+
+private:
+  /// Input (TODO: and output) tensors.
+  std::vector<Value> tensors;
+  /// The dim type array for each tensor.
+  std::vector<std::vector<SparseTensorEncodingAttr::DimLevelType>> dimTypes;
+  /// Sparse iteration information (by tensor and dim). These arrays
+  /// are updated to remain current within the current loop.
+  std::vector<std::vector<Value>> pidxs;
+  std::vector<std::vector<Value>> coord;
+  std::vector<std::vector<Value>> highs;
+  /// Universal dense indices and upper bounds (by index). The sizes array is
+  /// set once with the inferred dimension sizes.
+  std::vector<std::vector<Value>> ptrBuffer; // to_pointers
+  std::vector<std::vector<Value>> idxBuffer; // to_indices
+  std::vector<Value> valBuffer;              // to_value
+
+  /// Setups [lo, hi] for iterating tensor[dim], it assumes that tensor[0
+  /// ...dims-1] has already been setup.
+  void prepareLoopOverTensorAtDim(OpBuilder &builder, Location loc, size_t tid,
+                                  size_t dim);
+
+  /// Emits extra locals, since the locals might not be in simplified lattices
+  /// point used to generate the loops, but are still required to generates
+  /// expressions.
+  void emitExtraLocalsForTensorsAtDenseDims(OpBuilder &builder, Location loc,
+                                            ArrayRef<size_t> tids,
+                                            ArrayRef<size_t> dims);
+
+  /// Linearizes address for dense dimension (i.e., p = (i * d0) + j).
+  Value genAddress(OpBuilder &builder, Location loc, size_t tid, size_t dim,
+                   Value iv) {
+    Value p = dim == 0 ? constantIndex(builder, loc, 0) : pidxs[tid][dim - 1];
+    Value mul = builder.create<arith::MulIOp>(loc, highs[tid][dim], p);
+    Value add = builder.create<arith::AddIOp>(loc, mul, iv);
+    return add;
+  }
+
+  bool isLastOutput; // Is the last tensor output tensor
+  bool isOutputTensor(size_t tid) {
+    return isLastOutput && tid == tensors.size() - 1;
+  }
+
+  /// Exits a for loop, returns the reduction results, e.g.,
+  /// %ret = for () {
+  ///   ...
+  ///   yield %val
+  /// }
+  /// Return %ret to user, while %val is provided by users (`reduc`)
+  SmallVector<Value, 2> exitForLoop(OpBuilder &builder, Location loc,
+                                    ArrayRef<Value> reduc);
+
+  /// Exits a while loop, returns the reduction results.
+  SmallVector<Value, 2> exitCoiterationLoop(OpBuilder &builder, Location loc,
+                                            ArrayRef<Value> reduc);
+
+  struct LoopLevelInfo {
+    LoopLevelInfo(ArrayRef<size_t> tids, ArrayRef<size_t> dims, Operation *loop,
+                  Value iv)
+        : tids(tids), dims(dims), loop(loop), iv(iv) {}
+    const llvm::SmallVector<size_t, 4> tids;
+    const llvm::SmallVector<size_t, 4> dims;
+    const Operation *loop; // the loop operation
+    const Value iv;        // the induction variable for the loop
+  };
+
+  // Loop Stack, stores the information of all the nested loops that are alive.
+  std::vector<LoopLevelInfo> loopStack;
+
+  // Loop Sequence Stack, stores the unversial index for the current loop
+  // sequence.
+  std::vector<Value> loopSeqStack;
+
+  // TODO: not yet used, it should track the current level for each tensor
+  // to help eliminate `dim` paramters from above APIs.
+  // std::vector<size_t> curLv;
+};
+
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_

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

static Value genIndexLoad(OpBuilder &builder, Location loc, Value ptr,
   return load;
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse tensor loop emitter class implementations
 //===----------------------------------------------------------------------===//
 
 SparseTensorLoopEmitter::SparseTensorLoopEmitter(ValueRange tensors,
-                                                 bool isLastOutput)
-    : tensors(tensors.begin(), tensors.end()), dims(tensors.size()),
+                                                 bool isLastOutput,
+                                                 bool isSparseOut)
+    : tensors(tensors.begin(), tensors.end()), dimTypes(tensors.size()),
       pidxs(tensors.size()), coord(tensors.size()), highs(tensors.size()),
-      sizes(tensors.size()), ptrBuffer(tensors.size()),
-      idxBuffer(tensors.size()), valBuffer(tensors.size()),
-      isLastOutput(isLastOutput), loopStack(), curLv(tensors.size(), 0) {
-  for (size_t i = 0, e = tensors.size(); i < e; i++) {
-    auto t = tensors[i];
-    auto rtp = t.getType().dyn_cast<RankedTensorType>();
-    if (!rtp) // a scalar (0-dimension tensors)
+      ptrBuffer(tensors.size()), idxBuffer(tensors.size()),
+      valBuffer(tensors.size()), isLastOutput(isLastOutput), loopStack() {
+  for (size_t tid = 0, e = tensors.size(); tid < e; tid++) {
+    auto t = tensors[tid];
+    // a scalar or 0-dimension tensors
+    if (isZeroRankedTensorOrScalar(t.getType()))
       continue;
-
+    auto rtp = t.getType().cast<RankedTensorType>();
     auto rank = static_cast<size_t>(rtp.getRank());
     auto enc = getSparseTensorEncoding(rtp);
-    if (enc)
+    // We always treat sparse output tensor as dense so that we always iterate
+    // it based on dim size.
+    if (enc && !(isOutputTensor(tid) && isSparseOut))
       for (auto dimTp : enc.getDimLevelType())
-        dims[i].push_back(dimTp);
+        dimTypes[tid].push_back(dimTp);
     else
-      dims[i].assign(rank, SparseTensorEncodingAttr::DimLevelType::Dense);
+      dimTypes[tid].assign(rank, SparseTensorEncodingAttr::DimLevelType::Dense);
 
     // Initialize using empty value.
-    pidxs[i].assign(rank, Value());
-    coord[i].assign(rank, Value());
-    highs[i].assign(rank, Value());
-    sizes[i].assign(rank, Value());
-    ptrBuffer[i].assign(rank, Value());
-    idxBuffer[i].assign(rank, Value());
+    pidxs[tid].assign(rank, Value());
+    coord[tid].assign(rank, Value());
+    highs[tid].assign(rank, Value());
+    ptrBuffer[tid].assign(rank, Value());
+    idxBuffer[tid].assign(rank, Value());
   }
 }
 
-void SparseTensorLoopEmitter::initializeLoopEmit(OpBuilder &builder,
-                                                 Location loc) {
+void SparseTensorLoopEmitter::initializeLoopEmit(
+    OpBuilder &builder, Location loc,
+    SparseTensorLoopEmitter::OutputUpdater updater) {
   // For every tensor, find lower and upper bound on dimensions, set the
   // same bounds on loop indices, and obtain dense or sparse buffer(s).
-  // TODO: Provides ability to generate loop on output buffer (with undef
-  // dim level in Merger in GenericOp Sparsification).
   for (size_t t = 0, e = tensors.size(); t < e; t++) {
     auto tensor = tensors[t];
-    auto rtp = tensor.getType().cast<RankedTensorType>();
+    auto rtp = tensor.getType().dyn_cast<RankedTensorType>();
+    if (!rtp) // Scalar
+      continue;
     auto rank = rtp.getRank();
     auto shape = rtp.getShape();
     auto enc = getSparseTensorEncoding(rtp);
     auto dynShape = {ShapedType::kDynamicSize};
     // Scan all dimensions of current tensor.
     for (int64_t d = 0; d < rank; d++) {
       // This should be called only once at beginning.
-      assert(!ptrBuffer[t][d] && !idxBuffer[t][d] && !sizes[t][d] &&
-             !highs[t][d]);
+      assert(!ptrBuffer[t][d] && !idxBuffer[t][d] && !highs[t][d]);
       // Handle sparse storage schemes.
-      if (isCompressedDim(dims[t][d])) {
+      if (isCompressedDim(dimTypes[t][d])) {
         auto ptrTp =
             MemRefType::get(dynShape, getPointerOverheadType(builder, enc));
         auto indTp =
             MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
         auto dim = builder.getIndexAttr(d);
         // Generate sparse primitives to obtains pointer and indices.
         ptrBuffer[t][d] = builder.create<ToPointersOp>(loc, ptrTp, tensor, dim);
         idxBuffer[t][d] = builder.create<ToIndicesOp>(loc, indTp, tensor, dim);
-      } else if (isSingletonDim(dims[t][d])) {
+      } else if (isSingletonDim(dimTypes[t][d])) {
         // Singleton dimension, fetch indices.
         auto indTp =
             MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
         auto dim = builder.getIndexAttr(d);
         idxBuffer[t][d] = builder.create<ToIndicesOp>(loc, indTp, tensor, dim);
       } else {
         // Dense dimension, nothing to fetch.
-        assert(isDenseDim(dims[t][d]));
+        assert(isDenseDim(dimTypes[t][d]));
       }
 
       // Find upper bound in current dimension.
       unsigned p = toOrigDim(enc, d);
       Value up = mlir::linalg::createOrFoldDimOp(builder, loc, tensor, p);
-      sizes[t][d] = highs[t][d] = up;
+      highs[t][d] = up;
     }
+
     // Perform the required bufferization. Dense inputs materialize
-    // from the input tensors. Dense outputs need special handling.
-    // Sparse inputs use sparse primitives to obtain the values.
+    // from the input tensors. Sparse inputs use sparse primitives to obtain the
+    // values.
+    // Delegates extra output initialization to clients.
+    bool isOutput = isOutputTensor(t);
     Type elementType = rtp.getElementType();
-
     if (!enc) {
       // Non-annotated dense tensors.
       auto denseTp = MemRefType::get(shape, elementType);
-      if (isLastOutput && t == tensors.size() - 1)
-        llvm_unreachable("TODO: not yet handled");
-      else
-        valBuffer[t] =
-            builder.create<bufferization::ToMemrefOp>(loc, denseTp, tensor);
+      Value denseVal =
+      // Dense outputs need special handling.
+      if (isOutput && updater)
+        denseVal = updater(builder, loc, denseVal, tensor);
+
+      valBuffer[t] = denseVal;
     } else {
       // Annotated sparse tensors.
+      // We also need the value buffer for annotated all dense `sparse` tensor.
       auto dynShape = {ShapedType::kDynamicSize};
       auto sparseTp = MemRefType::get(dynShape, elementType);
       valBuffer[t] = builder.create<ToValuesOp>(loc, sparseTp, tensor);
     }
-    // Prepare to enter the first dim for all (input) tensors
-    prepareLoopOverTensorAtDim(builder, loc, t, 0);
+    // NOTE: we can also prepares for 0 dim here in advance, this will hosit
+    // some loop preparation from tensor iteration, but will also (undesirably)
+    // hosit the code ouside if conditions.
   }
 }
 
+void SparseTensorLoopEmitter::enterNewLoopSeq(OpBuilder &builder, Location loc,
+                                              ArrayRef<size_t> tids,
+                                              ArrayRef<size_t> dims) {
+  // Universal Index start from 0
+  assert(loopSeqStack.size() == loopStack.size());
+  // Universal index starts from 0
+  loopSeqStack.emplace_back(constantIndex(builder, loc, 0));
+  // Prepares for all the tensors used in the current loop sequence.
+  for (auto [tid, dim] : llvm::zip(tids, dims))
+    prepareLoopOverTensorAtDim(builder, loc, tid, dim);
+}
+
 Operation *SparseTensorLoopEmitter::enterLoopOverTensorAtDim(
     OpBuilder &builder, Location loc, size_t tid, size_t dim,
-    ArrayRef<Value> reduc) {
-  assert(dims[tid].size() > dim);
+    SmallVectorImpl<Value> &reduc, bool isParallel, ArrayRef<size_t> extraTids,
+    ArrayRef<size_t> extraDims) {
+  assert(dimTypes[tid].size() > dim);
   // We can not re-enter the same level.
   assert(!coord[tid][dim]);
+
   Value step = constantIndex(builder, loc, 1);
-  auto dimType = dims[tid][dim];
-  bool isSparse = isCompressedDim(dimType) || isSingletonDim(dimType);
+  auto dimType = dimTypes[tid][dim];
+  // Note that we generate loop based on dim bound for output tensor even on
+  // sparse dimension.
+  bool isSparseInput = isCompressedDim(dimType) || isSingletonDim(dimType);
+
   assert(isDenseDim(dimType) || isCompressedDim(dimType) ||
          isSingletonDim(dimType));
 
-  Value lo = isSparse ? pidxs[tid][dim] : constantIndex(builder, loc, 0);
+  Value lo = isSparseInput ? pidxs[tid][dim]      // current offset
+                           : loopSeqStack.back(); // univeral tid
   Value hi = highs[tid][dim];
 
-  // TODO: support reduction.
-  if (!reduc.empty())
-    llvm_unreachable("TODO: not implemented yet");
-
   scf::ForOp forOp = builder.create<scf::ForOp>(loc, lo, hi, step, reduc);
   builder.setInsertionPointToStart(forOp.getBody());
   Value iv = forOp.getInductionVar();
-  Operation *loop = forOp;
-
   assert(iv);
-  if (isSparse) {
+  if (isSparseInput) {
     pidxs[tid][dim] = iv;
     // Generating a load on the indices array yields the coordinate.
     Value ptr = idxBuffer[tid][dim];
-    // TODO: generates load for vector value.
     coord[tid][dim] = genIndexLoad(builder, loc, ptr, iv);
   } else {
     // Dense tensor, the coordinates is the inducation variable.
     coord[tid][dim] = iv;
     // generate pidx for dense dim (pidx = i * sz + j)
-    // TODO: handle vector loop.
-    Value p = dim == 0 ? constantIndex(builder, loc, 0) : pidxs[tid][dim - 1];
-    Value mul = builder.create<arith::MulIOp>(loc, sizes[tid][dim], p);
-    Value add = builder.create<arith::AddIOp>(loc, mul, iv);
-    pidxs[tid][dim] = add;
+    auto enc = getSparseTensorEncoding(tensors[tid].getType());
+    if (enc)
+      pidxs[tid][dim] = genAddress(builder, loc, tid, dim, iv);
   }
 
-  // Prepares for next dim if this is not currently the innermost dimension.
-  if (dim != dims[tid].size() - 1)
-    prepareLoopOverTensorAtDim(builder, loc, tid, dim + 1);
+  // NOTE: we can also prepares for next dim here in advance
+  // Push the loop into stack
+  loopStack.emplace_back(ArrayRef<size_t>(tid), ArrayRef<size_t>(dim), forOp,
+                         coord[tid][dim]);
+  // Emit extra locals.
+  emitExtraLocalsForTensorsAtDenseDims(builder, loc, extraTids, extraDims);
+
+  // In-place update on the reduction variable vector.
+  assert(forOp.getNumRegionIterArgs() == reduc.size());
+  for (int i = 0, e = reduc.size(); i < e; i++)
+    reduc[i] = forOp.getRegionIterArg(i);
+  return forOp;
+}
+
+Operation *SparseTensorLoopEmitter::enterCoiterationOverTensorsAtDims(
+    OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
+    ArrayRef<size_t> dims, SmallVectorImpl<Value> &reduc, bool needsUniv,
+    ArrayRef<size_t> extraTids, ArrayRef<size_t> extraDims) {
+  assert(tids.size() == dims.size());
+  SmallVector<Type, 4> types;
+  SmallVector<Value, 4> operands;
+  // Construct the while-loop with a parameter for each index.
+  Type indexType = builder.getIndexType();
+  for (auto [tid, dim] : llvm::zip(tids, dims)) {
+    if (isCompressedDim(dimTypes[tid][dim]) ||
+        isSingletonDim(dimTypes[tid][dim])) {
+      assert(pidxs[tid][dim]);
+      types.push_back(indexType);
+      operands.push_back(pidxs[tid][dim]);
+    }
+  }
+  // The position where user-supplied reduction variable starts.
+  for (Value rec : reduc) {
+    types.push_back(rec.getType());
+    operands.push_back(rec);
+  }
+  if (needsUniv) {
+    types.push_back(indexType);
+    // Update universal index.
+    operands.push_back(loopSeqStack.back());
+  }
+  assert(types.size() == operands.size());
+  scf::WhileOp whileOp = builder.create<scf::WhileOp>(loc, types, operands);
+
+  SmallVector<Location> locs(types.size(), loc);
+  Block *before = builder.createBlock(&whileOp.getBefore(), {}, types, locs);
+  Block *after = builder.createBlock(&whileOp.getAfter(), {}, types, locs);
+
+  // Build the "before" region, which effectively consists
+  // of a conjunction of "i < upper" tests on all induction.
+  builder.setInsertionPointToStart(&whileOp.getBefore().front());
+  Value cond;
+  unsigned o = 0;
+  for (auto [tid, dim] : llvm::zip(tids, dims)) {
+    if (isCompressedDim(dimTypes[tid][dim]) ||
+        isSingletonDim(dimTypes[tid][dim])) {
+      Value op1 = before->getArgument(o);
+      Value op2 = highs[tid][dim];
+      Value opc = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
+                                                op1, op2);
+      cond = cond ? builder.create<arith::AndIOp>(loc, cond, opc) : opc;
+      // Update
+      pidxs[tid][dim] = after->getArgument(o++);
+    }
+  }
+  builder.create<scf::ConditionOp>(loc, cond, before->getArguments());
+
+  // Generates while body.
+  builder.setInsertionPointToStart(&whileOp.getAfter().front());
+  Value min;
+  for (auto [tid, dim] : llvm::zip(tids, dims)) {
+    // Prepares for next level.
+    if (isCompressedDim(dimTypes[tid][dim]) ||
+        isSingletonDim(dimTypes[tid][dim])) {
+      Value ptr = idxBuffer[tid][dim];
+      Value s = pidxs[tid][dim];
+      Value load = genIndexLoad(builder, loc, ptr, s);
+      coord[tid][dim] = load;
+      if (!needsUniv) {
+        if (min) {
+          Value cmp = builder.create<arith::CmpIOp>(
+              loc, arith::CmpIPredicate::ult, load, min);
+          min = builder.create<arith::SelectOp>(loc, cmp, load, min);
+        } else {
+          min = load;
+        }
+      }
+    }
+  }
+
+  if (needsUniv) {
+    assert(!min);
+    // Otherwise, universal index is the minimal pidx.
+    min = after->getArguments().back();
+  }
 
-  loopStack.push_back(LoopLevelInfo({tid}, {dim}, coord[tid][dim]));
-  return loop;
+  for (auto [tid, dim] : llvm::zip(tids, dims)) {
+    // All dense dim (as well as sparse output tensor) shared the same pidx in
+    // the while loop.
+    if (isDenseDim(dimTypes[tid][dim])) {
+      pidxs[tid][dim] = min;
+      // generate pidx for dense dim (pidx = i * sz + j)
+      auto enc = getSparseTensorEncoding(tensors[tid].getType());
+      if (enc)
+        pidxs[tid][dim] = genAddress(builder, loc, tid, dim, min);
     }
+    // NOTE: we can also prepares for next dim here in advance
+  }
+  // Sets up the loop stack.
+  loopStack.emplace_back(tids, dims, whileOp, min);
+  assert(loopStack.size() == loopSeqStack.size());
+
+  // Emits extra locals
+  emitExtraLocalsForTensorsAtDenseDims(builder, loc, extraTids, extraDims);
 
-void SparseTensorLoopEmitter::enterCoiterationOverTensorsAtDims(
-    OpBuilder &builder, Location loc, ArrayRef<size_t> ts,
-    ArrayRef<size_t> ds) {
-  llvm_unreachable("TODO: unimplemented");
+  // Updates reduction variables
+  assert(after->getNumArguments() == o + reduc.size() + (needsUniv ? 1 : 0));
+  // In-place update on reduction variable.
+  for (unsigned i = 0, e = reduc.size(); i < e; i++)
+    reduc[i] = after->getArgument(o + i);
+
+  return whileOp;
 }
 
-bool SparseTensorLoopEmitter::prepareLoopOverTensorAtDim(OpBuilder &builder,
+void SparseTensorLoopEmitter::prepareLoopOverTensorAtDim(OpBuilder &builder,
                                                          Location loc,
                                                          size_t tid,
                                                          size_t dim) {
-  // TODO: generate loop iteration on output tensor based on the shape
-  // instead of pointer/indices arrays.
-  assert(dims[tid].size() > dim);
-  auto dimType = dims[tid][dim];
+  assert(dimTypes[tid].size() > dim);
+  auto dimType = dimTypes[tid][dim];
 
   if (isDenseDim(dimType))
-    return false;
-
+    // Note that output sparse tid adimalso iterate by loop bound.
+    return;
+  // assert(pidxs[tid][dim] == nullptr);
   // Either the first dimension, or the previous dimension has been set.
   assert(dim == 0 || pidxs[tid][dim - 1]);
   Value c0 = constantIndex(builder, loc, 0);
   Value c1 = constantIndex(builder, loc, 1);
   if (isCompressedDim(dimType)) {
     Value ptr = ptrBuffer[tid][dim];
 
     Value pLo = dim == 0 ? c0 : pidxs[tid][dim - 1];
-    Value pHi = builder.create<arith::AddIOp>(loc, pLo, c1);
-
     pidxs[tid][dim] = genIndexLoad(builder, loc, ptr, pLo);
+
+    Value pHi = builder.create<arith::AddIOp>(loc, pLo, c1);
     highs[tid][dim] = genIndexLoad(builder, loc, ptr, pHi);
-    return true;
+    return;
   }
   if (isSingletonDim(dimType)) {
     Value pLo = dim == 0 ? c0 : pidxs[tid][dim - 1];
     Value pHi = builder.create<arith::AddIOp>(loc, pLo, c1);
 
     pidxs[tid][dim] = pLo;
     highs[tid][dim] = pHi;
-    return true;
+    return;
   }
 
   llvm_unreachable("Unrecognizable dimesion type!");
 }
 
-Value SparseTensorLoopEmitter::emitExtraLocalsForTensorsAtDims(
-    OpBuilder &builder, Location loc, size_t tid, size_t dim) {
-  llvm_unreachable("TODO: not implemented yet");
+// FIXME: Make this call private
+void SparseTensorLoopEmitter::emitExtraLocalsForTensorsAtDenseDims(
+    OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
+    ArrayRef<size_t> dims) {
+  // Initialize dense positions. Note that we generate dense indices of the
+  // output tensor unconditionally, since they may not appear in the lattice,
+  // but may be needed for linearized codegen.
+  for (auto [tid, dim] : llvm::zip(tids, dims)) {
+    assert(isDenseDim(dimTypes[tid][dim]));
+    auto enc = getSparseTensorEncoding(tensors[tid].getType());
+    if (enc) {
+      bool validPidx = dim == 0 || pidxs[tid][dim - 1];
+      if (!validPidx) {
+        // We might not find the pidx for the sparse output tensor as it is
+        // unconditionally required by the sparsification.
+        assert(isOutputTensor(tid));
+        continue;
+      }
+      pidxs[tid][dim] = genAddress(builder, loc, tid, dim, loopStack.back().iv);
+      // NOTE: we can also prepares for next dim here in advance
+    }
+  }
 }
 
-void SparseTensorLoopEmitter::exitCurrentLoop() {
-  // Clean up the values, it would help use to discover potential bug at a
-  // earlier stage (instead of silently using a wrong value).
+SmallVector<Value, 2>
+SparseTensorLoopEmitter::exitForLoop(OpBuilder &builder, Location loc,
+                                     ArrayRef<Value> reduc) {
   LoopLevelInfo &loopInfo = loopStack.back();
-  assert(loopInfo.tensors.size() == loopInfo.dims.size());
-  for (auto info : llvm::zip(loopInfo.tensors, loopInfo.dims)) {
-    auto tid = std::get<0>(info);
-    auto dim = std::get<1>(info);
-    assert(pidxs[tid][dim] && coord[tid][dim] && highs[tid][dim]);
+  auto &dims = loopStack.back().dims;
+  auto &tids = loopStack.back().tids;
+  auto forOp = llvm::cast<scf::ForOp>(loopInfo.loop);
+  if (!reduc.empty()) {
+    assert(reduc.size() == forOp.getNumResults());
+    builder.setInsertionPointToEnd(forOp.getBody());
+    builder.create<scf::YieldOp>(loc, reduc);
+  }
+
+  // Finished iterating a tensor, clean up
+  // We only do the clean up on for loop as while loops do not necessarily
+  // finish the iteration on a sparse tensor
+  for (auto [tid, dim] : llvm::zip(tids, dims)) {
     // Reset to null.
-    pidxs[tid][dim] = Value();
     coord[tid][dim] = Value();
-    if (!isDenseDim(dims[tid][dim]))
-      // Dense dimension, high is fixed.
+    pidxs[tid][dim] = Value();
+    if (!isDenseDim(dimTypes[tid][dim]))
       highs[tid][dim] = Value();
   }
+  // exit the loop
+  builder.setInsertionPointAfter(forOp);
+  return forOp.getResults();
+}
+
+SmallVector<Value, 2>
+SparseTensorLoopEmitter::exitCoiterationLoop(OpBuilder &builder, Location loc,
+                                             ArrayRef<Value> reduc) {
+  auto whileOp = llvm::cast<scf::WhileOp>(loopStack.back().loop);
+  auto &dims = loopStack.back().dims;
+  auto &tids = loopStack.back().tids;
+  Value iv = loopStack.back().iv;
+  // Generation while loop induction at the end.
+  builder.setInsertionPointToEnd(&whileOp.getAfter().front());
+  // Finalize the induction. Note that the induction could be performed
+  // in the individual if-branches to avoid re-evaluating the conditions.
+  // However, that would result in a rather elaborate forest of yield
+  // instructions during code generation. Moreover, performing the induction
+  // after the if-statements more closely resembles code generated by TACO.
+  unsigned o = 0;
+  SmallVector<Value, 4> operands;
+  Value one = constantIndex(builder, loc, 1);
+  for (auto [tid, dim] : llvm::zip(tids, dims)) {
+    if (isCompressedDim(dimTypes[tid][dim]) ||
+        isSingletonDim(dimTypes[tid][dim])) {
+      Value op1 = coord[tid][dim];
+      Value op3 = pidxs[tid][dim];
+      Value cmp =
+          builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, op1, iv);
+      Value add = builder.create<arith::AddIOp>(loc, op3, one);
+      operands.push_back(builder.create<arith::SelectOp>(loc, cmp, add, op3));
+      // Following loops continue iteration from the break point of the
+      // current while loop.
+      pidxs[tid][dim] = whileOp->getResult(o++);
+      // The coordinates are invalid now.
+      coord[tid][dim] = nullptr;
+      // highs remains unchanged.
+    }
+  }
+
+  // Reduction value from users.
+  SmallVector<Value, 2> ret;
+  for (auto red : reduc) {
+    operands.push_back(red);
+    ret.push_back(whileOp->getResult(o++));
+  }
+
+  // An (optional) universal index.
+  if (operands.size() < whileOp.getNumResults()) {
+    assert(operands.size() + 1 == whileOp.getNumResults());
+    // The last one is the universial index.
+    operands.push_back(builder.create<arith::AddIOp>(loc, iv, one));
+    // update the loop starting point of current loop sequence
+    loopSeqStack.back() = whileOp->getResult(o++);
+  }
+
+  assert(o == operands.size());
+  builder.create<scf::YieldOp>(loc, operands);
+  builder.setInsertionPointAfter(whileOp);
+  return ret;
+}
+
+SmallVector<Value, 2>
+SparseTensorLoopEmitter::exitCurrentLoop(OpBuilder &builder, Location loc,
+                                         ArrayRef<Value> reduc) {
+  // Clean up the values, it would help use to discover potential bug at a
+  // earlier stage (instead of silently using a wrong value).
+  LoopLevelInfo &loopInfo = loopStack.back();
+  assert(loopInfo.tids.size() == loopInfo.dims.size());
+  SmallVector<Value, 2> red;
+  if (llvm::isa<scf::WhileOp>(loopInfo.loop)) {
+    red = exitCoiterationLoop(builder, loc, reduc);
+  } else {
+    red = exitForLoop(builder, loc, reduc);
+  }
+
+  assert(loopStack.size() == loopSeqStack.size());
   loopStack.pop_back();
+  return red;
 }
 
 //===----------------------------------------------------------------------===//
 // ExecutionEngine/SparseTensorUtils helper functions.
 //===----------------------------------------------------------------------===//
 
 OverheadType mlir::sparse_tensor::overheadTypeEncoding(unsigned width) {
   switch (width) {
void mlir::sparse_tensor::genReshapeDstShape(
   assert(reassociation.size() == srcShape.size());
   unsigned start = 0;
   // Expand the i-th dimension in srcShape.
   for (unsigned i = 0, size = srcShape.size(); i < size; i++) {
     const auto &map = reassociation[i];
     auto srcDim = srcShape[i];
     // Iterate through dimensions expanded from the i-th dimension.
     for (unsigned j = start; j < start + map.size(); j++) {
-      // There can be only one dynamic sized dimension among dimensions expanded
-      // from the i-th dimension in srcShape. For example, if srcDim = 8, then
-      // the expanded shape could be <2x?x2>, but not <2x?x?>.
+      // There can be only one dynamic sized dimension among dimensions
+      // expanded from the i-th dimension in srcShape. For example, if srcDim
+      // = 8, then the expanded shape could be <2x?x2>, but not <2x?x?>.
       if (staticDstShape[j] == ShapedType::kDynamicSize) {
         // The expanded dimension has dynamic size. We compute the dimension
         // by dividing srcDim by the product of the static dimensions.
         int64_t product = 1;
         for (unsigned k = start; k < start + map.size(); k++) {
           if (staticDstShape[k] != ShapedType::kDynamicSize) {
             product *= staticDstShape[k];
           }

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

LogicalResult matchAndRewrite(ForeachOp op,
     Value input = op.getTensor();
     auto rtp = input.getType().cast<RankedTensorType>();
     int64_t rank = rtp.getRank();
     auto enc = getSparseTensorEncoding(rtp);
 
     // 1. Generates loop for the sparse input.
     SparseTensorLoopEmitter loopEmitter(ValueRange{input});
     loopEmitter.initializeLoopEmit(rewriter, loc);
-    for (int64_t i = 0; i < rank; i++)
-      loopEmitter.enterLoopOverTensorAtDim(rewriter, loc, 0, i);
+    SmallVector<Value, 0> reduc{}; // no reduction variable
+    for (int64_t i = 0; i < rank; i++) {
+      // TODO: provide utility function for loop sequences that only contains
+      // one for loop?
+      loopEmitter.enterNewLoopSeq(rewriter, loc, 0, static_cast<size_t>(i));
+      loopEmitter.enterLoopOverTensorAtDim(rewriter, loc, 0, i, reduc);
+    }
 
-    Value vals = loopEmitter.getTensorValueBuffer(0);
-    Value idx = loopEmitter.getLastLevelTensorPointerIndex(0);
+    Value vals = loopEmitter.getValBuffer()[0];
+    Value idx = loopEmitter.getPidxs()[0].back();
     Value val = rewriter.create<memref::LoadOp>(op.getLoc(), vals, idx);
 
     SmallVector<Value, 4> coords;
     coords.reserve(rank);
     loopEmitter.getCoordinateArray(coords);
 
-    for (int64_t i = 0; i < rank; i++)
-      loopEmitter.exitCurrentLoop();
-
     // 2. Inline the block in the foreach operator.
     Block::iterator inlinePos = rewriter.getInsertionPoint();
     Block *srcBlock = op.getBody();
     // Remove sparse_tensor.yield.
     rewriter.eraseOp(srcBlock->getTerminator());
 
+    for (int64_t i = 0; i < rank; i++) {
+      loopEmitter.exitCurrentLoop(rewriter, loc);
+      loopEmitter.exitCurrentLoopSeq();
+    }
+
     SmallVector<Value, 4> args;
     // Remap coordinates.
     for (int64_t i = 0; i < rank; i++) {
       Value actual = coords[toOrigDim(enc, i)];
       args.push_back(actual);
     }
     // Remap value.
     args.push_back(val);

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

 using namespace mlir::sparse_tensor;
 
 //===----------------------------------------------------------------------===//
 // Declarations of data structures.
 //===----------------------------------------------------------------------===//
 
 namespace {
 
+constexpr unsigned INVALID_ID = std::numeric_limits<unsigned>::max();
+
 // Iteration graph sorting.
 enum SortMask {
   kSparseOnly = 0x0,
   kIncludeDense = 0x1,
   kIncludeUndef = 0x2,
   kIncludeAll = 0x3
 };
 
 // Reduction kinds.
 enum Reduction { kNoReduc, kSum, kProduct, kAnd, kOr, kXor, kCustom };
 
 // Code generation.
 struct CodeGen {
-  CodeGen(SparsificationOptions o, unsigned numTensors, unsigned numLoops,
-          OpOperand *op, unsigned nest, std::vector<unsigned> &ts)
-      : options(o), loops(numLoops), sizes(numLoops), buffers(numTensors),
-        pointers(numTensors, std::vector<Value>(numLoops)),
-        indices(numTensors, std::vector<Value>(numLoops)),
-        highs(numTensors, std::vector<Value>(numLoops)),
-        pidxs(numTensors, std::vector<Value>(numLoops)),
-        idxs(numTensors, std::vector<Value>(numLoops)), sparseOut(op),
-        outerParNest(nest), topSort(ts) {}
+  CodeGen(SparsificationOptions o, ValueRange tensors, unsigned numTensors,
+          unsigned numLoops, OpOperand *op, unsigned nest,
+          std::vector<unsigned> &ts)
+      : options(o), loopEmitter(tensors, /*isLastOutput=*/true,
+                                /*isSparseOut=*/op != nullptr),
+        sparseOut(op), outerParNest(nest), topSort(ts) {}
   /// Sparsification options.
   SparsificationOptions options;
-  /// Universal dense indices and upper bounds (by index). The loops array
-  /// is updated with the value of the universal dense index in the current
-  /// loop. The sizes array is set once with the inferred dimension sizes.
-  std::vector<Value> loops;
-  std::vector<Value> sizes;
-  /// Buffers for storing dense and sparse numerical values (by tensor).
-  /// This array is set once during bufferization of all tensors.
-  std::vector<Value> buffers;
-  /// Sparse storage schemes (1-D): pointers and indices (by tensor and index).
-  /// This array is set once during bufferization of all sparse tensors.
-  std::vector<std::vector<Value>> pointers;
-  std::vector<std::vector<Value>> indices;
-  /// Sparse iteration information (by tensor and index). These arrays
-  /// are updated to remain current within the current loop.
-  std::vector<std::vector<Value>> highs;
-  std::vector<std::vector<Value>> pidxs;
-  std::vector<std::vector<Value>> idxs;
+  /// Loop emitter helper class.
+  SparseTensorLoopEmitter loopEmitter;
   /// Current reduction, updated during code generation. When indices of a
   /// reduction are exhausted, all inner loops can use a scalarized reduction.
   unsigned redExp = -1u;
   Value redVal;
   Reduction redKind = kNoReduc;
   unsigned redCustom = -1u;
   // Sparse tensor as output. Implemented either through direct injective
   // insertion in lexicographic index order or through access pattern expansion
   // in the innermost loop nest (`expValues` through `expCount`).
   OpOperand *sparseOut;
   unsigned outerParNest;
   Value expValues;
   Value expFilled;
   Value expAdded;
   Value expCount;
   // Current vector length and mask.
   unsigned curVecLength = 1;
   Value curVecMask;
   // Topsort (reference should remain in scope).
   std::vector<unsigned> &topSort;
+
+  // From tensor id + loop id => dim id.
+  // TODO: This map should probably be maintained by Merger (it can be set up
+  // together with dimLvlType Map).
+  std::vector<std::vector<unsigned>> loopIdxToDim;
+
+  // Initialize the above two mapping.
+  void buildLoopIdxToDimMap(linalg::GenericOp op);
+
+  Value getLoopIdxValue(size_t loopIdx) const {
+    for (unsigned lv = 0; lv < topSort.size(); lv++)
+      if (topSort[lv] == loopIdx)
+        return loopEmitter.getLoopIV(lv);
+
+    llvm_unreachable("invalid loop index");
+  }
 };
 
+void CodeGen::buildLoopIdxToDimMap(linalg::GenericOp op) {
+  size_t numLoops = op.getNumLoops();
+  size_t numTensors = op.getNumOperands();
+  loopIdxToDim.assign(numTensors, std::vector<unsigned>(numLoops, INVALID_ID));
+
+  for (OpOperand &t : op->getOpOperands()) {
+    auto map = op.getMatchingIndexingMap(&t);
+    auto enc = getSparseTensorEncoding(t.get().getType());
+    // Scan all dimensions of current tensor.
+    unsigned tid = t.getOperandNumber();
+    for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
+      auto a = map.getResult(toOrigDim(enc, d)).dyn_cast<AffineDimExpr>();
+      if (a) {
+        unsigned loopId = a.getPosition();
+        // Fills the mapping.
+        loopIdxToDim[tid][loopId] = d;
+      }
+      // Else a compound affine, do nothing. (at least we are good for
+      // now, as we only support compound affine expr on non-annoated dense
+      // tensors).
+    }
+  }
+}
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler analysis methods.
 //===----------------------------------------------------------------------===//
 
 /// Helper method to construct a permuted dimension ordering
 /// that adheres to the given topological sort.
static DimLevelFormat toDimLevelFormat(const SparseTensorEncodingAttr &enc,
   }
   return DimLevelFormat(DimLvlType::kDense);
 }
 
 /// Helper method to inspect affine expressions. Rejects cases where the
 /// same index is used more than once. Also rejects compound affine
 /// expressions in sparse dimensions.
 static bool findAffine(Merger &merger, unsigned tensor, AffineExpr a,
-                       DimLevelFormat dim) {
+                       DimLevelFormat dim, bool setLvlFormat = true) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
     unsigned idx = a.cast<AffineDimExpr>().getPosition();
     if (!merger.isDimLevelType(tensor, idx, DimLvlType::kUndef))
       return false; // used more than once
+    if (setLvlFormat)
       merger.setDimLevelFormat(tensor, idx, dim);
     return true;
   }
   case AffineExprKind::Add:
   case AffineExprKind::Mul: {
     if (dim.levelType != DimLvlType::kDense)
       return false; // compound only in dense dim
     auto binOp = a.cast<AffineBinaryOpExpr>();
-    return findAffine(merger, tensor, binOp.getLHS(), dim) &&
-           findAffine(merger, tensor, binOp.getRHS(), dim);
+    // We do not set dim level format for affine expresssion like d0 + d1 on
+    // both loop index at d0 and d1,
+    return findAffine(merger, tensor, binOp.getLHS(), dim, false) &&
+           findAffine(merger, tensor, binOp.getRHS(), dim, false);
   }
   case AffineExprKind::Constant:
     return dim.levelType == DimLvlType::kDense; // const only in dense dim
   default:
     return false;
   }
 }
 
static bool isAdmissibleTensorExp(Merger &merger, linalg::GenericOp op,
   }
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler synthesis methods (reductions).
 //===----------------------------------------------------------------------===//
 
-/// Maps reduction kind to vector::CombiningKind.
-static vector::CombiningKind getCombiningKind(Reduction kind) {
-  switch (kind) {
-  case kNoReduc:
-  case kCustom:
-    break;
-  case kSum:
-    return vector::CombiningKind::ADD;
-  case kProduct:
-    return vector::CombiningKind::MUL;
-  case kAnd:
-    return vector::CombiningKind::AND;
-  case kOr:
-    return vector::CombiningKind::OR;
-  case kXor:
-    return vector::CombiningKind::XOR;
-  }
-  llvm_unreachable("unknown reduction kind");
-}
-
 /// Maps operation to reduction.
 static Reduction getReduction(Kind kind) {
   switch (kind) {
   case Kind::kAddF:
   case Kind::kAddC:
   case Kind::kAddI:
   case Kind::kSubF:
   case Kind::kSubC:
case Kind::kXorI:
     return kXor;
   case Kind::kReduce:
     return kCustom;
   default:
     llvm_unreachable("unexpected reduction operator");
   }
 }
 
-/// Generates an initial value for a vector reduction, following the scheme
-/// given in Chapter 5 of "The Software Vectorization Handbook", where the
-/// initial scalar value is correctly embedded in the vector reduction value,
-/// and a straightforward horizontal reduction will complete the operation.
-static Value genVectorReducInit(CodeGen &codegen, OpBuilder &builder,
-                                Location loc, VectorType vtp) {
-  Value r = codegen.redVal;
-  switch (codegen.redKind) {
-  case kNoReduc:
-  case kCustom:
-    break;
-  case kSum:
-  case kXor:
-    // Initialize reduction vector to: | 0 | .. | 0 | r |
-    return builder.create<vector::InsertElementOp>(
-        loc, r, constantZero(builder, loc, vtp),
-        constantIndex(builder, loc, 0));
-  case kProduct:
-    // Initialize reduction vector to: | 1 | .. | 1 | r |
-    return builder.create<vector::InsertElementOp>(
-        loc, r, constantOne(builder, loc, vtp), constantIndex(builder, loc, 0));
-  case kAnd:
-  case kOr:
-    // Initialize reduction vector to: | r | .. | r | r |
-    return builder.create<vector::BroadcastOp>(loc, vtp, r);
-  }
-  llvm_unreachable("unknown reduction kind");
-}
-
-/// Generates final value for a vector reduction.
-static Value genVectorReducEnd(CodeGen &codegen, OpBuilder &builder,
-                               Location loc, VectorType vtp) {
-  vector::CombiningKind kind = getCombiningKind(codegen.redKind);
-  return builder.create<vector::ReductionOp>(loc, kind, codegen.redVal);
-}
-
 /// Updates scalarized reduction value.
 static void updateReduc(Merger &merger, CodeGen &codegen, Value reduc) {
   assert(codegen.redKind != kNoReduc);
   codegen.redVal = merger.exp(codegen.redExp).val = reduc;
 }
 
 /// Extracts identity from custom reduce.
 static Value getCustomRedId(Operation *op) {
   return dyn_cast<sparse_tensor::ReduceOp>(op).getIdentity();
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler synthesis methods (statements and expressions).
 //===----------------------------------------------------------------------===//
 
-/// Generates buffer for the output tensor. Note that all sparse kernels
-/// assume that when all elements are written to (viz. x(i) = y(i) * z(i)),
-/// the output buffer is already initialized to all zeroes and only nonzeroes
-/// values are computed and written out. For updates (viz. x(i) += y(i) * z(i)),
-/// only nonzeroes values are used for the updates and no assumption on the
-/// original contents of the output buffer is necessary.
-static Value genOutputBuffer(CodeGen &codegen, OpBuilder &builder,
-                             linalg::GenericOp op, MemRefType denseTp,
-                             ArrayRef<Value> args) {
+/// Local bufferization of all dense and sparse data structures.
+static void genBuffers(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                       linalg::GenericOp op) {
   Location loc = op.getLoc();
+  assert(op.getNumOperands() == op.getNumInputs() + 1);
+
+  codegen.loopEmitter.initializeLoopEmit(
+      builder, loc,
+      /// Generates buffer for the output tensor. Note that all sparse kernels
+      /// assume that when all elements are written to (viz. x(i) = y(i) *
+      /// z(i)), the output buffer is already initialized to all zeroes and only
+      /// nonzeroes values are computed and written out. For updates (viz. x(i)
+      /// += y(i) * z(i)), only nonzeroes values are used for the updates and no
+      /// assumption on the original contents of the output buffer is necessary.
+      [&op](OpBuilder &builder, Location loc, Value memref,
+            Value tensor) -> Value {
+        // Must not be a sparse tensor.
+        assert(!getSparseTensorEncoding(tensor.getType()));
         OpOperand *lhs = op.getOutputOperand(0);
-  Value tensor = lhs->get();
+        // Two output tensors should match.
+        assert(lhs->get() == tensor);
         bool isInit = op.isInitTensor(lhs);
-  // An output tensor can simply materialize from the buffer of the tensor that
-  // appears in the outs() clause. For updates, this has the advantage that only
-  // the nonzero value are involved in the computation, keeping the operation
-  // O(nnz). In all other cases, we are forced to zero out the buffer to enforce
-  // the assumption above, which may negatively impact running complexity
-  // (viz. O(n^2 + nnz) vs. O(nnz) for matrices).
+        // An output tensor can simply materialize from the buffer of the tensor
+        // that appears in the outs() clause. For updates, this has the
+        // advantage that only the nonzero value are involved in the
+        // computation, keeping the operation O(nnz). In all other cases, we are
+        // forced to zero out the buffer to enforce the assumption above, which
+        // may negatively impact running complexity (viz. O(n^2 + nnz) vs.
+        // O(nnz) for matrices).
         // TODO: use better analysis to avoid zeroing out the buffer?
-  Value init = builder.create<bufferization::ToMemrefOp>(loc, denseTp, tensor);
+        Value init = memref;
         if (!isInit) {
-    Value zero = constantZero(builder, loc, denseTp.getElementType());
-    builder.create<linalg::FillOp>(loc, ValueRange{zero}, ValueRange{init});
+          Value zero = constantZero(builder, loc,
+                                    getElementTypeOrSelf(tensor.getType()));
+          builder.create<linalg::FillOp>(loc, ValueRange{zero},
+                                         ValueRange{init});
         }
         return init;
-}
-
-/// Local bufferization of all dense and sparse data structures.
-static void genBuffers(Merger &merger, CodeGen &codegen, OpBuilder &builder,
-                       linalg::GenericOp op) {
-  Location loc = op.getLoc();
-  assert(op->getNumOperands() == op.getNumInputs() + 1);
-  // For every tensor, find lower and upper bound on dimensions, set the
-  // same bounds on loop indices, and obtain dense or sparse buffer(s).
-  auto dynShape = {ShapedType::kDynamicSize};
-  SmallVector<Value, 4> args;
-  for (OpOperand &t : op->getOpOperands()) {
-    unsigned tensor = t.getOperandNumber();
-    auto shape = op.getShape(&t);
-    auto map = op.getMatchingIndexingMap(&t);
-    auto enc = getSparseTensorEncoding(t.get().getType());
-    // Scan all dimensions of current tensor.
-    args.clear();
-    for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
-      AffineExpr a = map.getResult(toOrigDim(enc, d));
-      if (a.getKind() != AffineExprKind::DimId)
-        continue; // compound
-      unsigned idx = a.cast<AffineDimExpr>().getPosition();
-      // Handle the different storage schemes.
-      if (merger.isDimLevelType(tensor, idx, DimLvlType::kCompressed)) {
-        // Compressed dimension, fetch pointer and indices.
-        auto ptrTp =
-            MemRefType::get(dynShape, getPointerOverheadType(builder, enc));
-        auto indTp =
-            MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
-        auto dim = builder.getIndexAttr(d);
-        codegen.pointers[tensor][idx] =
-            builder.create<ToPointersOp>(loc, ptrTp, t.get(), dim);
-        codegen.indices[tensor][idx] =
-            builder.create<ToIndicesOp>(loc, indTp, t.get(), dim);
-      } else if (merger.isDimLevelType(tensor, idx, DimLvlType::kSingleton)) {
-        // Singleton dimension, fetch indices.
-        auto indTp =
-            MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
-        auto dim = builder.getIndexAttr(d);
-        codegen.indices[tensor][idx] =
-            builder.create<ToIndicesOp>(loc, indTp, t.get(), dim);
-      } else {
-        // Dense dimension, nothing to fetch.
-        assert(merger.isDimLevelType(tensor, idx, DimLvlType::kDense));
-      }
-      // Find upper bound in current dimension.
-      unsigned p = toOrigDim(enc, d);
-      Value up = linalg::createOrFoldDimOp(builder, loc, t.get(), p);
-      if (ShapedType::isDynamic(shape[p]))
-        args.push_back(up);
-      assert(codegen.highs[tensor][idx] == nullptr);
-      codegen.sizes[idx] = codegen.highs[tensor][idx] = up;
-    }
-    // Perform the required bufferization. Dense inputs materialize
-    // from the input tensors. Dense outputs need special handling.
-    // Sparse inputs use sparse primitives to obtain the values.
-    Type elementType = getElementTypeOrSelf(t.get().getType());
-    if (!enc) {
-      // Non-annotated dense tensors.
-      auto denseTp = MemRefType::get(shape, elementType);
-      if (tensor < op.getNumInputs())
-        codegen.buffers[tensor] =
-            builder.create<bufferization::ToMemrefOp>(loc, denseTp, t.get());
-      else
-        codegen.buffers[tensor] =
-            genOutputBuffer(codegen, builder, op, denseTp, args);
-    } else if (&t != codegen.sparseOut) {
-      // Annotated sparse tensors (not involved in output).
-      auto sparseTp = MemRefType::get(dynShape, elementType);
-      codegen.buffers[tensor] =
-          builder.create<ToValuesOp>(loc, sparseTp, t.get());
-    }
-  }
+      });
 }
 
 /// Constructs vector type.
 static VectorType vectorType(CodeGen &codegen, Type etp) {
   unsigned numScalableDims = codegen.options.enableVLAVectorization;
   return VectorType::get(codegen.curVecLength, etp, numScalableDims);
 }
 
 /// Constructs vector type from pointer.
 static VectorType vectorType(CodeGen &codegen, Value ptr) {
   return vectorType(codegen, ptr.getType().cast<MemRefType>().getElementType());
 }
 
-/// Constructs vector iteration mask.
-static Value genVectorMask(CodeGen &codegen, OpBuilder &builder, Value iv,
-                           Value lo, Value hi, Value step) {
-  Location loc = iv.getLoc();
-  VectorType mtp = vectorType(codegen, builder.getI1Type());
-  // Special case if the vector length evenly divides the trip count (for
-  // example, "for i = 0, 128, 16"). A constant all-true mask is generated
-  // so that all subsequent masked memory operations are immediately folded
-  // into unconditional memory operations.
-  IntegerAttr loInt, hiInt, stepInt;
-  if (matchPattern(lo, m_Constant(&loInt)) &&
-      matchPattern(hi, m_Constant(&hiInt)) &&
-      matchPattern(step, m_Constant(&stepInt))) {
-    if (((hiInt.getInt() - loInt.getInt()) % stepInt.getInt()) == 0)
-      return builder.create<vector::BroadcastOp>(
-          loc, mtp, constantI1(builder, loc, true));
-  }
-  // Otherwise, generate a vector mask that avoids overrunning the upperbound
-  // during vector execution. Here we rely on subsequent loop optimizations to
-  // avoid executing the mask in all iterations, for example, by splitting the
-  // loop into an unconditional vector loop and a scalar cleanup loop.
-  auto minMap = AffineMap::get(
-      /*dimCount=*/2, /*symbolCount=*/1,
-      {builder.getAffineSymbolExpr(0),
-       builder.getAffineDimExpr(0) - builder.getAffineDimExpr(1)},
-      builder.getContext());
-  Value end =
-      builder.createOrFold<AffineMinOp>(loc, minMap, ValueRange{hi, iv, step});
-  return builder.create<vector::CreateMaskOp>(loc, mtp, end);
-}
-
 /// Generates a vectorized load lhs = a[ind[lo:hi]] or lhs = a[lo:hi].
 static Value genVectorLoad(CodeGen &codegen, OpBuilder &builder, Value ptr,
                            ArrayRef<Value> args) {
   Location loc = ptr.getLoc();
   VectorType vtp = vectorType(codegen, ptr);
   Value pass = constantZero(builder, loc, vtp);
   if (args.back().getType().isa<VectorType>()) {
     SmallVector<Value, 4> scalarArgs(args.begin(), args.end());

 //
 // TODO: generalize for sparse tensor subscripts
 //
 static Value genAffine(CodeGen &codegen, OpBuilder &builder, AffineExpr a,
                        Location loc) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
     unsigned idx = a.cast<AffineDimExpr>().getPosition();
-    return codegen.loops[idx]; // universal dense index
+    return codegen.getLoopIdxValue(idx); // universal dense index
   }
   case AffineExprKind::Add: {
     auto binOp = a.cast<AffineBinaryOpExpr>();
     return builder.create<arith::AddIOp>(
         loc, genAffine(codegen, builder, binOp.getLHS(), loc),
         genAffine(codegen, builder, binOp.getRHS(), loc));
   }
   case AffineExprKind::Mul: {

 
 /// Generates index for load/store on sparse tensor.
 static Value genIndex(CodeGen &codegen, linalg::GenericOp op, OpOperand *t) {
   auto map = op.getMatchingIndexingMap(t);
   auto enc = getSparseTensorEncoding(t->get().getType());
   AffineExpr a = map.getResult(toOrigDim(enc, map.getNumResults() - 1));
   assert(a.getKind() == AffineExprKind::DimId);
   unsigned idx = a.cast<AffineDimExpr>().getPosition();
-  return codegen.loops[idx];
+  return codegen.getLoopIdxValue(idx);
 }
 
 /// Generates subscript for load/store on a dense or sparse tensor.
 static Value genSubscript(CodeGen &codegen, OpBuilder &builder,
                           linalg::GenericOp op, OpOperand *t,
                           SmallVector<Value, 4> &args) {
   unsigned tensor = t->getOperandNumber();
   auto map = op.getMatchingIndexingMap(t);
   auto enc = getSparseTensorEncoding(t->get().getType());
   unsigned rank = map.getNumResults();
   if (enc) {
     // Note that currently, all sparse subscripts are simple.
     // TODO: accept affine too?
-    AffineExpr a = map.getResult(toOrigDim(enc, rank - 1));
-    assert(a.getKind() == AffineExprKind::DimId);
-    unsigned idx = a.cast<AffineDimExpr>().getPosition();
-    assert(codegen.pidxs[tensor][idx] != nullptr);
-    args.push_back(codegen.pidxs[tensor][idx]); // position index
+    assert(map.getResult(toOrigDim(enc, rank - 1)).getKind() ==
+           AffineExprKind::DimId);
+    Value pidx = codegen.loopEmitter.getPidxs()[tensor].back();
+    assert(pidx);
+    args.push_back(pidx); // position index
   } else {
     for (unsigned d = 0; d < rank; d++) {
       AffineExpr a = map.getResult(d);
       args.push_back(genAffine(codegen, builder, a, op.getLoc()));
     }
   }
-  return codegen.buffers[tensor];
+  return codegen.loopEmitter.getValBuffer()[tensor];
 }
 
 /// Generates insertion code to implement dynamic tensor load.
 static Value genInsertionLoad(CodeGen &codegen, OpBuilder &builder,
                               linalg::GenericOp op, OpOperand *t) {
   Location loc = op.getLoc();
   // Direct lexicographic index order, tensor loads as zero.
   if (!codegen.expValues) {

 static void genInsertionStore(CodeGen &codegen, OpBuilder &builder,
                               linalg::GenericOp op, OpOperand *t, Value rhs) {
   Location loc = op.getLoc();
   // Direct insertion in lexicographic index order.
   if (!codegen.expValues) {
     unsigned rank = op.getRank(t);
     SmallVector<Value, 4> indices;
     for (unsigned i = 0; i < rank; i++) {
-      assert(codegen.loops[codegen.topSort[i]]);
-      indices.push_back(codegen.loops[codegen.topSort[i]]);
+      assert(codegen.loopEmitter.getLoopIV(i));
+      indices.push_back(codegen.loopEmitter.getLoopIV(i));
     }
     builder.create<InsertOp>(loc, rhs, t->get(), indices);
     return;
   }
   // Generates insertion code along expanded access pattern.
   //   if (!expFilled[i]) then
   //     expFilled[i] = true
   //     expAdded[inserts++] = i
static void genTensorStore(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   SmallVector<Value, 4> args;
   Value ptr = genSubscript(codegen, builder, op, t, args);
   if (codegen.curVecLength > 1)
     genVectorStore(codegen, builder, rhs, ptr, args);
   else
     builder.create<memref::StoreOp>(loc, rhs, ptr, args);
 }
 
-/// Generates a pointer/index load from the sparse storage scheme. Narrower
-/// data types need to be zero extended before casting the value into the
-/// index type used for looping and indexing.
-static Value genLoad(CodeGen &codegen, OpBuilder &builder, Location loc,
-                     Value ptr, Value s) {
-  // See https://llvm.org/docs/GetElementPtr.html for some background on
-  // the complications described below.
-  if (codegen.curVecLength > 1) {
-    // Since the index vector is used in a subsequent gather/scatter operations,
-    // which effectively defines an unsigned pointer + signed index, we must
-    // zero extend the vector to an index width. For 8-bit and 16-bit values,
-    // an 32-bit index width suffices. For 32-bit values, zero extending the
-    // elements into 64-bit loses some performance since the 32-bit indexed
-    // gather/scatter is more efficient than the 64-bit index variant (if the
-    // negative 32-bit index space is unused, the enableSIMDIndex32 flag can
-    // preserve this performance). For 64-bit values, there is no good way
-    // to state that the indices are unsigned, with creates the potential of
-    // incorrect address calculations in the unlikely case we need such
-    // extremely large offsets.
-    Type etp = ptr.getType().cast<MemRefType>().getElementType();
-    Value vload = genVectorLoad(codegen, builder, ptr, {s});
-    if (!etp.isa<IndexType>()) {
-      if (etp.getIntOrFloatBitWidth() < 32)
-        vload = builder.create<arith::ExtUIOp>(
-            loc, vectorType(codegen, builder.getI32Type()), vload);
-      else if (etp.getIntOrFloatBitWidth() < 64 &&
-               !codegen.options.enableSIMDIndex32)
-        vload = builder.create<arith::ExtUIOp>(
-            loc, vectorType(codegen, builder.getI64Type()), vload);
-    }
-    return vload;
-  }
-  // For the scalar case, we simply zero extend narrower indices into 64-bit
-  // values before casting to index without a performance penalty. Here too,
-  // however, indices that already are 64-bit, in theory, cannot express the
-  // full range as explained above.
-  Value load = builder.create<memref::LoadOp>(loc, ptr, s);
-  if (!load.getType().isa<IndexType>()) {
-    if (load.getType().getIntOrFloatBitWidth() < 64)
-      load = builder.create<arith::ExtUIOp>(loc, builder.getI64Type(), load);
-    load =
-        builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), load);
-  }
-  return load;
-}
-
 /// Generates an invariant value.
 static Value genInvariantValue(Merger &merger, CodeGen &codegen,
                                OpBuilder &builder, unsigned exp) {
   Value val = merger.exp(exp).val;
   if (codegen.curVecLength > 1)
     return genVectorInvariantValue(codegen, builder, val);
   return val;
 }
 
-/// Generates an address computation "sz * p + i".
-static Value genAddress(CodeGen &codegen, OpBuilder &builder, Location loc,
-                        Value size, Value p, Value i) {
-  Value mul = builder.create<arith::MulIOp>(loc, size, p);
-  if (auto vtp = i.getType().dyn_cast<VectorType>()) {
-    Value inv =
-        builder.create<arith::IndexCastOp>(loc, vtp.getElementType(), mul);
-    mul = genVectorInvariantValue(codegen, builder, inv);
-  }
-  return builder.create<arith::AddIOp>(loc, mul, i);
-}
-
 /// Generates an index value.
 static Value genIndexValue(CodeGen &codegen, OpBuilder &builder, unsigned idx,
                            unsigned ldx) {
-  Value ival = codegen.loops[idx];
+  Value ival = codegen.getLoopIdxValue(idx);
   Type itype = ival.getType();
   // During vectorization, we either encounter:
-  // (1) indices already in vector form, as in ... = ind[lo:hi], good to go, or
-  // (2) single index, as in ... = i, must convert to [i, i+1, ...] for inner i.
+  // (1) indices already in vector form, as in ... = ind[lo:hi], good to go,
+  // or (2) single index, as in ... = i, must convert to [i, i+1, ...] for
+  // inner i.
   unsigned vl = codegen.curVecLength;
   if (vl > 1 && !itype.isa<VectorType>()) {
     Location loc = ival.getLoc();
     VectorType vtp = vectorType(codegen, itype);
     ival = builder.create<vector::BroadcastOp>(loc, vtp, ival);
     if (idx == ldx) {
       Value incr;
       if (vtp.isScalable()) {

 /// Determines if affine expression is invariant.
 static bool isInvariantAffine(const CodeGen &codegen, AffineExpr a,
                               unsigned ldx, bool &atLevel) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
     unsigned idx = a.cast<AffineDimExpr>().getPosition();
     if (idx == ldx)
       atLevel = true;
-    return codegen.loops[idx] != nullptr; // no longer in play?
+    return codegen.getLoopIdxValue(idx) != nullptr; // no longer in play?
   }
   case AffineExprKind::Add:
   case AffineExprKind::Mul: {
     auto binOp = a.cast<AffineBinaryOpExpr>();
     return isInvariantAffine(codegen, binOp.getLHS(), ldx, atLevel) &&
            isInvariantAffine(codegen, binOp.getRHS(), ldx, atLevel);
   }
   default:

 
 /// Generates an expanded access pattern in innermost dimension.
 static void genExpansion(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                          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
+  assert(codegen.redVal == nullptr);
   // 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, builder.getI1Type());
     Type t3 = MemRefType::get(dynShape, builder.getIndexType());
     Type t4 = builder.getIndexType();
     auto res =
         builder.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);
     SmallVector<Value, 4> indices;
     for (unsigned i = 0; i < at; i++) {
-      assert(codegen.loops[codegen.topSort[i]]);
-      indices.push_back(codegen.loops[codegen.topSort[i]]);
+      assert(codegen.loopEmitter.getLoopIV(i));
+      indices.push_back(codegen.loopEmitter.getLoopIV(i));
     }
     builder.create<CompressOp>(loc, codegen.expValues, codegen.expFilled,
                                codegen.expAdded, codegen.expCount, tensor,
                                indices);
     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.
-static bool genInit(Merger &merger, CodeGen &codegen, OpBuilder &builder,
-                    linalg::GenericOp op, unsigned at, BitVector &inits) {
-  std::vector<unsigned> &topSort(codegen.topSort);
-  bool needsUniv = false;
-  Location loc = op.getLoc();
-  unsigned idx = topSort[at];
-
-  // Initialize sparse positions.
-  for (unsigned b = 0, be = inits.size(); b < be; b++) {
-    if (!inits[b])
-      continue;
-    unsigned tensor = merger.tensor(b);
-    assert(idx == merger.index(b));
-    if (merger.isDimLevelType(b, DimLvlType::kCompressed)) {
-      // Initialize sparse index that will implement the iteration:
-      //   for pidx_idx = pointers(pidx_idx-1), pointers(1+pidx_idx-1)
-      unsigned pat = at;
-      for (; pat != 0; pat--) {
-        if (codegen.pidxs[tensor][topSort[pat - 1]])
-          break;
-      }
-      Value ptr = codegen.pointers[tensor][idx];
-      Value one = constantIndex(builder, loc, 1);
-      Value p0 = (pat == 0) ? constantIndex(builder, loc, 0)
-                            : codegen.pidxs[tensor][topSort[pat - 1]];
-      codegen.pidxs[tensor][idx] = genLoad(codegen, builder, loc, ptr, p0);
-      Value p1 = builder.create<arith::AddIOp>(loc, p0, one);
-      codegen.highs[tensor][idx] = genLoad(codegen, builder, loc, ptr, p1);
-    } else if (merger.isDimLevelType(b, DimLvlType::kSingleton)) {
-      // Initialize sparse index that will implement the "iteration":
-      //   for pidx_idx = pidx_idx-1, 1+pidx_idx-1
-      // We rely on subsequent loop unrolling to get rid of the loop
-      // if it is not involved in co-iteration with anything else.
-      unsigned pat = at;
-      for (; pat != 0; pat--) {
-        if (codegen.pidxs[tensor][topSort[pat - 1]])
-          break;
-      }
-      Value one = constantIndex(builder, loc, 1);
-      Value p0 = (pat == 0) ? constantIndex(builder, loc, 0)
-                            : codegen.pidxs[tensor][topSort[pat - 1]];
-      codegen.pidxs[tensor][idx] = p0;
-      codegen.highs[tensor][idx] = builder.create<arith::AddIOp>(loc, p0, one);
-    } else {
-      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
-             merger.isDimLevelType(b, DimLvlType::kUndef));
-      // Dense index still in play.
-      needsUniv = true;
-    }
-  }
-
-  // Initialize the universal dense index.
-  codegen.loops[idx] = constantIndex(builder, loc, 0);
-  return needsUniv;
-}
-
-/// Returns vectorization strategy. Any implicit inner loop in the Linalg
-/// operation is a candidate. Whether it is actually converted to SIMD code
-/// depends on the requested strategy.
-static bool isVectorFor(CodeGen &codegen, bool isInner, bool isReduction,
-                        bool isSparse) {
-  // Reject vectorization of sparse output, unless innermost is reduction.
-  if (codegen.sparseOut && !isReduction)
-    return false;
-  // Inspect strategy.
-  switch (codegen.options.vectorizationStrategy) {
-  case SparseVectorizationStrategy::kNone:
-    return false;
-  case SparseVectorizationStrategy::kDenseInnerLoop:
-    return isInner && !isSparse;
-  case SparseVectorizationStrategy::kAnyStorageInnerLoop:
-    return isInner;
-  }
-  llvm_unreachable("unexpected vectorization strategy");
-}
-
-/// Returns parallelization strategy. Any implicit loop in the Linalg operation
-/// that is marked "parallel" is a candidate. Whether it is actually converted
-/// to a parallel operation depends on the requested strategy.
+/// Returns parallelization strategy. Any implicit loop in the Linalg
+/// operation that is marked "parallel" is a candidate. Whether it is actually
+/// converted to a parallel operation depends on the requested strategy.
 static bool isParallelFor(CodeGen &codegen, bool isOuter, bool isReduction,
-                          bool isSparse, bool isVector) {
+                          bool isSparse) {
   // Reject parallelization of sparse output.
   if (codegen.sparseOut)
     return false;
   // Inspect strategy.
   switch (codegen.options.parallelizationStrategy) {
   case SparseParallelizationStrategy::kNone:
     return false;
   case SparseParallelizationStrategy::kDenseOuterLoop:
-    return isOuter && !isSparse && !isReduction && !isVector;
+    return isOuter && !isSparse && !isReduction;
   case SparseParallelizationStrategy::kAnyStorageOuterLoop:
-    return isOuter && !isReduction && !isVector;
+    return isOuter && !isReduction;
   case SparseParallelizationStrategy::kDenseAnyLoop:
-    return !isSparse && !isReduction && !isVector;
+    return !isSparse && !isReduction;
   case SparseParallelizationStrategy::kAnyStorageAnyLoop:
-    return !isReduction && !isVector;
+    return !isReduction;
   }
   llvm_unreachable("unexpected parallelization strategy");
 }
 
-/// Checks unit stride for dense tensors. The iteration graph may have ignored
-/// dense access patterns in order to avoid cycles (sparse access patterns are
-/// always placed innermost), but that means dense access has become strided.
-/// This prevents effective vectorization.
-static bool denseUnitStrides(Merger &merger, linalg::GenericOp op,
-                             unsigned idx) {
-  for (OpOperand &t : op->getOpOperands()) {
-    if (!getSparseTensorEncoding(t.get().getType())) {
-      auto map = op.getMatchingIndexingMap(&t);
-      for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
-        AffineExpr a = map.getResult(d);
-        // Report non-unit stride if innermost index appears at an outer
-        // dimension (true non-unit stride) or if the innermost index appears
-        // in a compound subscript in the innermost dimension. Even if the
-        // latter is unit stride, it does not play well with scatter/gather.
-        // TODO: accept unit stride affine innermost like a[i,j+k+1]?
-        if (a.isFunctionOfDim(idx) &&
-            ((d != rank - 1) || (a.getKind() != AffineExprKind::DimId)))
-          return false;
-      }
-    }
-  }
-  return true;
-}
-
 /// Generates a for-loop on a single index.
 static Operation *genFor(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                          linalg::GenericOp op, bool isOuter, bool isInner,
-                         unsigned idx, BitVector &indices) {
-  unsigned fb = indices.find_first();
-  unsigned tensor = merger.tensor(fb);
-  assert(idx == merger.index(fb));
+                         unsigned idx, size_t tid, size_t dim,
+                         ArrayRef<size_t> extraTids,
+                         ArrayRef<size_t> extraDims) {
+  Location loc = op.getLoc();
   auto iteratorTypes = op.getIteratorTypesArray();
   bool isReduction = linalg::isReductionIterator(iteratorTypes[idx]);
-  bool isSparse = merger.isDimLevelType(fb, DimLvlType::kCompressed) ||
-                  merger.isDimLevelType(fb, DimLvlType::kSingleton);
-  bool isVector = isVectorFor(codegen, isInner, isReduction, isSparse) &&
-                  denseUnitStrides(merger, op, idx);
-  bool isParallel =
-      isParallelFor(codegen, isOuter, isReduction, isSparse, isVector);
-
-  // Prepare vector length.
-  if (isVector)
-    codegen.curVecLength = codegen.options.vectorLength;
-
-  // Loop bounds and increment.
-  Location loc = op.getLoc();
-  Value lo = isSparse ? codegen.pidxs[tensor][idx] : codegen.loops[idx];
-  Value hi = isSparse ? codegen.highs[tensor][idx] : codegen.sizes[idx];
-  Value step = constantIndex(builder, loc, codegen.curVecLength);
-  if (isVector && codegen.options.enableVLAVectorization) {
-    Value vscale = builder.create<vector::VectorScaleOp>(
-        loc, IndexType::get(builder.getContext()));
-    step = builder.create<arith::MulIOp>(loc, vscale, step);
-  }
-
-  // Emit a parallel loop.
-  if (isParallel) {
-    assert(!isVector);
-    scf::ParallelOp parOp = builder.create<scf::ParallelOp>(loc, lo, hi, step);
-    if (isSparse)
-      codegen.pidxs[tensor][idx] = parOp.getInductionVars()[0];
-    else
-      codegen.loops[idx] = parOp.getInductionVars()[0];
-    builder.setInsertionPointToStart(parOp.getBody());
-    return parOp;
-  }
+  bool isSparse = merger.isDimLevelType(tid, idx, DimLvlType::kCompressed) ||
+                  merger.isDimLevelType(tid, idx, DimLvlType::kSingleton);
+  bool isParallel = isParallelFor(codegen, isOuter, isReduction, isSparse);
+  assert(!isParallel);
 
   // Emit a sequential or vector loop.
   SmallVector<Value, 4> operands;
-  if (codegen.redVal) {
-    // In a vector loop, bring reduction into SIMD form, if not already.
-    if (isVector && !codegen.redVal.getType().isa<VectorType>()) {
-      VectorType vtp = vectorType(codegen, codegen.redVal.getType());
-      Value vred = genVectorReducInit(codegen, builder, loc, vtp);
-      updateReduc(merger, codegen, vred);
-    }
+  if (codegen.redVal)
     operands.push_back(codegen.redVal);
-  }
   if (codegen.expValues)
     operands.push_back(codegen.expCount);
-  scf::ForOp forOp = builder.create<scf::ForOp>(loc, lo, hi, step, operands);
+
+  Operation *loop = codegen.loopEmitter.enterLoopOverTensorAtDim(
+      builder, loc, tid, dim, operands, false, extraTids, extraDims);
+
+  // The operands should be updated by loop emitter already.
   if (codegen.redVal)
-    updateReduc(merger, codegen, forOp.getRegionIterArgs().front());
+    updateReduc(merger, codegen, operands.front());
   if (codegen.expValues)
-    codegen.expCount = forOp.getRegionIterArgs().back();
-  // Assign induction variable to sparse or dense index.
-  Value iv = forOp.getInductionVar();
-  if (isSparse)
-    codegen.pidxs[tensor][idx] = iv;
-  else
-    codegen.loops[idx] = iv;
-  builder.setInsertionPointToStart(forOp.getBody());
-  // Share vector iteration mask between all subsequent loads/stores.
-  if (isVector)
-    codegen.curVecMask = genVectorMask(codegen, builder, iv, lo, hi, step);
-  return forOp;
+    codegen.expCount = operands.back();
+
+  return loop;
 }
 
 /// Emit a while-loop for co-iteration over multiple indices.
 static Operation *genWhile(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                            linalg::GenericOp op, unsigned idx, bool needsUniv,
-                           BitVector &indices) {
-  SmallVector<Type, 4> types;
+                           ArrayRef<size_t> condTids, ArrayRef<size_t> condDims,
+                           ArrayRef<size_t> extraTids,
+                           ArrayRef<size_t> extraDims) {
+
   SmallVector<Value, 4> operands;
   // Construct the while-loop with a parameter for each index.
-  Type indexType = builder.getIndexType();
-  for (unsigned b = 0, be = indices.size(); b < be; b++) {
-    if (!indices[b])
-      continue;
-    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
-        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
-      unsigned tensor = merger.tensor(b);
-      assert(idx == merger.index(b));
-      types.push_back(indexType);
-      operands.push_back(codegen.pidxs[tensor][idx]);
-    } else {
-      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
-             merger.isDimLevelType(b, DimLvlType::kUndef));
-    }
-  }
-  if (codegen.redVal) {
-    types.push_back(codegen.redVal.getType());
+
+  if (codegen.redVal)
     operands.push_back(codegen.redVal);
-  }
-  if (codegen.expValues) {
-    types.push_back(indexType);
+  if (codegen.expValues)
     operands.push_back(codegen.expCount);
-  }
-  if (needsUniv) {
-    types.push_back(indexType);
-    operands.push_back(codegen.loops[idx]);
-  }
-  assert(types.size() == operands.size());
-  Location loc = op.getLoc();
-  scf::WhileOp whileOp = builder.create<scf::WhileOp>(loc, types, operands);
 
-  SmallVector<Location> locs(types.size(), loc);
-  Block *before = builder.createBlock(&whileOp.getBefore(), {}, types, locs);
-  Block *after = builder.createBlock(&whileOp.getAfter(), {}, types, locs);
-
-  // Build the "before" region, which effectively consists
-  // of a conjunction of "i < upper" tests on all induction.
-  builder.setInsertionPointToStart(&whileOp.getBefore().front());
-  Value cond;
-  unsigned o = 0;
-  for (unsigned b = 0, be = indices.size(); b < be; b++) {
-    if (!indices[b])
-      continue;
-    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
-        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
-      unsigned tensor = merger.tensor(b);
-      assert(idx == merger.index(b));
-      Value op1 = before->getArgument(o);
-      Value op2 = codegen.highs[tensor][idx];
-      Value opc = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
-                                                op1, op2);
-      cond = cond ? builder.create<arith::AndIOp>(loc, cond, opc) : opc;
-      codegen.pidxs[tensor][idx] = after->getArgument(o++);
-    } else {
-      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
-             merger.isDimLevelType(b, DimLvlType::kUndef));
-    }
-  }
+  Operation *loop = codegen.loopEmitter.enterCoiterationOverTensorsAtDims(
+      builder, op.getLoc(), condTids, condDims, operands, needsUniv, extraTids,
+      extraDims);
+
   if (codegen.redVal)
-    updateReduc(merger, codegen, after->getArgument(o++));
+    updateReduc(merger, codegen, operands.front());
   if (codegen.expValues)
-    codegen.expCount = after->getArgument(o++);
-  if (needsUniv)
-    codegen.loops[idx] = after->getArgument(o++);
-  assert(o == operands.size());
-  builder.create<scf::ConditionOp>(loc, cond, before->getArguments());
-  builder.setInsertionPointToStart(&whileOp.getAfter().front());
-  return whileOp;
+    codegen.expCount = operands.back();
+
+  return loop;
 }
 
 /// Generates a for-loop or a while-loop, depending on whether it implements
 /// singleton iteration or co-iteration over the given conjunction.
 static Operation *genLoop(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                           linalg::GenericOp op, unsigned at, bool needsUniv,
-                          BitVector &indices) {
+                          ArrayRef<size_t> condTids, ArrayRef<size_t> condDims,
+                          ArrayRef<size_t> extraTids,
+                          ArrayRef<size_t> extraDims) {
+  assert(condTids.size() == condDims.size());
+  assert(extraTids.size() == extraDims.size());
   unsigned idx = codegen.topSort[at];
-  if (indices.count() == 1) {
+  if (condTids.size() == 1) {
     bool isOuter = at == 0;
     bool isInner = at == codegen.topSort.size() - 1;
-    return genFor(merger, codegen, builder, op, isOuter, isInner, idx, indices);
-  }
-  return genWhile(merger, codegen, builder, op, idx, needsUniv, indices);
-}
-
-/// Generates the local variables for this loop, consisting of the sparse
-/// indices, restored universal dense index, and dense positions.
-static void genLocals(Merger &merger, CodeGen &codegen, OpBuilder &builder,
-                      linalg::GenericOp op, unsigned at, bool needsUniv,
-                      BitVector &locals) {
-  std::vector<unsigned> &topSort(codegen.topSort);
-  Location loc = op.getLoc();
-  unsigned idx = topSort[at];
-
-  // Initialize sparse indices.
-  Value min;
-  for (unsigned b = 0, be = locals.size(); b < be; b++) {
-    if (!locals[b])
-      continue;
-    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
-        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
-      unsigned tensor = merger.tensor(b);
-      assert(idx == merger.index(b));
-      Value ptr = codegen.indices[tensor][idx];
-      Value s = codegen.pidxs[tensor][idx];
-      Value load = genLoad(codegen, builder, loc, ptr, s);
-      codegen.idxs[tensor][idx] = load;
-      if (!needsUniv) {
-        if (min) {
-          Value cmp = builder.create<arith::CmpIOp>(
-              loc, arith::CmpIPredicate::ult, load, min);
-          min = builder.create<arith::SelectOp>(loc, cmp, load, min);
-        } else {
-          min = load;
-        }
-      }
-    } else {
-      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
-             merger.isDimLevelType(b, DimLvlType::kUndef));
-    }
-  }
-
-  // Merge dense universal index over minimum.
-  if (min) {
-    assert(!needsUniv);
-    codegen.loops[idx] = min;
-  }
-
-  // Initialize dense positions. Note that we generate dense indices of the
-  // output tensor unconditionally, since they may not appear in the lattice,
-  // but may be needed for linearized codegen.
-  for (unsigned b = 0, be = locals.size(); b < be; b++) {
-    if ((locals[b] || merger.isOutTensor(b, idx)) &&
-        merger.isDimLevelType(b, DimLvlType::kDense)) {
-      unsigned tensor = merger.tensor(b);
-      assert(idx == merger.index(b));
-      unsigned pat = at;
-      for (; pat != 0; pat--)
-        if (codegen.pidxs[tensor][topSort[pat - 1]])
-          break;
-      Value p = (pat == 0) ? constantIndex(builder, loc, 0)
-                           : codegen.pidxs[tensor][topSort[pat - 1]];
-      codegen.pidxs[tensor][idx] = genAddress(
-          codegen, builder, loc, codegen.sizes[idx], p, codegen.loops[idx]);
-    }
+    return genFor(merger, codegen, builder, op, isOuter, isInner, idx,
+                  condTids.front(), condDims.front(), extraTids, extraDims);
   }
+  return genWhile(merger, codegen, builder, op, idx, needsUniv, condTids,
+                  condDims, extraTids, extraDims);
 }
 
 /// Generates the induction structure for a while-loop.
-static void genWhileInduction(Merger &merger, CodeGen &codegen,
-                              OpBuilder &builder, linalg::GenericOp op,
-                              unsigned idx, bool needsUniv,
-                              BitVector &induction, scf::WhileOp whileOp) {
+static void finalizeWhileOp(Merger &merger, CodeGen &codegen,
+                            unsigned idx, bool needsUniv, BitVector &induction,
+                            scf::WhileOp whileOp) {
   Location loc = op.getLoc();
   // Finalize each else branch of all if statements.
   if (codegen.redVal || codegen.expValues) {
     while (auto ifOp = dyn_cast_or_null<scf::IfOp>(
                builder.getInsertionBlock()->getParentOp())) {
       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());
       builder.create<scf::YieldOp>(loc, yields);
       builder.setInsertionPointAfter(ifOp);
     }
   }
   builder.setInsertionPointToEnd(&whileOp.getAfter().front());
-  // Finalize the induction. Note that the induction could be performed
-  // in the individual if-branches to avoid re-evaluating the conditions.
-  // However, that would result in a rather elaborate forest of yield
-  // instructions during code generation. Moreover, performing the induction
-  // after the if-statements more closely resembles code generated by TACO.
-  unsigned o = 0;
-  SmallVector<Value, 4> operands;
-  Value one = constantIndex(builder, loc, 1);
-  for (unsigned b = 0, be = induction.size(); b < be; b++) {
-    if (!induction[b])
-      continue;
-    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
-        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
-      unsigned tensor = merger.tensor(b);
-      assert(idx == merger.index(b));
-      Value op1 = codegen.idxs[tensor][idx];
-      Value op2 = codegen.loops[idx];
-      Value op3 = codegen.pidxs[tensor][idx];
-      Value cmp = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
-                                                op1, op2);
-      Value add = builder.create<arith::AddIOp>(loc, op3, one);
-      operands.push_back(builder.create<arith::SelectOp>(loc, cmp, add, op3));
-      codegen.pidxs[tensor][idx] = whileOp->getResult(o++);
-    } else {
-      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
-             merger.isDimLevelType(b, DimLvlType::kUndef));
-    }
-  }
-  if (codegen.redVal) {
-    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(
-        builder.create<arith::AddIOp>(loc, codegen.loops[idx], one));
-    codegen.loops[idx] = whileOp->getResult(o++);
-  }
-  assert(o == operands.size());
-  builder.create<scf::YieldOp>(loc, operands);
-  builder.setInsertionPointAfter(whileOp);
-}
-
-/// Generates the induction structure for a for-loop.
-static void genForInduction(Merger &merger, CodeGen &codegen,
-                            OpBuilder &builder, linalg::GenericOp op,
-                            Operation *loop) {
-  Location loc = op.getLoc();
-  unsigned o = 0;
-  SmallVector<Value, 4> operands;
-  if (codegen.redVal) {
-    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)
-    builder.create<scf::YieldOp>(loc, operands);
-  builder.setInsertionPointAfter(loop);
 }
 
 /// Generates a single if-statement within a while-loop.
 static scf::IfOp genIf(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                        linalg::GenericOp op, unsigned idx,
                        BitVector &conditions) {
   Location loc = op.getLoc();
   SmallVector<Type, 4> types;
   Value cond;
   for (unsigned b = 0, be = conditions.size(); b < be; b++) {
     if (!conditions[b])
       continue;
     unsigned tensor = merger.tensor(b);
     assert(idx == merger.index(b));
     Value clause;
     if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
         merger.isDimLevelType(b, DimLvlType::kSingleton)) {
-      Value op1 = codegen.idxs[tensor][idx];
-      Value op2 = codegen.loops[idx];
+      auto dim = codegen.loopIdxToDim[tensor][idx];
+      assert(dim != INVALID_ID);
+      Value op1 = codegen.loopEmitter.getCoord()[tensor][dim];
+      Value op2 = codegen.getLoopIdxValue(idx);
       clause = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, op1,
                                              op2);
     } else {
       assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
              merger.isDimLevelType(b, DimLvlType::kUndef));
       clause = constantI1(builder, loc, true);
     }
     cond = cond ? builder.create<arith::AndIOp>(loc, cond, clause) : clause;

 // Sparse compiler synthesis methods (loop sequence).
 //===----------------------------------------------------------------------===//
 
 /// Starts a loop sequence at given level. Returns true if
 /// the universal loop index must be maintained at this level.
 static bool startLoopSeq(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                          linalg::GenericOp op, unsigned exp, unsigned at,
                          unsigned idx, unsigned ldx, unsigned lts) {
+
   assert(codegen.curVecLength == 1);
-  assert(!codegen.loops[idx]);
+  assert(!codegen.getLoopIdxValue(idx));
   // Emit invariants at this loop sequence level.
   genInvariants(merger, codegen, builder, op, exp, ldx, /*atStart=*/true);
   // Emit access pattern expansion for sparse tensor output.
   genExpansion(merger, codegen, builder, op, at, /*atStart=*/true);
   // Emit further intitialization at this loop sequence level.
   unsigned l0 = merger.set(lts)[0];
-  bool needsUniv =
-      genInit(merger, codegen, builder, op, at, merger.lat(l0).bits);
+  bool needsUniv = false;
+
+  SmallVector<size_t, 4> ts;
+  SmallVector<size_t, 4> ds;
+  for (auto b : merger.lat(l0).bits.set_bits()) {
+    if (merger.isDimLevelType(b, DimLvlType::kDense) ||
+        merger.isDimLevelType(b, DimLvlType::kUndef)) {
+      needsUniv = true;
+    } else {
+      unsigned tensor = merger.tensor(b);
+      assert(idx == merger.index(b));
+      size_t dim = codegen.loopIdxToDim[tensor][idx];
+      assert(dim != INVALID_ID);
+      ts.push_back(tensor);
+      ds.push_back(dim);
+    }
+  }
+
+  codegen.loopEmitter.enterNewLoopSeq(builder, op.getLoc(), ts, ds);
+
   // Maintain the universal index only if it is actually
   // consumed by a subsequent lattice point.
   if (needsUniv) {
     unsigned lsize = merger.set(lts).size();
     for (unsigned i = 1; i < lsize; i++) {
       unsigned li = merger.set(lts)[i];
       if (!merger.hasAnySparse(merger.lat(li).simple))
         return true;
     }
   }
   return false;
 }
 
 /// Starts a single loop in current sequence.
 static Operation *startLoop(Merger &merger, CodeGen &codegen,
                             OpBuilder &builder, linalg::GenericOp op,
                             unsigned at, unsigned li, bool needsUniv) {
   assert(codegen.curVecLength == 1);
+  const BitVector &simple = merger.lat(li).simple;
+  const BitVector &all = merger.lat(li).bits;
+  assert(simple.size() == all.size());
+  // The set of tensors + dims to generate loops on
+  SmallVector<size_t, 4> condTids, condDims;
+  // The set of (dense) tensors that is optimized from condition, yet still need
+  // extra locals to iterate on them.
+  SmallVector<size_t, 4> extraTids, extraDims;
+  // First converts bits to array + dim pair
+  for (unsigned b = 0, e = simple.size(); b < e; b++) {
+    size_t tid = merger.tensor(b);
+    size_t idx = codegen.topSort[at];
+    if (simple.test(b)) {
+      // the simplified condition must be a subset of the original condition.
+      assert(all.test(b));
+      assert(merger.index(b) == idx);
+      if (merger.isDimLevelType(b, DimLvlType::kUndef)) {
+        // This could be a synthetic tensor (for invariants and sparse output
+        // tensor).
+        // In both cases, we mean to generate loops over output tensor.
+        // e.g.,
+        // out[i][j] = invariant;
+        if (merger.getSynTensorID() == tid)
+          tid = merger.getOutTensorID();
+      }
+      auto dim = codegen.loopIdxToDim[tid][idx];
+      if (dim != INVALID_ID) {
+        // dim could be invalid if this is a zero ranked tensor
+        condTids.push_back(tid);
+        condDims.push_back(dim);
+      }
+    } else if ((all.test(b) || merger.isOutTensor(b, idx)) &&
+               merger.isDimLevelType(b, DimLvlType::kDense)) {
+      // Note that we generate dense indices of the output tensor
+      // unconditionally, since they may not appear in the lattice, but may be
+      // needed for linearized codegen.
+      assert(merger.index(b) == idx);
+      // Only dense dimensions should be optimized from conditions.
+      assert(merger.isDimLevelType(tid, idx, DimLvlType::kDense));
+      auto dim = codegen.loopIdxToDim[tid][idx];
+      assert(dim != INVALID_ID);
+      extraTids.push_back(tid);
+      extraDims.push_back(dim);
+    }
+  }
   // Emit the for/while-loop control.
   Operation *loop = genLoop(merger, codegen, builder, op, at, needsUniv,
-                            merger.lat(li).simple);
+                            condTids, condDims, extraTids, extraDims);
   // Emit the locals for this loop.
-  genLocals(merger, codegen, builder, op, at, needsUniv, merger.lat(li).bits);
+  // genLocals(merger, codegen, builder, op, at, needsUniv,
+  // merger.lat(li).bits);
   return loop;
 }
 
 /// Ends a single loop in current sequence. Returns new values for needsUniv.
 static bool endLoop(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                     linalg::GenericOp op, Operation *loop, unsigned idx,
                     unsigned li, bool needsUniv) {
   codegen.curVecLength = 1;
+
   // End a while-loop.
   if (auto whileOp = dyn_cast<scf::WhileOp>(loop)) {
-    genWhileInduction(merger, codegen, builder, op, idx, needsUniv,
-                      merger.lat(li).bits, whileOp);
-    return needsUniv;
+    finalizeWhileOp(merger, codegen, builder, op, idx, needsUniv,
+  } else {
+    needsUniv = false;
   }
-  // End a for-loop.
-  genForInduction(merger, codegen, builder, op, loop);
-  return false;
+
+  SmallVector<Value, 2> reduc;
+  if (codegen.redVal)
+    reduc.push_back(codegen.redVal);
+  if (codegen.expValues)
+    reduc.push_back(codegen.expCount);
+
+  auto loopRet =
+      codegen.loopEmitter.exitCurrentLoop(builder, op.getLoc(), reduc);
+  assert(reduc.size() == loopRet.size());
+
+  if (codegen.redVal)
+    updateReduc(merger, codegen, loopRet.front());
+  if (codegen.expValues)
+    codegen.expCount = loopRet.back();
+
+  return needsUniv;
 }
 
 /// Ends a loop sequence at given level.
 static void endLoopSeq(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                        linalg::GenericOp op, unsigned exp, unsigned at,
                        unsigned idx, unsigned ldx) {
   assert(codegen.curVecLength == 1);
-  assert(codegen.loops[idx]);
-  codegen.loops[idx] = Value();
+  assert(codegen.getLoopIdxValue(idx) == nullptr);
+
+  codegen.loopEmitter.exitCurrentLoopSeq();
   // Bring a pending reduction back from SIMD form when sequence ends.
-  if (codegen.redVal)
-    if (auto vtp = codegen.redVal.getType().dyn_cast<VectorType>())
-      updateReduc(merger, codegen,
-                  genVectorReducEnd(codegen, builder, op.getLoc(), vtp));
+  // TODO: support vectorizaion in loop emitter.
+  //  if (codegen.redVal)
+  //    if (auto vtp = codegen.redVal.getType().dyn_cast<VectorType>())
+  //      updateReduc(merger, codegen,
+  //                  genVectorReducEnd(codegen, builder, op.getLoc(), vtp));
   // Unmark bookkeeping of invariants and loop index.
   genInvariants(merger, codegen, builder, op, exp, ldx, /*atStart=*/false);
   // Finalize access pattern expansion for sparse tensor output.
   genExpansion(merger, codegen, builder, op, at, /*atStart=*/false);
 }
 
 /// Recursively generates code while computing iteration lattices in order
 /// to manage the complexity of implementing co-iteration over unions
static void genResult(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
   if (getSparseTensorEncoding(resType)) {
     // The sparse tensor rematerializes from the original sparse tensor's
     // underlying sparse storage format.
     rewriter.replaceOpWithNewOp<LoadOp>(op, resType, lhs->get(),
                                         codegen.sparseOut == lhs);
   } else {
     // To rematerialize an non-annotated tensor, simply load it
     // from the bufferized value.
-    Value val = codegen.buffers.back(); // value array
+    Value val = codegen.loopEmitter.getValBuffer().back(); // value array
     rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, resType, val);
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler rewriting methods.
 //===----------------------------------------------------------------------===//
 
LogicalResult matchAndRewrite(linalg::GenericOp op,
 
     if (hasCycle)
       // Give it one last shot to resolve the cycle.
       return resolveCycle(merger, rewriter, op);
     if (!isAdmissible)
       // Inadmissible expression, reject.
       return failure();
 
-    // Recursively generates code if admissible.
     merger.setHasSparseOut(sparseOut != nullptr);
-    CodeGen codegen(options, numTensors, numLoops, sparseOut, outerParNest,
-                    topSort);
+
+    SmallVector<Value, 4> tensors;
+    for (OpOperand &t : op->getOpOperands())
+      tensors.push_back(t.get());
+
+    // Recursively generates code if admissible.
+    CodeGen codegen(options, tensors, numTensors, numLoops, sparseOut,
+                    outerParNest, topSort);
+    // TODO: maybe merger should be responsible of maintaining the map.
+    codegen.buildLoopIdxToDimMap(op);
     genBuffers(merger, codegen, rewriter, op);
     genStmt(merger, codegen, rewriter, op, exp, 0);
     genResult(merger, codegen, rewriter, op);
     return success();
   }
 
 private:
   // Last resort cycle resolution.

mlir/test/Dialect/SparseTensor/sorted_coo.mlir

 // CHECK-SAME:    %[[VAL_1:.*]]: tensor<64xf64>,
 // CHECK-SAME:    %[[VAL_2:.*]]: tensor<32xf64>) -> tensor<32xf64> {
 // CHECK-DAG:     %[[VAL_3:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[VAL_4:.*]] = arith.constant 1 : index
 // CHECK-DAG:     %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
 // CHECK-DAG:     %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
 // CHECK-DAG:     %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
 // CHECK-DAG:     %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xf64>
-// CHECK:         %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<64xf64>
-// CHECK:         %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32xf64>
-// CHECK:         %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:         %[[VAL_12:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:     %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<64xf64>
+// CHECK-DAG:     %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32xf64>
+// CHECK-DAG:     %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK-DAG:     %[[VAL_12:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
 // CHECK:         scf.for %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_4]] {
 // CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_13]]] : memref<?xindex>
 // CHECK:           %[[VAL_15:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_14]]] : memref<32xf64>
 // CHECK:           %[[VAL_16:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xindex>
 // CHECK:           %[[VAL_17:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_13]]] : memref<?xf64>
 // CHECK:           %[[VAL_18:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_16]]] : memref<64xf64>
 // CHECK:           %[[VAL_19:.*]] = arith.mulf %[[VAL_17]], %[[VAL_18]] : f64
 // CHECK:           %[[VAL_20:.*]] = arith.addf %[[VAL_15]], %[[VAL_19]] : f64

 // CHECK-DAG:     %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
 // CHECK-DAG:     %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
 // CHECK-DAG:     %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
 // CHECK-DAG:     %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xf64>
 // CHECK-DAG:     %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
 // CHECK-DAG:     %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
 // CHECK-DAG:     %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
 // CHECK-DAG:     %[[VAL_13:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xf64>
-// CHECK:         %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x64xf64>
-// CHECK:         linalg.fill ins(%[[VAL_3]] : f64) outs(%[[VAL_14]] : memref<32x64xf64>)
-// CHECK:         %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:         %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
-// CHECK:         %[[VAL_17:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:         %[[VAL_18:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK-DAG:     %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x64xf64>
+// CHECK-DAG:     linalg.fill ins(%[[VAL_3]] : f64) outs(%[[VAL_14]] : memref<32x64xf64>)
+// CHECK-DAG:     %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:     %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK-DAG:     %[[VAL_17:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:     %[[VAL_18:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_5]]] : memref<?xindex>
 // CHECK:         %[[VAL_19:.*]]:2 = scf.while (%[[VAL_20:.*]] = %[[VAL_15]], %[[VAL_21:.*]] = %[[VAL_17]]) : (index, index) -> (index, index) {
 // CHECK:           %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_16]] : index
 // CHECK:           %[[VAL_23:.*]] = arith.cmpi ult, %[[VAL_21]], %[[VAL_18]] : index
 // CHECK:           %[[VAL_24:.*]] = arith.andi %[[VAL_22]], %[[VAL_23]] : i1
 // CHECK:           scf.condition(%[[VAL_24]]) %[[VAL_20]], %[[VAL_21]] : index, index
 // CHECK:         } do {
 // CHECK:         ^bb0(%[[VAL_25:.*]]: index, %[[VAL_26:.*]]: index):
 // CHECK:           %[[VAL_27:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_25]]] : memref<?xindex>

mlir/test/Dialect/SparseTensor/sparse_1d.mlir

 // CHECK-SAME:      %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
 // CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 32 : index
 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = arith.constant true
 // CHECK-DAG:       %[[VAL_6:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
+// CHECK-DAG:       %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
 // CHECK-DAG:       %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]]
-// CHECK:           linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_11]] : memref<32xf32>)
-// CHECK:           %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
+// CHECK-DAG:       linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_11]] : memref<32xf32>)
 // CHECK:           %[[VAL_14:.*]]:2 = scf.while (%[[VAL_15:.*]] = %[[VAL_12]], %[[VAL_16:.*]] = %[[VAL_4]]) : (index, index) -> (index, index) {
 // CHECK:             %[[VAL_17:.*]] = arith.cmpi ult, %[[VAL_15]], %[[VAL_13]] : index
 // CHECK:             scf.condition(%[[VAL_17]]) %[[VAL_15]], %[[VAL_16]] : index, index
 // CHECK:           } do {
 // CHECK:           ^bb0(%[[VAL_18:.*]]: index, %[[VAL_19:.*]]: index):
 // CHECK:             %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xindex>
 // CHECK:             %[[VAL_21:.*]] = arith.cmpi eq, %[[VAL_20]], %[[VAL_19]] : index
 // CHECK:             scf.if %[[VAL_21]] {

 // CHECK-SAME:      %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
 // CHECK-SAME:      %[[VAL_1:.*]]: tensor<32xf32>) -> tensor<32xf32> {
 // CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_4:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_5:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]]
-// CHECK:           linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_8]] : memref<32xf32>)
-// CHECK:           %[[VAL_9:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex>
-// CHECK:           %[[VAL_10:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_9:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_10:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK-DAG:       linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_8]] : memref<32xf32>)
 // CHECK:           scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_3]] {
 // CHECK:             %[[VAL_12:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xindex>
 // CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf32>
 // CHECK:             %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf32>
 // CHECK:             %[[VAL_15:.*]] = arith.addf %[[VAL_13]], %[[VAL_14]] : f32
 // CHECK:             %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf32>
 // CHECK:             %[[VAL_17:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf32>
 // CHECK:             %[[VAL_18:.*]] = arith.addf %[[VAL_16]], %[[VAL_17]] : f32

 // CHECK-SAME:      %[[VAL_1:.*]]: f32,
 // CHECK-SAME:      %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
 // CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_2]]
-// CHECK:           linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_9]] : memref<32xf32>)
-// CHECK:           %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:           %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_9]] : memref<32xf32>)
+// CHECK-DAG:       %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
 // CHECK:           scf.for %[[VAL_12:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_4]] {
 // CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex>
 // CHECK:             %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_12]]] : memref<?xf32>
 // CHECK:             %[[VAL_15:.*]] = arith.mulf %[[VAL_14]], %[[VAL_1]] : f32
 // CHECK:             memref.store %[[VAL_15]], %[[VAL_9]]{{\[}}%[[VAL_13]]] : memref<32xf32>
 // CHECK:           }
 // CHECK:           %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xf32>
 // CHECK:           return %[[VAL_16]] : tensor<32xf32>

 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = arith.constant true
 // CHECK-DAG:       %[[VAL_6:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_7:.*]] = bufferization.to_memref %[[VAL_0]] : memref<32xf32>
 // CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]]
-// CHECK:           linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_12]] : memref<32xf32>)
-// CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_6]]] : memref<?xindex>
+// CHECK-DAG:       linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_12]] : memref<32xf32>)
+// CHECK-DAG:       %[[VAL_13:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_14:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_6]]] : memref<?xindex>
 // CHECK:           %[[VAL_15:.*]]:2 = scf.while (%[[VAL_16:.*]] = %[[VAL_13]], %[[VAL_17:.*]] = %[[VAL_4]]) : (index, index) -> (index, index) {
 // CHECK:             %[[VAL_18:.*]] = arith.cmpi ult, %[[VAL_16]], %[[VAL_14]] : index
 // CHECK:             scf.condition(%[[VAL_18]]) %[[VAL_16]], %[[VAL_17]] : index, index
 // CHECK:           } do {
 // CHECK:           ^bb0(%[[VAL_19:.*]]: index, %[[VAL_20:.*]]: index):
 // CHECK:             %[[VAL_21:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_19]]] : memref<?xindex>
 // CHECK:             %[[VAL_22:.*]] = arith.cmpi eq, %[[VAL_21]], %[[VAL_20]] : index
 // CHECK:             scf.if %[[VAL_22]] {

 // CHECK-SAME:      %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
 // CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = bufferization.to_memref %[[VAL_0]] : memref<32xf32>
 // CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]]
-// CHECK:           linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_10]] : memref<32xf32>)
-// CHECK:           %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:           %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_10]] : memref<32xf32>)
+// CHECK-DAG:       %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
 // CHECK:           scf.for %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_4]] {
 // CHECK:             %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xindex>
 // CHECK:             %[[VAL_15:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_14]]] : memref<32xf32>
 // CHECK:             %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_13]]] : memref<?xf32>
 // CHECK:             %[[VAL_17:.*]] = arith.mulf %[[VAL_15]], %[[VAL_16]] : f32
 // CHECK:             memref.store %[[VAL_17]], %[[VAL_10]]{{\[}}%[[VAL_14]]] : memref<32xf32>
 // CHECK:           }
 // CHECK:           %[[VAL_18:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32xf32>

 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = arith.constant true
 // CHECK-DAG:       %[[VAL_6:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32xf32>
 // CHECK-DAG:       %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]]
-// CHECK:           linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_12]] : memref<32xf32>)
-// CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
+// CHECK-DAG:       linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_12]] : memref<32xf32>)
+// CHECK-DAG:       %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
 // CHECK:           %[[VAL_15:.*]]:2 = scf.while (%[[VAL_16:.*]] = %[[VAL_13]], %[[VAL_17:.*]] = %[[VAL_4]]) : (index, index) -> (index, index) {
 // CHECK:             %[[VAL_18:.*]] = arith.cmpi ult, %[[VAL_16]], %[[VAL_14]] : index
 // CHECK:             scf.condition(%[[VAL_18]]) %[[VAL_16]], %[[VAL_17]] : index, index
 // CHECK:           } do {
 // CHECK:           ^bb0(%[[VAL_19:.*]]: index, %[[VAL_20:.*]]: index):
 // CHECK:             %[[VAL_21:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_19]]] : memref<?xindex>
 // CHECK:             %[[VAL_22:.*]] = arith.cmpi eq, %[[VAL_21]], %[[VAL_20]] : index
 // CHECK:             scf.if %[[VAL_22]] {

 // CHECK-SAME:      %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
 // CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32xf32>
 // CHECK-DAG:       %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]]
-// CHECK:           linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_10]] : memref<32xf32>)
-// CHECK:           %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:           %[[VAL_12:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_10]] : memref<32xf32>)
+// CHECK-DAG:       %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_12:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
 // CHECK:           scf.for %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_4]] {
 // CHECK:             %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_13]]] : memref<?xindex>
 // CHECK:             %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xf32>
 // CHECK:             %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_14]]] : memref<32xf32>
 // CHECK:             %[[VAL_17:.*]] = arith.mulf %[[VAL_15]], %[[VAL_16]] : f32
 // CHECK:             memref.store %[[VAL_17]], %[[VAL_10]]{{\[}}%[[VAL_14]]] : memref<32xf32>
 // CHECK:           }
 // CHECK:           %[[VAL_18:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32xf32>

 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]]
-// CHECK:           linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_12]] : memref<32xf32>)
-// CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:           %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_12]] : memref<32xf32>)
+// CHECK-DAG:       %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_14:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex>
 // CHECK:           %[[VAL_17:.*]]:2 = scf.while (%[[VAL_18:.*]] = %[[VAL_13]], %[[VAL_19:.*]] = %[[VAL_15]]) : (index, index) -> (index, index) {
 // CHECK:             %[[VAL_20:.*]] = arith.cmpi ult, %[[VAL_18]], %[[VAL_14]] : index
 // CHECK:             %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_16]] : index
 // CHECK:             %[[VAL_22:.*]] = arith.andi %[[VAL_20]], %[[VAL_21]] : i1
 // CHECK:             scf.condition(%[[VAL_22]]) %[[VAL_18]], %[[VAL_19]] : index, index
 // CHECK:           } do {
 // CHECK:           ^bb0(%[[VAL_23:.*]]: index, %[[VAL_24:.*]]: index):
 // CHECK:             %[[VAL_25:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_23]]] : memref<?xindex>

 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]]
-// CHECK:           linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_12]] : memref<32xf32>)
-// CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:           %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_12]] : memref<32xf32>)
+// CHECK-DAG:       %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_14:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex>
 // CHECK:           %[[VAL_17:.*]]:2 = scf.while (%[[VAL_18:.*]] = %[[VAL_13]], %[[VAL_19:.*]] = %[[VAL_15]]) : (index, index) -> (index, index) {
 // CHECK:             %[[VAL_20:.*]] = arith.cmpi ult, %[[VAL_18]], %[[VAL_14]] : index
 // CHECK:             %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_16]] : index
 // CHECK:             %[[VAL_22:.*]] = arith.andi %[[VAL_20]], %[[VAL_21]] : i1
 // CHECK:             scf.condition(%[[VAL_22]]) %[[VAL_18]], %[[VAL_19]] : index, index
 // CHECK:           } do {
 // CHECK:           ^bb0(%[[VAL_23:.*]]: index, %[[VAL_24:.*]]: index):
 // CHECK:             %[[VAL_25:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_23]]] : memref<?xindex>

 // CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_3]]
-// CHECK:           linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_13]] : memref<16xf32>)
-// CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
-// CHECK:           %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_17:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK-DAG:       linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_13]] : memref<16xf32>)
+// CHECK-DAG:       %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_17:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex>
 // CHECK:           %[[VAL_18:.*]]:2 = scf.while (%[[VAL_19:.*]] = %[[VAL_14]], %[[VAL_20:.*]] = %[[VAL_16]]) : (index, index) -> (index, index) {
 // CHECK:             %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_15]] : index
 // CHECK:             %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_17]] : index
 // CHECK:             %[[VAL_23:.*]] = arith.andi %[[VAL_21]], %[[VAL_22]] : i1
 // CHECK:             scf.condition(%[[VAL_23]]) %[[VAL_19]], %[[VAL_20]] : index, index
 // CHECK:           } do {
 // CHECK:           ^bb0(%[[VAL_24:.*]]: index, %[[VAL_25:.*]]: index):
 // CHECK:             %[[VAL_26:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_24]]] : memref<?xindex>

 // CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_3]]
-// CHECK:           linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_13]] : memref<16xf32>)
-// CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
-// CHECK:           %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_17:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK-DAG:       linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_13]] : memref<16xf32>)
+// CHECK-DAG:       %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_17:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex>
 // CHECK:           %[[VAL_18:.*]]:2 = scf.while (%[[VAL_19:.*]] = %[[VAL_14]], %[[VAL_20:.*]] = %[[VAL_16]]) : (index, index) -> (index, index) {
 // CHECK:             %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_15]] : index
 // CHECK:             %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_17]] : index
 // CHECK:             %[[VAL_23:.*]] = arith.andi %[[VAL_21]], %[[VAL_22]] : i1
 // CHECK:             scf.condition(%[[VAL_23]]) %[[VAL_19]], %[[VAL_20]] : index, index
 // CHECK:           } do {
 // CHECK:           ^bb0(%[[VAL_24:.*]]: index, %[[VAL_25:.*]]: index):
 // CHECK:             %[[VAL_26:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_24]]] : memref<?xindex>

 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<f32>
-// CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_11]][] : memref<f32>
-// CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:           %[[VAL_15:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:           %[[VAL_17:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:           %[[VAL_13:.*]] = memref.load %[[VAL_11]][] : memref<f32>
+// CHECK-DAG:           %[[VAL_14:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK-DAG:           %[[VAL_15:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:           %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK-DAG:           %[[VAL_17:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex>
 // CHECK:           %[[VAL_18:.*]]:3 = scf.while (%[[VAL_19:.*]] = %[[VAL_14]], %[[VAL_20:.*]] = %[[VAL_16]], %[[VAL_21:.*]] = %[[VAL_13]]) : (index, index, f32) -> (index, index, f32) {
 // CHECK:             %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_15]] : index
 // CHECK:             %[[VAL_23:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_17]] : index
 // CHECK:             %[[VAL_24:.*]] = arith.andi %[[VAL_22]], %[[VAL_23]] : i1
 // CHECK:             scf.condition(%[[VAL_24]]) %[[VAL_19]], %[[VAL_20]], %[[VAL_21]] : index, index, f32
 // CHECK:           } do {
 // CHECK:           ^bb0(%[[VAL_25:.*]]: index, %[[VAL_26:.*]]: index, %[[VAL_27:.*]]: f32):
 // CHECK:             %[[VAL_28:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_25]]] : memref<?xindex>

 // CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<f32>
 // CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_2]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_2]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_12:.*]] = sparse_tensor.values %[[VAL_2]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_3]] : memref<f32>
-// CHECK:           %[[VAL_15:.*]] = memref.load %[[VAL_13]][] : memref<f32>
-// CHECK:           %[[VAL_16:.*]] = memref.load %[[VAL_9]][] : memref<f32>
-// CHECK:           %[[VAL_17:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_18:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
-// CHECK:           %[[VAL_19:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_20:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_15:.*]] = memref.load %[[VAL_13]][] : memref<f32>
+// CHECK-DAG:       %[[VAL_16:.*]] = memref.load %[[VAL_9]][] : memref<f32>
+// CHECK-DAG:       %[[VAL_17:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_18:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_19:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_20:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_5]]] : memref<?xindex>
 // CHECK:           %[[VAL_21:.*]]:3 = scf.while (%[[VAL_22:.*]] = %[[VAL_17]], %[[VAL_23:.*]] = %[[VAL_19]], %[[VAL_24:.*]] = %[[VAL_15]]) : (index, index, f32) -> (index, index, f32) {
 // CHECK:             %[[VAL_25:.*]] = arith.cmpi ult, %[[VAL_22]], %[[VAL_18]] : index
 // CHECK:             %[[VAL_26:.*]] = arith.cmpi ult, %[[VAL_23]], %[[VAL_20]] : index
 // CHECK:             %[[VAL_27:.*]] = arith.andi %[[VAL_25]], %[[VAL_26]] : i1
 // CHECK:             scf.condition(%[[VAL_27]]) %[[VAL_22]], %[[VAL_23]], %[[VAL_24]] : index, index, f32
 // CHECK:           } do {
 // CHECK:           ^bb0(%[[VAL_28:.*]]: index, %[[VAL_29:.*]]: index, %[[VAL_30:.*]]: f32):
 // CHECK:             %[[VAL_31:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_28]]] : memref<?xindex>

 // CHECK-DAG:       %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_0]] : memref<?xf64>
 // CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf64>
 // CHECK-DAG:       %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] : memref<?xf64>
 // CHECK-DAG:       %[[VAL_13:.*]] = sparse_tensor.pointers %[[VAL_3]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_14:.*]] = sparse_tensor.indices %[[VAL_3]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_15:.*]] = sparse_tensor.values %[[VAL_3]] : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf64>
-// CHECK-DAG:       %[[VAL_16:.*]] = tensor.dim %[[VAL_4]], %[[VAL_5]] : tensor<?xf64>
+// CHECK-DAG:       %[[VAL_16:.*]] = tensor.dim %[[VAL_0]], %[[VAL_5]] : tensor<?xf64>
 // CHECK-DAG:       %[[VAL_18:.*]] = bufferization.to_memref %[[VAL_4]]
-// CHECK:           linalg.fill ins(%{{.*}} : f64) outs(%[[VAL_18]] : memref<?xf64>)
-// CHECK:           %[[VAL_19:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex>
-// CHECK:           %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_7]]] : memref<?xindex>
-// CHECK:           %[[VAL_21:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_5]]] : memref<?xindex>
-// CHECK:           %[[VAL_22:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_7]]] : memref<?xindex>
+// CHECK-DAG:       linalg.fill ins(%{{.*}} : f64) outs(%[[VAL_18]] : memref<?xf64>)
+// CHECK-DAG:       %[[VAL_19:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_7]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_21:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_22:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_7]]] : memref<?xindex>
 // CHECK:           %[[VAL_23:.*]]:3 = scf.while (%[[VAL_24:.*]] = %[[VAL_19]], %[[VAL_25:.*]] = %[[VAL_21]], %[[VAL_26:.*]] = %[[VAL_5]]) : (index, index, index) -> (index, index, index) {
 // CHECK:             %[[VAL_27:.*]] = arith.cmpi ult, %[[VAL_24]], %[[VAL_20]] : index
 // CHECK:             %[[VAL_28:.*]] = arith.cmpi ult, %[[VAL_25]], %[[VAL_22]] : index
 // CHECK:             %[[VAL_29:.*]] = arith.andi %[[VAL_27]], %[[VAL_28]] : i1
 // CHECK:             scf.condition(%[[VAL_29]]) %[[VAL_24]], %[[VAL_25]], %[[VAL_26]] : index, index, index
 // CHECK:           } do {
 // CHECK:           ^bb0(%[[VAL_30:.*]]: index, %[[VAL_31:.*]]: index, %[[VAL_32:.*]]: index):
 // CHECK:             %[[VAL_33:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_30]]] : memref<?xindex>

 // CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf64>
 // CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf64>
 // CHECK-DAG:       %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_2]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_13:.*]] = sparse_tensor.indices %[[VAL_2]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_14:.*]] = sparse_tensor.values %[[VAL_2]] : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf64>
 // CHECK-DAG:       %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_3]] : memref<f64>
-// CHECK:           %[[VAL_17:.*]] = memref.load %[[VAL_15]][] : memref<f64>
-// CHECK:           %[[VAL_18:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_19:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
-// CHECK:           %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_21:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex>
-// CHECK:           %[[VAL_22:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           %[[VAL_23:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_17:.*]] = memref.load %[[VAL_15]][] : memref<f64>
+// CHECK-DAG:       %[[VAL_18:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_19:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_21:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_22:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_23:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_5]]] : memref<?xindex>
 // CHECK:           %[[VAL_24:.*]]:4 = scf.while (%[[VAL_25:.*]] = %[[VAL_18]], %[[VAL_26:.*]] = %[[VAL_20]], %[[VAL_27:.*]] = %[[VAL_22]], %[[VAL_28:.*]] = %[[VAL_17]]) : (index, index, index, f64) -> (index, index, index, f64) {
 // CHECK:             %[[VAL_29:.*]] = arith.cmpi ult, %[[VAL_25]], %[[VAL_19]] : index
 // CHECK:             %[[VAL_30:.*]] = arith.cmpi ult, %[[VAL_26]], %[[VAL_21]] : index
 // CHECK:             %[[VAL_31:.*]] = arith.andi %[[VAL_29]], %[[VAL_30]] : i1
 // CHECK:             %[[VAL_32:.*]] = arith.cmpi ult, %[[VAL_27]], %[[VAL_23]] : index
 // CHECK:             %[[VAL_33:.*]] = arith.andi %[[VAL_31]], %[[VAL_32]] : i1
 // CHECK:             scf.condition(%[[VAL_33]]) %[[VAL_25]], %[[VAL_26]], %[[VAL_27]], %[[VAL_28]] : index, index, index, f64
 // CHECK:           } do {

mlir/test/Dialect/SparseTensor/sparse_2d.mlir

 // CHECK-SAME:      %[[VAL_0:.*]]: tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>,
 // CHECK-SAME:      %[[VAL_1:.*]]: tensor<?x?xf64>) -> tensor<?x?xf64> {
 // CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 2.000000e+00 : f64
 // CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xf64>
-// CHECK-DAG:       %[[VAL_8:.*]] = tensor.dim %[[VAL_1]], %[[VAL_3]] : tensor<?x?xf64>
+// CHECK-DAG:       %[[VAL_8:.*]] = tensor.dim %[[VAL_0]], %[[VAL_3]] : tensor<?x?xf64
 // CHECK-DAG:       %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_1]] : memref<?x?xf64>
 // CHECK:           linalg.fill ins(%{{.*}} : f64) outs(%[[VAL_11]] : memref<?x?xf64>)
 // CHECK:           scf.for %[[VAL_12:.*]] = %[[VAL_3]] to %[[VAL_8]] step %[[VAL_4]] {
 // CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_12]]] : memref<?xindex>
 // CHECK:             %[[VAL_14:.*]] = arith.addi %[[VAL_12]], %[[VAL_4]] : index
 // CHECK:             %[[VAL_15:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_14]]] : memref<?xindex>
 // CHECK:             scf.for %[[VAL_16:.*]] = %[[VAL_13]] to %[[VAL_15]] step %[[VAL_4]] {
 // CHECK:               %[[VAL_17:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_16]]] : memref<?xindex>

 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_1]] : memref<?x?xf32>
-// CHECK-DAG:       %[[VAL_12:.*]] = tensor.dim %[[VAL_2]], %[[VAL_4]] : tensor<?x?xf32>
+// CHECK-DAG:       %[[VAL_12:.*]] = tensor.dim %[[VAL_1]], %[[VAL_5]] : tensor<?x?xf32>
 // CHECK-DAG:       %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_2]] : memref<?x?xf32>
 // CHECK-DAG:       %[[VAL_17:.*]] = bufferization.to_memref %[[VAL_3]] : memref<?x?xf32>
 // CHECK:           %[[VAL_18:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
 // CHECK:           %[[VAL_19:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
 // CHECK:           scf.for %[[VAL_20:.*]] = %[[VAL_18]] to %[[VAL_19]] step %[[VAL_5]] {
 // CHECK:             %[[VAL_21:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_20]]] : memref<?xindex>
 // CHECK:             %[[VAL_22:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_20]]] : memref<?xindex>
 // CHECK:             %[[VAL_23:.*]] = arith.addi %[[VAL_20]], %[[VAL_5]] : index

 // CHECK-DAG:       %[[VAL_14:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 1 : index} : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_15:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
 // CHECK:           %[[VAL_16:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_17:.*]] = sparse_tensor.pointers %[[VAL_2]] {dimension = 1 : index} : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_18:.*]] = sparse_tensor.indices %[[VAL_2]] {dimension = 1 : index} : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_19:.*]] = sparse_tensor.values %[[VAL_2]] : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xf32>
 // CHECK-DAG:       %[[VAL_20:.*]] = bufferization.to_memref %[[VAL_3]] : memref<?xf32>
 // CHECK-DAG:       %[[VAL_21:.*]] = bufferization.to_memref %[[VAL_4]] : memref<f32>
-// CHECK-DAG:       %[[VAL_22:.*]] = tensor.dim %[[VAL_5]], %[[VAL_6]] : tensor<?xf32>
+// CHECK-DAG:       %[[VAL_22:.*]] = tensor.dim %[[VAL_2]], %[[VAL_6]] : tensor<?x?xf32,
 // CHECK-DAG:       %[[VAL_24:.*]] = bufferization.to_memref %[[VAL_5]] : memref<?xf32>
 // CHECK:           %[[VAL_25:.*]] = memref.load %[[VAL_21]][] : memref<f32>
 // CHECK:           %[[VAL_26:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_6]]] : memref<?xindex>
 // CHECK:           %[[VAL_27:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_7]]] : memref<?xindex>
 // CHECK:           %[[VAL_28:.*]]:2 = scf.while (%[[VAL_29:.*]] = %[[VAL_26]], %[[VAL_30:.*]] = %[[VAL_6]]) : (index, index) -> (index, index) {
 // CHECK:             %[[VAL_31:.*]] = arith.cmpi ult, %[[VAL_29]], %[[VAL_27]] : index
 // CHECK:             scf.condition(%[[VAL_31]]) %[[VAL_29]], %[[VAL_30]] : index, index
 // CHECK:           } do {

mlir/test/Dialect/SparseTensor/sparse_3d.mlir

 // CHECK-SAME:      %[[VAL_1:.*1]]: tensor<?x?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>>,
 // CHECK-SAME:      %[[VAL_2:.*2]]: tensor<?x?xf32>,
 // CHECK-SAME:      %[[VAL_3:.*3]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
 // CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_6:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 2 : index} : tensor<?x?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 2 : index} : tensor<?x?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xindex>
 // CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?x?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xf32>
-// CHECK-DAG:       %[[VAL_10:.*]] = tensor.dim %[[VAL_2]], %[[VAL_5]] : tensor<?x?xf32>
+// CHECK-DAG:       %[[VAL_10:.*]] = tensor.dim %[[VAL_1]], %[[VAL_6]] : tensor<?x?x?xf32
 // CHECK-DAG:       %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<?x?xf32>
 // CHECK-DAG:       %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_3]] : memref<?x?xf32>
-// CHECK-DAG:       %[[VAL_13:.*]] = tensor.dim %[[VAL_0]], %[[VAL_5]] : tensor<?x?xf32>
-// CHECK-DAG:       %[[VAL_14:.*]] = tensor.dim %[[VAL_0]], %[[VAL_6]] : tensor<?x?xf32>
+// CHECK-DAG:       %[[VAL_13:.*]] = tensor.dim %[[VAL_1]], %[[VAL_5]] : tensor<?x?x?xf32
+// CHECK-DAG:       %[[VAL_14:.*]] = tensor.dim %[[VAL_2]], %[[VAL_6]] : tensor<?x?xf32>
 // CHECK-DAG:       %[[VAL_16:.*]] = bufferization.to_memref %[[VAL_0]] : memref<?x?xf32>
 // CHECK:           scf.for %[[VAL_17:.*]] = %[[VAL_5]] to %[[VAL_13]] step %[[VAL_6]] {
 // CHECK:             scf.for %[[VAL_18:.*]] = %[[VAL_5]] to %[[VAL_10]] step %[[VAL_6]] {
 // CHECK:               %[[VAL_19:.*]] = arith.muli %[[VAL_10]], %[[VAL_17]] : index
 // CHECK:               %[[VAL_20:.*]] = arith.addi %[[VAL_19]], %[[VAL_18]] : index
 // CHECK:               %[[VAL_21:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_20]]] : memref<?xindex>
 // CHECK:               %[[VAL_22:.*]] = arith.addi %[[VAL_20]], %[[VAL_6]] : index
 // CHECK:               %[[VAL_23:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_22]]] : memref<?xindex>

 // CHECK-SAME:      %[[VAL_1:.*]]: tensor<?xf32, #sparse_tensor.encoding<{{{.*}}}>>
 // CHECK-SAME:      %[[VAL_2:.*]]: tensor<f32>) -> tensor<f32> {
 // CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 2 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_6:.*]] = tensor.dim %[[VAL_0]], %[[VAL_3]] : tensor<?x?x?xf32>
 // CHECK-DAG:       %[[VAL_7:.*]] = tensor.dim %[[VAL_0]], %[[VAL_4]] : tensor<?x?x?xf32>
 // CHECK-DAG:       %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_0]] : memref<?x?x?xf32>
-// CHECK-DAG:       %[[VAL_9:.*]] = tensor.dim %[[VAL_1]], %[[VAL_5]] : tensor<?xf32, #sparse_tensor.encoding<{{{.*}}}>>
+// CHECK-DAG:       %[[VAL_9:.*]] = tensor.dim %[[VAL_0]], %[[VAL_5]] : tensor<?x?x?xf32>
 // CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xf32, #sparse_tensor.encoding<{{{.*}}}>>
 // CHECK-DAG:       %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] : memref<f32>
 // CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_12]][] : memref<f32>
 // CHECK:           %[[VAL_14:.*]] = scf.for %[[VAL_15:.*]] = %[[VAL_5]] to %[[VAL_9]] step %[[VAL_3]] iter_args(%[[VAL_16:.*]] = %[[VAL_13]]) -> (f32) {
 // CHECK:             %[[VAL_17:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_15]]] : memref<?xf32>
 // CHECK:             %[[VAL_18:.*]] = scf.for %[[VAL_19:.*]] = %[[VAL_5]] to %[[VAL_6]] step %[[VAL_3]] iter_args(%[[VAL_20:.*]] = %[[VAL_16]]) -> (f32) {
 // CHECK:               %[[VAL_21:.*]] = scf.for %[[VAL_22:.*]] = %[[VAL_5]] to %[[VAL_7]] step %[[VAL_3]] iter_args(%[[VAL_23:.*]] = %[[VAL_20]]) -> (f32) {
 // CHECK:                 %[[VAL_24:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_15]], %[[VAL_19]], %[[VAL_22]]] : memref<?x?x?xf32>

mlir/test/Dialect/SparseTensor/sparse_concat_codegen.mlir

 //       CHECK:  %[[TMP_2:.*]] = sparse_tensor.indices %[[TMP_arg0]] {dimension = 0 : index} : tensor<2x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_3:.*]] = sparse_tensor.pointers %[[TMP_arg0]] {dimension = 1 : index} : tensor<2x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_4:.*]] = sparse_tensor.indices %[[TMP_arg0]] {dimension = 1 : index} : tensor<2x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_5:.*]] = sparse_tensor.values %[[TMP_arg0]] : tensor<2x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_6:.*]] = memref.load %[[TMP_1]][%[[TMP_c0]]] : memref<?xindex>
 //       CHECK:  %[[TMP_7:.*]] = memref.load %[[TMP_1]][%[[TMP_c1]]] : memref<?xindex>
 //       CHECK:  scf.for %[[TMP_arg3:.*]] = %[[TMP_6]] to %[[TMP_7]] step %[[TMP_c1]] {
 //       CHECK:    %[[TMP_23:.*]] = memref.load %[[TMP_2]][%[[TMP_arg3]]] : memref<?xindex>
-//       CHECK:    %[[TMP_24:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
 //       CHECK:    %[[TMP_25:.*]] = memref.load %[[TMP_3]][%[[TMP_arg3]]] : memref<?xindex>
+//       CHECK:    %[[TMP_24:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
 //       CHECK:    %[[TMP_26:.*]] = memref.load %[[TMP_3]][%[[TMP_24]]] : memref<?xindex>
 //       CHECK:    scf.for %[[TMP_arg4:.*]] = %[[TMP_25]] to %[[TMP_26]] step %[[TMP_c1]] {
 //       CHECK:      %[[TMP_27:.*]] = memref.load %[[TMP_4]][%[[TMP_arg4]]] : memref<?xindex>
 //       CHECK:      %[[TMP_28:.*]] = memref.load %[[TMP_5]][%[[TMP_arg4]]] : memref<?xf64>
 //       CHECK:      sparse_tensor.insert %[[TMP_28]] into %[[TMP_0]][%[[TMP_23]], %[[TMP_27]]] : tensor<9x4xf64, #sparse_tensor
 //       CHECK:    }
 //       CHECK:  }
 //       CHECK:  %[[TMP_8:.*]] = sparse_tensor.pointers %[[TMP_arg1]] {dimension = 0 : index} : tensor<3x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_9:.*]] = sparse_tensor.indices %[[TMP_arg1]] {dimension = 0 : index} : tensor<3x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_10:.*]] = sparse_tensor.pointers %[[TMP_arg1]] {dimension = 1 : index} : tensor<3x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_11:.*]] = sparse_tensor.indices %[[TMP_arg1]] {dimension = 1 : index} : tensor<3x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_12:.*]] = sparse_tensor.values %[[TMP_arg1]] : tensor<3x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_13:.*]] = memref.load %[[TMP_8]][%[[TMP_c0]]] : memref<?xindex>
 //       CHECK:  %[[TMP_14:.*]] = memref.load %[[TMP_8]][%[[TMP_c1]]] : memref<?xindex>
 //       CHECK:  scf.for %[[TMP_arg3:.*]] = %[[TMP_13]] to %[[TMP_14]] step %[[TMP_c1]] {
 //       CHECK:    %[[TMP_23:.*]] = memref.load %[[TMP_9]][%[[TMP_arg3]]] : memref<?xindex>
-//       CHECK:    %[[TMP_24:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
 //       CHECK:    %[[TMP_25:.*]] = memref.load %[[TMP_10]][%[[TMP_arg3]]] : memref<?xindex>
+//       CHECK:    %[[TMP_24:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
 //       CHECK:    %[[TMP_26:.*]] = memref.load %[[TMP_10]][%[[TMP_24]]] : memref<?xindex>
 //       CHECK:    scf.for %[[TMP_arg4:.*]] = %[[TMP_25]] to %[[TMP_26]] step %[[TMP_c1]] {
 //       CHECK:      %[[TMP_27:.*]] = memref.load %[[TMP_11]][%[[TMP_arg4]]] : memref<?xindex>
 //       CHECK:      %[[TMP_28:.*]] = memref.load %[[TMP_12]][%[[TMP_arg4]]] : memref<?xf64>
 //       CHECK:      %[[TMP_29:.*]] = arith.addi %[[TMP_23]], %[[TMP_c2]] : index
 //       CHECK:      sparse_tensor.insert %[[TMP_28]] into %[[TMP_0]][%[[TMP_29]], %[[TMP_27]]] : tensor<9x4xf64, #sparse_tensor
 //       CHECK:    }
 //       CHECK:  }
 //       CHECK:  %[[TMP_15:.*]] = sparse_tensor.pointers %[[TMP_arg2]] {dimension = 0 : index} : tensor<4x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_16:.*]] = sparse_tensor.indices %[[TMP_arg2]] {dimension = 0 : index} : tensor<4x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_17:.*]] = sparse_tensor.pointers %[[TMP_arg2]] {dimension = 1 : index} : tensor<4x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_18:.*]] = sparse_tensor.indices %[[TMP_arg2]] {dimension = 1 : index} : tensor<4x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_19:.*]] = sparse_tensor.values %[[TMP_arg2]] : tensor<4x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_20:.*]] = memref.load %[[TMP_15]][%[[TMP_c0]]] : memref<?xindex>
 //       CHECK:  %[[TMP_21:.*]] = memref.load %[[TMP_15]][%[[TMP_c1]]] : memref<?xindex>
 //       CHECK:  scf.for %[[TMP_arg3:.*]] = %[[TMP_20]] to %[[TMP_21]] step %[[TMP_c1]] {
 //       CHECK:    %[[TMP_23:.*]] = memref.load %[[TMP_16]][%[[TMP_arg3]]] : memref<?xindex>
-//       CHECK:    %[[TMP_24:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
 //       CHECK:    %[[TMP_25:.*]] = memref.load %[[TMP_17]][%[[TMP_arg3]]] : memref<?xindex>
+//       CHECK:    %[[TMP_24:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
 //       CHECK:    %[[TMP_26:.*]] = memref.load %[[TMP_17]][%[[TMP_24]]] : memref<?xindex>
 //       CHECK:    scf.for %[[TMP_arg4:.*]] = %[[TMP_25]] to %[[TMP_26]] step %[[TMP_c1]] {
 //       CHECK:      %[[TMP_27:.*]] = memref.load %[[TMP_18]][%[[TMP_arg4]]] : memref<?xindex>
 //       CHECK:      %[[TMP_28:.*]] = memref.load %[[TMP_19]][%[[TMP_arg4]]] : memref<?xf64>
 //       CHECK:      %[[TMP_29:.*]] = arith.addi %[[TMP_23]], %[[TMP_c5]] : index
 //       CHECK:      sparse_tensor.insert %[[TMP_28]] into %[[TMP_0]][%[[TMP_29]], %[[TMP_27]]] : tensor<9x4xf64, #sparse_tensor
 //       CHECK:    }
 //       CHECK:  }

mlir/test/Dialect/SparseTensor/sparse_index.mlir

 // CHECK-LABEL:   func.func @dense_index(
 // CHECK-SAME:      %[[VAL_0:.*]]: tensor<?x?xi64, #sparse_tensor.encoding
 // CHECK-DAG:       %[[VAL_1:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_3:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor<?x?xi64, #sparse_tensor.encoding
 // CHECK-DAG:       %[[VAL_4:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor<?x?xi64, #sparse_tensor.encoding
 // CHECK-DAG:       %[[VAL_5:.*]] = bufferization.alloc_tensor(%[[VAL_3]], %[[VAL_4]]) : tensor<?x?xi64, #sparse_tensor.encoding
 // CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?x?xi64, #sparse_tensor.encoding
-// CHECK-DAG:       %[[VAL_7:.*]] = tensor.dim %[[VAL_5]], %[[VAL_1]] : tensor<?x?xi64, #sparse_tensor.encoding
-// CHECK-DAG:       %[[VAL_8:.*]] = tensor.dim %[[VAL_5]], %[[VAL_2]] : tensor<?x?xi64, #sparse_tensor.encoding
+// CHECK-DAG:       %[[VAL_7:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor<?x?xi64, #sparse_tensor.encoding
+// CHECK-DAG:       %[[VAL_8:.*]] = tensor.dim %[[VAL_0]], %[[VAL_2]] : tensor<?x?xi64, #sparse_tensor.encoding
+// CHECK-DAG:       %[[VAL_24:.*]] = tensor.dim %[[VAL_5]], %[[VAL_2]] : tensor<?x?xi64, #sparse_tensor.encoding
 // CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_5]] : tensor<?x?xi64, #sparse_tensor.encoding
 // CHECK:           scf.for %[[VAL_10:.*]] = %[[VAL_1]] to %[[VAL_7]] step %[[VAL_2]] {
 // CHECK:             scf.for %[[VAL_11:.*]] = %[[VAL_1]] to %[[VAL_8]] step %[[VAL_2]] {
 // CHECK:               %[[VAL_12:.*]] = arith.muli %[[VAL_8]], %[[VAL_10]] : index
 // CHECK:               %[[VAL_13:.*]] = arith.addi %[[VAL_12]], %[[VAL_11]] : index
-// CHECK:               %[[VAL_14:.*]] = arith.muli %[[VAL_8]], %[[VAL_10]] : index
+// CHECK:               %[[VAL_14:.*]] = arith.muli %[[VAL_24]], %[[VAL_10]] : index
 // CHECK:               %[[VAL_15:.*]] = arith.addi %[[VAL_14]], %[[VAL_11]] : index
 // CHECK:               %[[VAL_16:.*]] = arith.index_cast %[[VAL_11]] : index to i64
 // CHECK:               %[[VAL_17:.*]] = arith.index_cast %[[VAL_10]] : index to i64
 // CHECK:               %[[VAL_18:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_13]]] : memref<?xi64>
 // CHECK:               %[[VAL_19:.*]] = arith.muli %[[VAL_17]], %[[VAL_18]] : i64
 // CHECK:               %[[VAL_20:.*]] = arith.muli %[[VAL_16]], %[[VAL_19]] : i64
 // CHECK:               memref.store %[[VAL_20]], %[[VAL_9]]{{\[}}%[[VAL_15]]] : memref<?xi64>
 // CHECK:             }

mlir/test/Dialect/SparseTensor/sparse_lower_col.mlir

 // CHECK-HIR-DAG:       %[[VAL_4:.*]] = arith.constant 0 : index
 // CHECK-HIR-DAG:       %[[VAL_5:.*]] = arith.constant 1 : index
 // CHECK-HIR-DAG:       %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)> }>> to memref<?xindex>
 // CHECK-HIR-DAG:       %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)> }>> to memref<?xindex>
 // CHECK-HIR-DAG:       %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)> }>> to memref<?xf64>
 // CHECK-HIR-DAG:       %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<64xf64>
 // CHECK-HIR-DAG:       %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32xf64>
 // CHECK-HIR:           scf.for %[[VAL_12:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
-// CHECK-HIR:             %[[VAL_13:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_12]]] : memref<64xf64>
-// CHECK-HIR:             %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex>
-// CHECK-HIR:             %[[VAL_15:.*]] = arith.addi %[[VAL_12]], %[[VAL_5]] : index
-// CHECK-HIR:             %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_15]]] : memref<?xindex>
+// CHECK-HIR-DAG:         %[[VAL_13:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_12]]] : memref<64xf64>
+// CHECK-HIR-DAG:         %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex>
+// CHECK-HIR-DAG:         %[[VAL_15:.*]] = arith.addi %[[VAL_12]], %[[VAL_5]] : index
+// CHECK-HIR-DAG:         %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_15]]] : memref<?xindex>
 // CHECK-HIR:             scf.for %[[VAL_17:.*]] = %[[VAL_14]] to %[[VAL_16]] step %[[VAL_5]] {
-// CHECK-HIR:               %[[VAL_18:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_17]]] : memref<?xindex>
-// CHECK-HIR:               %[[VAL_19:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_18]]] : memref<32xf64>
-// CHECK-HIR:               %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_17]]] : memref<?xf64>
+// CHECK-HIR-DAG:           %[[VAL_18:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_17]]] : memref<?xindex>
+// CHECK-HIR-DAG:           %[[VAL_19:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_18]]] : memref<32xf64>
+// CHECK-HIR-DAG:           %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_17]]] : memref<?xf64>
 // CHECK-HIR:               %[[VAL_21:.*]] = arith.mulf %[[VAL_20]], %[[VAL_13]] : f64
 // CHECK-HIR:               %[[VAL_22:.*]] = arith.addf %[[VAL_19]], %[[VAL_21]] : f64
 // CHECK-HIR:               memref.store %[[VAL_22]], %[[VAL_11]]{{\[}}%[[VAL_18]]] : memref<32xf64>
 // CHECK-HIR:             }
 // CHECK-HIR:           }
 // CHECK-HIR:           %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xf64>
 // CHECK-HIR:           return %[[VAL_23]] : tensor<32xf64>
 // CHECK-HIR:         }

mlir/test/Dialect/SparseTensor/sparse_parallel.mlir

 // RUN: mlir-opt %s -sparsification="parallelization-strategy=none" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR0
-// RUN: mlir-opt %s -sparsification="parallelization-strategy=dense-outer-loop" | \
-// RUN:   FileCheck %s --check-prefix=CHECK-PAR1
-// RUN: mlir-opt %s -sparsification="parallelization-strategy=any-storage-outer-loop" | \
-// RUN:   FileCheck %s --check-prefix=CHECK-PAR2
-// RUN: mlir-opt %s -sparsification="parallelization-strategy=dense-any-loop" | \
-// RUN:   FileCheck %s --check-prefix=CHECK-PAR3
-// RUN: mlir-opt %s -sparsification="parallelization-strategy=any-storage-any-loop" | \
-// RUN:   FileCheck %s --check-prefix=CHECK-PAR4
+// FIXME: we do not support vectorization/parallel loops in loop emitter right now
+// R_U_N: mlir-opt %s -sparsification="parallelization-strategy=dense-outer-loop" | \
+// R_U_N:   FileCheck %s --check-prefix=CHECK-PAR1
+// R_U_N: mlir-opt %s -sparsification="parallelization-strategy=any-storage-outer-loop" | \
+// R_U_N:   FileCheck %s --check-prefix=CHECK-PAR2
+// R_U_N: mlir-opt %s -sparsification="parallelization-strategy=dense-any-loop" | \
+// R_U_N:   FileCheck %s --check-prefix=CHECK-PAR3
+// R_U_N: mlir-opt %s -sparsification="parallelization-strategy=any-storage-any-loop" | \
+// R_U_N:   FileCheck %s --check-prefix=CHECK-PAR4
 
 #DenseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "dense" ]
 }>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ]
 }>

mlir/test/Dialect/SparseTensor/sparse_perm.mlir

 // CHECK-LABEL:   func @sparse_dynamic_dims(
 // CHECK-SAME:                          %[[VAL_0:.*]]: tensor<?x?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>,
 // CHECK-SAME:                          %[[VAL_1:.*]]: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
 // CHECK-DAG:       %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
 // CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 2 : index
 // CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?x?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
-// CHECK-DAG:       %[[VAL_6:.*]] = tensor.dim %[[VAL_1]], %[[VAL_3]] : tensor<?x?x?xf32>
-// CHECK-DAG:       %[[VAL_7:.*]] = tensor.dim %[[VAL_1]], %[[VAL_4]] : tensor<?x?x?xf32>
-// CHECK-DAG:       %[[VAL_8:.*]] = tensor.dim %[[VAL_1]], %[[VAL_2]] : tensor<?x?x?xf32>
+// CHECK-DAG:       %[[VAL_6:.*]] = tensor.dim %[[VAL_0]], %[[VAL_2]] : tensor<?x?x?xf32
+// CHECK-DAG:       %[[VAL_7:.*]] = tensor.dim %[[VAL_0]], %[[VAL_3]] : tensor<?x?x?xf32
+// CHECK-DAG:       %[[VAL_8:.*]] = tensor.dim %[[VAL_0]], %[[VAL_4]] : tensor<?x?x?xf32
 // CHECK-DAG:       %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<?x?x?xf32>
 // CHECK:           linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_10]] : memref<?x?x?xf32>)
-// CHECK:           scf.for %[[VAL_11:.*]] = %[[VAL_3]] to %[[VAL_7]] step %[[VAL_4]] {
-// CHECK:             scf.for %[[VAL_12:.*]] = %[[VAL_3]] to %[[VAL_8]] step %[[VAL_4]] {
-// CHECK:               %[[VAL_13:.*]] = arith.muli %[[VAL_8]], %[[VAL_11]] : index
+// CHECK:           scf.for %[[VAL_11:.*]] = %[[VAL_3]] to %[[VAL_6]] step %[[VAL_4]] {
+// CHECK:             scf.for %[[VAL_12:.*]] = %[[VAL_3]] to %[[VAL_7]] step %[[VAL_4]] {
+// CHECK:               %[[VAL_13:.*]] = arith.muli %[[VAL_7]], %[[VAL_11]] : index
 // CHECK:               %[[VAL_14:.*]] = arith.addi %[[VAL_13]], %[[VAL_12]] : index
-// CHECK:               scf.for %[[VAL_15:.*]] = %[[VAL_3]] to %[[VAL_6]] step %[[VAL_4]] {
-// CHECK:                 %[[VAL_16:.*]] = arith.muli %[[VAL_6]], %[[VAL_14]] : index
+// CHECK:               scf.for %[[VAL_15:.*]] = %[[VAL_3]] to %[[VAL_8]] step %[[VAL_4]] {
+// CHECK:                 %[[VAL_16:.*]] = arith.muli %[[VAL_8]], %[[VAL_14]] : index
 // CHECK:                 %[[VAL_17:.*]] = arith.addi %[[VAL_16]], %[[VAL_15]] : index
 // CHECK:                 %[[VAL_18:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_17]]] : memref<?xf32>
 // CHECK:                 memref.store %[[VAL_18]], %[[VAL_10]]{{\[}}%[[VAL_15]], %[[VAL_11]], %[[VAL_12]]] : memref<?x?x?xf32>
 // CHECK:               }
 // CHECK:             }
 // CHECK:           }
 // CHECK:           %[[VAL_19:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<?x?x?xf32>
 // CHECK:           return %[[VAL_19]] : tensor<?x?x?xf32>

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

 // CHECK-RWT:         %[[I0:.*]] = sparse_tensor.indices %[[S]] {dimension = 0 : index}
 // CHECK-RWT:          %[[P1:.*]] = sparse_tensor.pointers %[[S]] {dimension = 1 : index}
 // CHECK-RWT:          %[[I1:.*]] = sparse_tensor.indices %[[S]] {dimension = 1 : index}
 // CHECK-RWT:          %[[V:.*]] = sparse_tensor.values %[[S]]
 // CHECK-RWT:          %[[S0:.*]] = memref.load %[[P0]]{{\[}}%[[C0]]] : memref<?xindex>
 // CHECK-RWT:          %[[E0:.*]] = memref.load %[[P0]]{{\[}}%[[C1]]] : memref<?xindex>
 // CHECK-RWT:          scf.for %[[I:.*]] = %[[S0]] to %[[E0]] step %[[C1]] {
 // CHECK-RWT:            %[[SI0:.*]] = memref.load %[[I0]]{{\[}}%[[I]]] : memref<?xindex>
-// CHECK-RWT:            %[[PE1:.*]] = arith.addi %[[I]], %[[C1]] : index
 // CHECK-RWT:            %[[S1:.*]] = memref.load %[[P1]]{{\[}}%[[I]]] : memref<?xindex>
+// CHECK-RWT:            %[[PE1:.*]] = arith.addi %[[I]], %[[C1]] : index
 // CHECK-RWT:            %[[E1:.*]] = memref.load %[[P1]]{{\[}}%[[PE1]]] : memref<?xindex>
 // CHECK-RWT:            scf.for %[[J:.*]] = %[[S1]] to %[[E1]] step %[[C1]] {
 // CHECK-RWT:              %[[SI1:.*]] = memref.load %[[I1]]{{\[}}%[[J]]] : memref<?xindex>
 // CHECK-RWT:              %[[SV:.*]] = memref.load %[[V]]{{\[}}%[[J]]] : memref<?xf64>
 // CHECK-RWT:              %[[T:.*]] = arith.muli %[[SI0]], %[[C10]] : index
 // CHECK-RWT:              %[[DI:.*]] = arith.addi %[[T]], %[[SI1]] : index
 // CHECK-RWT:              sparse_tensor.insert %[[SV]] into %[[B]]{{\[}}%[[DI]]]
 // CHECK-RWT             }

 // CHECK-RWT:         %[[I0:.*]] = sparse_tensor.indices %[[S]] {dimension = 0 : index}
 // CHECK-RWT:          %[[P1:.*]] = sparse_tensor.pointers %[[S]] {dimension = 1 : index}
 // CHECK-RWT:          %[[I1:.*]] = sparse_tensor.indices %[[S]] {dimension = 1 : index}
 // CHECK-RWT:          %[[V:.*]] = sparse_tensor.values %[[S]]
 // CHECK-RWT:          %[[S0:.*]] = memref.load %[[P0]]{{\[}}%[[C0]]] : memref<?xindex>
 // CHECK-RWT:          %[[E0:.*]] = memref.load %[[P0]]{{\[}}%[[C1]]] : memref<?xindex>
 // CHECK-RWT:          scf.for %[[I:.*]] = %[[S0]] to %[[E0]] step %[[C1]] {
 // CHECK-RWT:            %[[SI0:.*]] = memref.load %[[I0]]{{\[}}%[[I]]] : memref<?xindex>
-// CHECK-RWT:            %[[PE1:.*]] = arith.addi %[[I]], %[[C1]] : index
 // CHECK-RWT:            %[[S1:.*]] = memref.load %[[P1]]{{\[}}%[[I]]] : memref<?xindex>
+// CHECK-RWT:            %[[PE1:.*]] = arith.addi %[[I]], %[[C1]] : index
 // CHECK-RWT:            %[[E1:.*]] = memref.load %[[P1]]{{\[}}%[[PE1]]] : memref<?xindex>
 // CHECK-RWT:            scf.for %[[J:.*]] = %[[S1]] to %[[E1]] step %[[C1]] {
 // CHECK-RWT:              %[[SI1:.*]] = memref.load %[[I1]]{{\[}}%[[J]]] : memref<?xindex>
 // CHECK-RWT:              %[[SV:.*]] = memref.load %[[V]]{{\[}}%[[J]]] : memref<?xf64>
 // CHECK-RWT:               %[[T1:.*]] = arith.divui %[[DD0]], %[[C10]] : index
 // CHECK-RWT:               %[[T2:.*]] = arith.muli %[[SI0]], %[[T1]] : index
 // CHECK-RWT:               %[[T3:.*]] = arith.divui %[[T1]], %[[SD1]] : index
 // CHECK-RWT:               %[[T4:.*]] = arith.muli %[[SI1]], %[[T3]] : index

mlir/test/Dialect/SparseTensor/sparse_scalars.mlir

 // CHECK-SAME:              %[[VAL_0:.*0]]: tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>>,
 // CHECK-SAME:              %[[VAL_1:.*1]]: tensor<f32>,
 // CHECK-SAME:              %[[VAL_2:.*2]]: f32,
 // CHECK-SAME:              %[[VAL_3:.*3]]: f32,
 // CHECK-SAME:              %[[VAL_4:.*4]]: tensor<32x16xf32>) -> tensor<32x16xf32> {
 // CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 2.200000e+00 : f32
 // CHECK-DAG:       %[[VAL_6:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_7:.*]] = arith.constant 1 : index
-// CHECK:           %[[VAL_8:.*]] = arith.addf %[[VAL_2]], %[[VAL_3]] : f32
-// CHECK:           %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
-// CHECK:           %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
-// CHECK:           %[[VAL_11:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
-// CHECK:           %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
-// CHECK:           %[[VAL_13:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xf32>
-// CHECK:           %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_1]] : memref<f32>
-// CHECK:           %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_4]] : memref<32x16xf32>
-// CHECK:           %[[VAL_16:.*]] = memref.load %[[VAL_14]][] : memref<f32>
-// CHECK:           %[[VAL_17:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_6]]] : memref<?xindex>
-// CHECK:           %[[VAL_18:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_7]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_8:.*]] = arith.addf %[[VAL_2]], %[[VAL_3]] : f32
+// CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
+// CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
+// CHECK-DAG:       %[[VAL_11:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
+// CHECK-DAG:       %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
+// CHECK-DAG:       %[[VAL_13:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xf32>
+// CHECK-DAG:       %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_1]] : memref<f32>
+// CHECK-DAG:       %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_4]] : memref<32x16xf32>
+// CHECK-DAG:       %[[VAL_16:.*]] = memref.load %[[VAL_14]][] : memref<f32>
+// CHECK-DAG:       %[[VAL_17:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_6]]] : memref<?xindex>
+// CHECK-DAG:       %[[VAL_18:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_7]]] : memref<?xindex>
 // CHECK:           scf.for %[[VAL_19:.*]] = %[[VAL_17]] to %[[VAL_18]] step %[[VAL_7]] {
 // CHECK:             %[[VAL_20:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_19]]] : memref<?xindex>
 // CHECK:             %[[VAL_21:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_19]]] : memref<?xindex>
 // CHECK:             %[[VAL_22:.*]] = arith.addi %[[VAL_19]], %[[VAL_7]] : index
 // CHECK:             %[[VAL_23:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_22]]] : memref<?xindex>
 // CHECK:             scf.for %[[VAL_24:.*]] = %[[VAL_21]] to %[[VAL_23]] step %[[VAL_7]] {
 // CHECK:               %[[VAL_25:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_24]]] : memref<?xindex>
 // CHECK:               %[[VAL_26:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_24]]] : memref<?xf32>

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

+// FIXME: support vector in sparse loop emitter.
 // RUN: mlir-opt %s -sparsification="vectorization-strategy=none vl=16" -cse -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC0
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=dense-inner-loop vl=16" -cse -split-input-file | \
-// RUN:   FileCheck %s --check-prefix=CHECK-VEC1
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=16" -cse -split-input-file | \
-// RUN:   FileCheck %s --check-prefix=CHECK-VEC2
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=16 enable-simd-index32=true" -cse -split-input-file | \
-// RUN:   FileCheck %s --check-prefix=CHECK-VEC3
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=4 enable-vla-vectorization=true" -cse -split-input-file | \
-// RUN:   FileCheck %s --check-prefix=CHECK-VEC4
+// R_U_N: mlir-opt %s -sparsification="vectorization-strategy=dense-inner-loop vl=16" -cse -split-input-file | \
+// R_U_N:   FileCheck %s --check-prefix=CHECK-VEC1
+// R_U_N: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=16" -cse -split-input-file | \
+// R_U_N:   FileCheck %s --check-prefix=CHECK-VEC2
+// R_U_N: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=16 enable-simd-index32=true" -cse -split-input-file | \
+// R_U_N:   FileCheck %s --check-prefix=CHECK-VEC3
+// R_U_N: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=4 enable-vla-vectorization=true" -cse -split-input-file | \
+// R_U_N:   FileCheck %s --check-prefix=CHECK-VEC4
 
 #DenseVector = #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>
 
 #trait_scale_d = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>   // x (out)
   ],

mlir/test/Dialect/SparseTensor/sparse_vector_chain.mlir

-// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
-
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=8" -canonicalize | \
-// RUN:   FileCheck %s
-
-#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["dense","compressed"]}>
-
-#trait = {
-  indexing_maps = [
-    affine_map<(i,j) -> (i,j)>,  // a (in)
-    affine_map<(i,j) -> (i,j)>,  // b (in)
-    affine_map<(i,j) -> ()>      // x (out)
-  ],
-  iterator_types = ["reduction", "reduction"]
-}
-
-// Verifies that the SIMD reductions in the two for-loops after the
-// while-loop are chained before horizontally reducing these back to scalar.
-//
-// CHECK-LABEL:   func @sparse_matrix_sum(
-// CHECK-SAME:      %[[VAL_0:.*]]: tensor<f64>,
-// CHECK-SAME:      %[[VAL_1:.*]]: tensor<64x32xf64, #sparse_tensor.encoding<{{{.*}}}>>,
-// CHECK-SAME:      %[[VAL_2:.*]]: tensor<64x32xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<f64> {
-// CHECK-DAG:       %[[VAL_3:.*]] = arith.constant dense<0.000000e+00> : vector<8xf64>
-// CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 8 : index
-// CHECK-DAG:       %[[VAL_6:.*]] = arith.constant 0 : index
-// CHECK-DAG:       %[[VAL_7:.*]] = arith.constant 64 : index
-// CHECK-DAG:       %[[VAL_8:.*]] = arith.constant 1 : index
-// CHECK:           %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 1 : index} : tensor<64x32xf64, #sparse_tensor.encoding<{{{.*}}}>>
-// CHECK:           %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<64x32xf64, #sparse_tensor.encoding<{{{.*}}}>>
-// CHECK:           %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<64x32xf64, #sparse_tensor.encoding<{{{.*}}}>>
-// CHECK:           %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_2]] {dimension = 1 : index} : tensor<64x32xf64, #sparse_tensor.encoding<{{{.*}}}>>
-// CHECK:           %[[VAL_13:.*]] = sparse_tensor.indices %[[VAL_2]] {dimension = 1 : index} : tensor<64x32xf64, #sparse_tensor.encoding<{{{.*}}}>>
-// CHECK:           %[[VAL_14:.*]] = sparse_tensor.values %[[VAL_2]] : tensor<64x32xf64, #sparse_tensor.encoding<{{{.*}}}>>
-// CHECK:           %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_0]] : memref<f64>
-// CHECK:           %[[VAL_16:.*]] = tensor.extract %[[VAL_0]][] : tensor<f64>
-// CHECK:           %[[VAL_17:.*]] = scf.for %[[VAL_18:.*]] = %[[VAL_6]] to %[[VAL_7]] step %[[VAL_8]] iter_args(%[[VAL_19:.*]] = %[[VAL_16]]) -> (f64) {
-// CHECK:             %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_18]]] : memref<?xindex>
-// CHECK:             %[[VAL_21:.*]] = arith.addi %[[VAL_18]], %[[VAL_8]] : index
-// CHECK:             %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_21]]] : memref<?xindex>
-// CHECK:             %[[VAL_23:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_18]]] : memref<?xindex>
-// CHECK:             %[[VAL_24:.*]] = arith.addi %[[VAL_18]], %[[VAL_8]] : index
-// CHECK:             %[[VAL_25:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_24]]] : memref<?xindex>
-// CHECK:             %[[VAL_26:.*]]:3 = scf.while (%[[VAL_27:.*]] = %[[VAL_20]], %[[VAL_28:.*]] = %[[VAL_23]], %[[VAL_29:.*]] = %[[VAL_19]]) : (index, index, f64) -> (index, index, f64) {
-// CHECK:               %[[VAL_30:.*]] = arith.cmpi ult, %[[VAL_27]], %[[VAL_22]] : index
-// CHECK:               %[[VAL_31:.*]] = arith.cmpi ult, %[[VAL_28]], %[[VAL_25]] : index
-// CHECK:               %[[VAL_32:.*]] = arith.andi %[[VAL_30]], %[[VAL_31]] : i1
-// CHECK:               scf.condition(%[[VAL_32]]) %[[VAL_27]], %[[VAL_28]], %[[VAL_29]] : index, index, f64
-// CHECK:             } do {
-// CHECK:             ^bb0(%[[VAL_33:.*]]: index, %[[VAL_34:.*]]: index, %[[VAL_35:.*]]: f64):
-// CHECK:               %[[VAL_36:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_33]]] : memref<?xindex>
-// CHECK:               %[[VAL_37:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_34]]] : memref<?xindex>
-// CHECK:               %[[VAL_38:.*]] = arith.cmpi ult, %[[VAL_37]], %[[VAL_36]] : index
-// CHECK:               %[[VAL_39:.*]] = arith.select %[[VAL_38]], %[[VAL_37]], %[[VAL_36]] : index
-// CHECK:               %[[VAL_40:.*]] = arith.cmpi eq, %[[VAL_36]], %[[VAL_39]] : index
-// CHECK:               %[[VAL_41:.*]] = arith.cmpi eq, %[[VAL_37]], %[[VAL_39]] : index
-// CHECK:               %[[VAL_42:.*]] = arith.andi %[[VAL_40]], %[[VAL_41]] : i1
-// CHECK:               %[[VAL_43:.*]] = scf.if %[[VAL_42]] -> (f64) {
-// CHECK:                 %[[VAL_44:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_33]]] : memref<?xf64>
-// CHECK:                 %[[VAL_45:.*]] = memref.load %[[VAL_14]]{{\[}}%[[VAL_34]]] : memref<?xf64>
-// CHECK:                 %[[VAL_46:.*]] = arith.addf %[[VAL_44]], %[[VAL_45]] : f64
-// CHECK:                 %[[VAL_47:.*]] = arith.addf %[[VAL_35]], %[[VAL_46]] : f64
-// CHECK:                 scf.yield %[[VAL_47]] : f64
-// CHECK:               } else {
-// CHECK:                 %[[VAL_48:.*]] = arith.cmpi eq, %[[VAL_36]], %[[VAL_39]] : index
-// CHECK:                 %[[VAL_49:.*]] = scf.if %[[VAL_48]] -> (f64) {
-// CHECK:                   %[[VAL_50:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_33]]] : memref<?xf64>
-// CHECK:                   %[[VAL_51:.*]] = arith.addf %[[VAL_35]], %[[VAL_50]] : f64
-// CHECK:                   scf.yield %[[VAL_51]] : f64
-// CHECK:                 } else {
-// CHECK:                   %[[VAL_52:.*]] = arith.cmpi eq, %[[VAL_37]], %[[VAL_39]] : index
-// CHECK:                   %[[VAL_53:.*]] = scf.if %[[VAL_52]] -> (f64) {
-// CHECK:                     %[[VAL_54:.*]] = memref.load %[[VAL_14]]{{\[}}%[[VAL_34]]] : memref<?xf64>
-// CHECK:                     %[[VAL_55:.*]] = arith.addf %[[VAL_35]], %[[VAL_54]] : f64
-// CHECK:                     scf.yield %[[VAL_55]] : f64
-// CHECK:                   } else {
-// CHECK:                     scf.yield %[[VAL_35]] : f64
-// CHECK:                   }
-// CHECK:                   scf.yield %[[VAL_56:.*]] : f64
-// CHECK:                 }
-// CHECK:                 scf.yield %[[VAL_57:.*]] : f64
-// CHECK:               }
-// CHECK:               %[[VAL_58:.*]] = arith.cmpi eq, %[[VAL_36]], %[[VAL_39]] : index
-// CHECK:               %[[VAL_59:.*]] = arith.addi %[[VAL_33]], %[[VAL_8]] : index
-// CHECK:               %[[VAL_60:.*]] = arith.select %[[VAL_58]], %[[VAL_59]], %[[VAL_33]] : index
-// CHECK:               %[[VAL_61:.*]] = arith.cmpi eq, %[[VAL_37]], %[[VAL_39]] : index
-// CHECK:               %[[VAL_62:.*]] = arith.addi %[[VAL_34]], %[[VAL_8]] : index
-// CHECK:               %[[VAL_63:.*]] = arith.select %[[VAL_61]], %[[VAL_62]], %[[VAL_34]] : index
-// CHECK:               scf.yield %[[VAL_60]], %[[VAL_63]], %[[VAL_64:.*]] : index, index, f64
-// CHECK:             }
-// CHECK:             %[[VAL_65:.*]] = vector.insertelement %[[VAL_66:.*]]#2, %[[VAL_3]]{{\[}}%[[VAL_6]] : index] : vector<8xf64>
-// CHECK:             %[[VAL_67:.*]] = scf.for %[[VAL_68:.*]] = %[[VAL_66]]#0 to %[[VAL_22]] step %[[VAL_4]] iter_args(%[[VAL_69:.*]] = %[[VAL_65]]) -> (vector<8xf64>) {
-// CHECK:               %[[VAL_70:.*]] = affine.min #map(%[[VAL_22]], %[[VAL_68]])
-// CHECK:               %[[VAL_71:.*]] = vector.create_mask %[[VAL_70]] : vector<8xi1>
-// CHECK:               %[[VAL_72:.*]] = vector.maskedload %[[VAL_11]]{{\[}}%[[VAL_68]]], %[[VAL_71]], %[[VAL_3]] : memref<?xf64>, vector<8xi1>, vector<8xf64> into vector<8xf64>
-// CHECK:               %[[VAL_73:.*]] = arith.addf %[[VAL_69]], %[[VAL_72]] : vector<8xf64>
-// CHECK:               %[[VAL_74:.*]] = arith.select %[[VAL_71]], %[[VAL_73]], %[[VAL_69]] : vector<8xi1>, vector<8xf64>
-// CHECK:               scf.yield %[[VAL_74]] : vector<8xf64>
-// CHECK:             }
-// CHECK:             %[[VAL_75:.*]] = scf.for %[[VAL_76:.*]] = %[[VAL_66]]#1 to %[[VAL_25]] step %[[VAL_4]] iter_args(%[[VAL_77:.*]] = %[[VAL_78:.*]]) -> (vector<8xf64>) {
-// CHECK:               %[[VAL_79:.*]] = affine.min #map(%[[VAL_25]], %[[VAL_76]])
-// CHECK:               %[[VAL_80:.*]] = vector.create_mask %[[VAL_79]] : vector<8xi1>
-// CHECK:               %[[VAL_81:.*]] = vector.maskedload %[[VAL_14]]{{\[}}%[[VAL_76]]], %[[VAL_80]], %[[VAL_3]] : memref<?xf64>, vector<8xi1>, vector<8xf64> into vector<8xf64>
-// CHECK:               %[[VAL_82:.*]] = arith.addf %[[VAL_77]], %[[VAL_81]] : vector<8xf64>
-// CHECK:               %[[VAL_83:.*]] = arith.select %[[VAL_80]], %[[VAL_82]], %[[VAL_77]] : vector<8xi1>, vector<8xf64>
-// CHECK:               scf.yield %[[VAL_83]] : vector<8xf64>
-// CHECK:             }
-// CHECK:             %[[VAL_84:.*]] = vector.reduction <add>, %[[VAL_85:.*]] : vector<8xf64> into f64
-// CHECK:             scf.yield %[[VAL_84]] : f64
-// CHECK:           }
-// CHECK:           memref.store %[[VAL_86:.*]], %[[VAL_15]][] : memref<f64>
-// CHECK:           %[[VAL_87:.*]] = bufferization.to_tensor %[[VAL_15]] : memref<f64>
-// CHECK:           return %[[VAL_87]] : tensor<f64>
-// CHECK:         }
-func.func @sparse_matrix_sum(%argx: tensor<f64>,
-                             %arga: tensor<64x32xf64, #SparseMatrix>,
-                             %argb: tensor<64x32xf64, #SparseMatrix>) -> tensor<f64> {
-  %0 = linalg.generic #trait
-     ins(%arga, %argb: tensor<64x32xf64, #SparseMatrix>,
-                       tensor<64x32xf64, #SparseMatrix>)
-      outs(%argx: tensor<f64>) {
-      ^bb(%a: f64, %b: f64, %x: f64):
-        %m = arith.addf %a, %b : f64
-        %t = arith.addf %x, %m : f64
-        linalg.yield %t : f64
-  } -> tensor<f64>
-  return %0 : tensor<f64>
-}

mlir/test/Dialect/SparseTensor/sparse_vector_index.mlir

-// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
-
-// The script is designed to make adding checks to
-// a test case fast, it is *not* designed to be authoritative
-// about what constitutes a good test! The CHECK should be
-// minimized and named to reflect the test intent.
-
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=8" -canonicalize | \
-// RUN:   FileCheck %s
-
-#SparseVector = #sparse_tensor.encoding<{
-  dimLevelType = ["compressed"]
-}>
-
-#trait_1d = {
-  indexing_maps = [
-    affine_map<(i) -> (i)>,  // a
-    affine_map<(i) -> (i)>   // x (out)
-  ],
-  iterator_types = ["parallel"],
-  doc = "X(i) = a(i) op i"
-}
-
-// CHECK-LABEL:   func @sparse_index_1d_conj(
-// CHECK-SAME:      %[[VAL_0:.*]]: tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<8xi64> {
-// CHECK-DAG:       %[[VAL_1:.*]] = arith.constant dense<0> : vector<8xi64>
-// CHECK-DAG:       %[[VAL_2:.*]] = arith.constant dense<0> : vector<8xindex>
-// CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 8 : index
-// CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 1 : index
-// CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 0 : i64
-// CHECK-DAG:       %[[VAL_6:.*]] = arith.constant 0 : index
-// CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xi64>
-// CHECK-DAG:       %[[VAL_10a:.*]] = tensor.empty() : tensor<8xi64>
-// CHECK-DAG:       %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_10a]] : memref<8xi64>
-// CHECK-DAG:       linalg.fill ins(%[[VAL_5]] : i64) outs(%[[VAL_10]] : memref<8xi64>)
-// CHECK-DAG:       %[[VAL_11:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
-// CHECK-DAG:       %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           scf.for %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_3]] {
-// CHECK:             %[[VAL_14:.*]] = affine.min #map0(%[[VAL_13]]){{\[}}%[[VAL_12]]]
-// CHECK:             %[[VAL_15:.*]] = vector.create_mask %[[VAL_14]] : vector<8xi1>
-// CHECK:             %[[VAL_16:.*]] = vector.maskedload %[[VAL_8]]{{\[}}%[[VAL_13]]], %[[VAL_15]], %[[VAL_2]] : memref<?xindex>, vector<8xi1>, vector<8xindex> into vector<8xindex>
-// CHECK:             %[[VAL_17:.*]] = vector.maskedload %[[VAL_9]]{{\[}}%[[VAL_13]]], %[[VAL_15]], %[[VAL_1]] : memref<?xi64>, vector<8xi1>, vector<8xi64> into vector<8xi64>
-// CHECK:             %[[VAL_18:.*]] = arith.index_cast %[[VAL_16]] : vector<8xindex> to vector<8xi64>
-// CHECK:             %[[VAL_19:.*]] = arith.muli %[[VAL_17]], %[[VAL_18]] : vector<8xi64>
-// CHECK:             vector.scatter %[[VAL_10]]{{\[}}%[[VAL_6]]] {{\[}}%[[VAL_16]]], %[[VAL_15]], %[[VAL_19]] : memref<8xi64>, vector<8xindex>, vector<8xi1>, vector<8xi64>
-// CHECK:           }
-// CHECK:           %[[VAL_20:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<8xi64>
-// CHECK:           return %[[VAL_20]] : tensor<8xi64>
-// CHECK:         }
-func.func @sparse_index_1d_conj(%arga: tensor<8xi64, #SparseVector>) -> tensor<8xi64> {
-  %init = tensor.empty() : tensor<8xi64>
-  %r = linalg.generic #trait_1d
-      ins(%arga: tensor<8xi64, #SparseVector>)
-     outs(%init: tensor<8xi64>) {
-      ^bb(%a: i64, %x: i64):
-        %i = linalg.index 0 : index
-        %ii = arith.index_cast %i : index to i64
-        %m1 = arith.muli %a, %ii : i64
-        linalg.yield %m1 : i64
-  } -> tensor<8xi64>
-  return %r : tensor<8xi64>
-}
-
-// CHECK-LABEL:   func @sparse_index_1d_disj(
-// CHECK-SAME:      %[[VAL_0:.*]]: tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<8xi64> {
-// CHECK-DAG:       %[[VAL_1:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex>
-// CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 1 : index
-// CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 0 : i64
-// CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 8 : index
-// CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 0 : index
-// CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xi64>
-// CHECK-DAG:       %[[VAL_9a:.*]] = tensor.empty() : tensor<8xi64>
-// CHECK-DAG:       %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_9a]] : memref<8xi64>
-// CHECK-DAG:       linalg.fill ins(%[[VAL_3]] : i64) outs(%[[VAL_9]] : memref<8xi64>)
-// CHECK-DAG:       %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
-// CHECK-DAG:       %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_2]]] : memref<?xindex>
-// CHECK:           %[[VAL_12:.*]]:2 = scf.while (%[[VAL_13:.*]] = %[[VAL_10]], %[[VAL_14:.*]] = %[[VAL_5]]) : (index, index) -> (index, index) {
-// CHECK:             %[[VAL_15:.*]] = arith.cmpi ult, %[[VAL_13]], %[[VAL_11]] : index
-// CHECK:             scf.condition(%[[VAL_15]]) %[[VAL_13]], %[[VAL_14]] : index, index
-// CHECK:           } do {
-// CHECK:           ^bb0(%[[VAL_16:.*]]: index, %[[VAL_17:.*]]: index):
-// CHECK:             %[[VAL_18:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_16]]] : memref<?xindex>
-// CHECK:             %[[VAL_19:.*]] = arith.cmpi eq, %[[VAL_18]], %[[VAL_17]] : index
-// CHECK:             scf.if %[[VAL_19]] {
-// CHECK:               %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_16]]] : memref<?xi64>
-// CHECK:               %[[VAL_21:.*]] = arith.index_cast %[[VAL_17]] : index to i64
-// CHECK:               %[[VAL_22:.*]] = arith.addi %[[VAL_20]], %[[VAL_21]] : i64
-// CHECK:               memref.store %[[VAL_22]], %[[VAL_9]]{{\[}}%[[VAL_17]]] : memref<8xi64>
-// CHECK:             } else {
-// CHECK:               %[[VAL_23:.*]] = arith.index_cast %[[VAL_17]] : index to i64
-// CHECK:               memref.store %[[VAL_23]], %[[VAL_9]]{{\[}}%[[VAL_17]]] : memref<8xi64>
-// CHECK:             }
-// CHECK:             %[[VAL_24:.*]] = arith.cmpi eq, %[[VAL_18]], %[[VAL_17]] : index
-// CHECK:             %[[VAL_25:.*]] = arith.addi %[[VAL_16]], %[[VAL_2]] : index
-// CHECK:             %[[VAL_26:.*]] = arith.select %[[VAL_24]], %[[VAL_25]], %[[VAL_16]] : index
-// CHECK:             %[[VAL_27:.*]] = arith.addi %[[VAL_17]], %[[VAL_2]] : index
-// CHECK:             scf.yield %[[VAL_26]], %[[VAL_27]] : index, index
-// CHECK:           }
-// CHECK:           scf.for %[[VAL_28:.*]] = %[[VAL_29:.*]]#1 to %[[VAL_4]] step %[[VAL_4]] {
-// CHECK:             %[[VAL_30:.*]] = affine.min #map1(%[[VAL_28]])
-// CHECK:             %[[VAL_31:.*]] = vector.create_mask %[[VAL_30]] : vector<8xi1>
-// CHECK:             %[[VAL_32:.*]] = vector.broadcast %[[VAL_28]] : index to vector<8xindex>
-// CHECK:             %[[VAL_33:.*]] = arith.addi %[[VAL_32]], %[[VAL_1]] : vector<8xindex>
-// CHECK:             %[[VAL_34:.*]] = arith.index_cast %[[VAL_33]] : vector<8xindex> to vector<8xi64>
-// CHECK:             vector.maskedstore %[[VAL_9]]{{\[}}%[[VAL_28]]], %[[VAL_31]], %[[VAL_34]] : memref<8xi64>, vector<8xi1>, vector<8xi64>
-// CHECK:           }
-// CHECK:           %[[VAL_35:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<8xi64>
-// CHECK:           return %[[VAL_35]] : tensor<8xi64>
-// CHECK:         }
-func.func @sparse_index_1d_disj(%arga: tensor<8xi64, #SparseVector>) -> tensor<8xi64> {
-  %init = tensor.empty() : tensor<8xi64>
-  %r = linalg.generic #trait_1d
-      ins(%arga: tensor<8xi64, #SparseVector>)
-     outs(%init: tensor<8xi64>) {
-      ^bb(%a: i64, %x: i64):
-        %i = linalg.index 0 : index
-        %ii = arith.index_cast %i : index to i64
-        %m1 = arith.addi %a, %ii : i64
-        linalg.yield %m1 : i64
-  } -> tensor<8xi64>
-  return %r : tensor<8xi64>
-}

mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir

-// RUN: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=16" -scf-for-loop-peeling -canonicalize | \
-// RUN:   FileCheck %s
-
-#SparseVector = #sparse_tensor.encoding<{
-  dimLevelType = [ "compressed" ],
-  pointerBitWidth = 32,
-  indexBitWidth = 32
-}>
-
-#trait_mul_s = {
-  indexing_maps = [
-    affine_map<(i) -> (i)>,  // a
-    affine_map<(i) -> (i)>,  // b
-    affine_map<(i) -> (i)>   // x (out)
-  ],
-  iterator_types = ["parallel"],
-  doc = "x(i) = a(i) * b(i)"
-}
-
-// CHECK-DAG:   #[[$map0:.*]] = affine_map<()[s0, s1] -> (s0 + ((-s0 + s1) floordiv 16) * 16)>
-// CHECK-DAG:   #[[$map1:.*]] = affine_map<(d0)[s0] -> (-d0 + s0)>
-// CHECK-LABEL: func @mul_s
-// CHECK-DAG:   %[[c0:.*]] = arith.constant 0 : index
-// CHECK-DAG:   %[[c1:.*]] = arith.constant 1 : index
-// CHECK-DAG:   %[[c16:.*]] = arith.constant 16 : index
-// CHECK:       %[[p:.*]] = memref.load %{{.*}}[%[[c0]]] : memref<?xi32>
-// CHECK:       %[[a:.*]] = arith.extui %[[p]] : i32 to i64
-// CHECK:       %[[q:.*]] = arith.index_cast %[[a]] : i64 to index
-// CHECK:       %[[r:.*]] = memref.load %{{.*}}[%[[c1]]] : memref<?xi32>
-// CHECK:       %[[b:.*]] = arith.extui %[[r]] : i32 to i64
-// CHECK:       %[[s:.*]] = arith.index_cast %[[b]] : i64 to index
-// CHECK:       %[[boundary:.*]] = affine.apply #[[$map0]]()[%[[q]], %[[s]]]
-// CHECK:       scf.for %[[i:.*]] = %[[q]] to %[[boundary]] step %[[c16]] {
-// CHECK:         %[[mask:.*]] = vector.constant_mask [16] : vector<16xi1>
-// CHECK:         %[[li:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xi32>, vector<16xi32>
-// CHECK:         %[[zi:.*]] = arith.extui %[[li]] : vector<16xi32> to vector<16xi64>
-// CHECK:         %[[la:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xf32>, vector<16xf32>
-// CHECK:         %[[lb:.*]] = vector.gather %{{.*}}[%[[c0]]] [%[[zi]]], %[[mask]], %{{.*}} : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> into vector<16xf32>
-// CHECK:         %[[m:.*]] = arith.mulf %[[la]], %[[lb]] : vector<16xf32>
-// CHECK:         vector.scatter %{{.*}}[%[[c0]]] [%[[zi]]], %[[mask]], %[[m]] : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32>
-// CHECK:       }
-// CHECK:       scf.for %[[i2:.*]] = %[[boundary]] to %[[s]] step %[[c16]] {
-// CHECK:         %[[sub:.*]] = affine.apply #[[$map1]](%[[i2]])[%[[s]]]
-// CHECK:         %[[mask2:.*]] = vector.create_mask %[[sub]] : vector<16xi1>
-// CHECK:         %[[li2:.*]] = vector.maskedload %{{.*}}[%[[i2]]], %[[mask2]], %{{.*}} : memref<?xi32>, vector<16xi1>, vector<16xi32> into vector<16xi32>
-// CHECK:         %[[zi2:.*]] = arith.extui %[[li2]] : vector<16xi32> to vector<16xi64>
-// CHECK:         %[[la2:.*]] = vector.maskedload %{{.*}}[%[[i2]]], %[[mask2]], %{{.*}} : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
-// CHECK:         %[[lb2:.*]] = vector.gather %{{.*}}[%[[c0]]] [%[[zi2]]], %[[mask2]], %{{.*}} : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> into vector<16xf32>
-// CHECK:         %[[m2:.*]] = arith.mulf %[[la2]], %[[lb2]] : vector<16xf32>
-// CHECK:         vector.scatter %{{.*}}[%[[c0]]] [%[[zi2]]], %[[mask2]], %[[m2]] : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32>
-// CHECK:       }
-// CHECK:       return
-//
-func.func @mul_s(%arga: tensor<1024xf32, #SparseVector>, %argb: tensor<1024xf32>, %argx: tensor<1024xf32>) -> tensor<1024xf32> {
-  %0 = linalg.generic #trait_mul_s
-    ins(%arga, %argb: tensor<1024xf32, #SparseVector>, tensor<1024xf32>)
-    outs(%argx: tensor<1024xf32>) {
-      ^bb(%a: f32, %b: f32, %x: f32):
-        %0 = arith.mulf %a, %b : f32
-        linalg.yield %0 : f32
-  } -> tensor<1024xf32>
-  return %0 : tensor<1024xf32>
-}

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
-//
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=2" | \
-// RUN: mlir-cpu-runner \
-// RUN:  -e entry -entry-point-result=void  \
-// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
-// RUN: FileCheck %s
+// FIXME: we do not support vectorization/parallel loops in loop emitter right now
+// R_U_N: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=2" | \
+// R_U_N: mlir-cpu-runner \
+// R_U_N:  -e entry -entry-point-result=void  \
+// R_U_N:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// R_U_N: FileCheck %s
 
 #SV = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
 
 #trait_cast = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // A (in)
     affine_map<(i) -> (i)>   // X (out)
   ],

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: TENSOR0="%mlir_src_dir/test/Integration/data/test.tns" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
-//
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
-// RUN: TENSOR0="%mlir_src_dir/test/Integration/data/test.tns" \
-// RUN: mlir-cpu-runner \
-// RUN:  -e entry -entry-point-result=void  \
-// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
-// RUN: FileCheck %s
+// TODO: support vectorizaion in loop emitter?
+// R_U_N: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
+// R_U_N: TENSOR0="%mlir_src_dir/test/Integration/data/test.tns" \
+// R_U_N: mlir-cpu-runner \
+// R_U_N:  -e entry -entry-point-result=void  \
+// R_U_N:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// R_U_N: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #SparseTensor = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed", "compressed", "compressed",
                    "compressed", "compressed", "compressed", "compressed" ],
   // Note that any dimOrdering permutation should give the same results
   // since, even though it impacts the sparse storage scheme layout,

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index_dense.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
-//
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
-// RUN: mlir-cpu-runner -e entry -entry-point-result=void \
-// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
-// RUN: FileCheck %s
+// FIXME: support vectorization in loop emitter
+// R_U_N: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
+// R_U_N: mlir-cpu-runner -e entry -entry-point-result=void \
+// R_U_N:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// R_U_N: FileCheck %s
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"]
 }>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = ["compressed", "compressed"]
 }>

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
-//
-// RUN: mlir-opt %s \
-// RUN:   --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=16 enable-simd-index32" | \
-// RUN: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
-// RUN: mlir-cpu-runner \
-// RUN:  -e entry -entry-point-result=void  \
-// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
-// RUN: FileCheck %s
+// FIXME: support vector in sparse loop emitter.
+// R_U_N: mlir-opt %s \
+// R_U_N:   --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=16 enable-simd-index32" | \
+// R_U_N: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
+// R_U_N: mlir-cpu-runner \
+// R_U_N:  -e entry -entry-point-result=void  \
+// R_U_N:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// R_U_N: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ],
   pointerBitWidth = 8,
   indexBitWidth = 8
 }>

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: TENSOR0="%mlir_src_dir/test/Integration/data/mttkrp_b.tns" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
-//
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
-// RUN: TENSOR0="%mlir_src_dir/test/Integration/data/mttkrp_b.tns" \
-// RUN: mlir-cpu-runner \
-// RUN:  -e entry -entry-point-result=void  \
-// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
-// RUN: FileCheck %s
+// FIXME: support vector in sparse loop emitter.
+// R_U_N: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
+// R_U_N: TENSOR0="%mlir_src_dir/test/Integration/data/mttkrp_b.tns" \
+// R_U_N: mlir-cpu-runner \
+// R_U_N:  -e entry -entry-point-result=void  \
+// R_U_N:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// R_U_N: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #SparseTensor = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed", "compressed" ]
 }>
 
 #mttkrp = {

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: TENSOR0="%mlir_src_dir/test/Integration/data/test.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
-//
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
-// RUN: TENSOR0="%mlir_src_dir/test/Integration/data/test.mtx" \
-// RUN: mlir-cpu-runner \
-// RUN:  -e entry -entry-point-result=void  \
-// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
-// RUN: FileCheck %s
+// FIXME: support vector in sparse loop emitter.
+// R_U_N: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
+// R_U_N: TENSOR0="%mlir_src_dir/test/Integration/data/test.mtx" \
+// R_U_N: mlir-cpu-runner \
+// R_U_N:  -e entry -entry-point-result=void  \
+// R_U_N:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// R_U_N: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #DCSR = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ],
   dimOrdering = affine_map<(i,j) -> (i,j)>
 }>
 

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
-//
-// RUN: mlir-opt %s -sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=8" | \
-// RUN: mlir-cpu-runner -e entry -entry-point-result=void \
-// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
-// RUN: FileCheck %s
+// FIXME: we do not support vectorization/parallel loops in loop emitter right now
+// R_U_N: mlir-opt %s -sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=8" | \
+// R_U_N: mlir-cpu-runner -e entry -entry-point-result=void \
+// R_U_N:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// R_U_N: FileCheck %s
 
 #SV = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
 #DV = #sparse_tensor.encoding<{ dimLevelType = [ "dense"      ] }>
 
 #trait_reduction = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> ()>    // x (scalar out)

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
-//
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
-// RUN: mlir-cpu-runner \
-// RUN:  -e entry -entry-point-result=void  \
-// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
-// RUN: FileCheck %s
+// FIXME: support vectorization in loop emitter.
+// R_U_N: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
+// R_U_N: mlir-cpu-runner \
+// R_U_N:  -e entry -entry-point-result=void  \
+// R_U_N:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// R_U_N: FileCheck %s
 
 #CSR = #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>
 
 #trait_scale = {
   indexing_maps = [
     affine_map<(i,j) -> (i,j)>   // X (out)
   ],
   iterator_types = ["parallel", "parallel"],

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
-//
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=2" | \
-// RUN: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
-// RUN: mlir-cpu-runner \
-// RUN:  -e entry -entry-point-result=void  \
-// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
-// RUN: FileCheck %s
+// FIXME: we do not support vectorization/parallel loops in loop emitter right now
+// R_U_N: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=2" | \
+// R_U_N: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
+// R_U_N: mlir-cpu-runner \
+// R_U_N:  -e entry -entry-point-result=void  \
+// R_U_N:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// R_U_N: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ]
 }>
 
 #spmm = {

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: TENSOR0="%mlir_src_dir/test/Integration/data/test_symmetric.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
-//
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=2" | \
-// RUN: TENSOR0="%mlir_src_dir/test/Integration/data/test_symmetric.mtx" \
-// RUN: mlir-cpu-runner \
-// RUN:  -e entry -entry-point-result=void  \
-// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
-// RUN: FileCheck %s
+// FIXME: support vector in sparse loop emitter.
+// R_U_N: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=2" | \
+// R_U_N: TENSOR0="%mlir_src_dir/test/Integration/data/test_symmetric.mtx" \
+// R_U_N: mlir-cpu-runner \
+// R_U_N:  -e entry -entry-point-result=void  \
+// R_U_N:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// R_U_N: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ]
 }>
 
 #trait_sum_reduce = {

mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py

with ir.Context() as ctx, ir.Location.unknown():
     levels = [[st.DimLevelType.dense, st.DimLevelType.dense],
               [st.DimLevelType.dense, st.DimLevelType.compressed],
               [st.DimLevelType.compressed, st.DimLevelType.dense],
               [st.DimLevelType.compressed, st.DimLevelType.compressed]]
     orderings = [
         ir.AffineMap.get_permutation([0, 1]),
         ir.AffineMap.get_permutation([1, 0])
     ]
-    vec_strategy = ['none', 'dense-inner-loop']
+    vec_strategy = ['none'] #, 'dense-inner-loop']
     for level in levels:
       for ordering in orderings:
         for pwidth in [32]:
           for iwidth in [32]:
             for vec in vec_strategy:
               for e in [True]:
                 vl = 1 if vec == 0 else 16
                 attr = st.EncodingAttr.get(level, ordering, None, pwidth,
                                            iwidth)
                 opt = (f'parallelization-strategy=none '
                        f'vectorization-strategy={vec} '
                        f'vl={vl} enable-simd-index32={e}')
                 compiler = sparse_compiler.SparseCompiler(
                     options=opt, opt_level=0, shared_libs=[support_lib])
                 build_compile_and_run_SDDMMM(attr, compiler)
                 count = count + 1
-  # CHECK: Passed 16 tests
+  # CHECK: Passed 8 tests
   print('Passed ', count, 'tests')
 
 
 if __name__ == '__main__':
   main()