[mlir][sparse] Pass down constant coefficients of affine index expressions to LoopEmitter.

[aartbik]

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

62

it -> its

63

I would remove "linalg". The affine map concept exists outside the linalg dialect

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

110

use plural: list of loop indices .... list of ... pairs

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

87–92

This was already there, but "affine expression index" reads a bit verbose. Just index in affine expression or so?

87–92

h

87–92

x

343

why is this absolute now?

727–729

Why is this conceptual?

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

328

This TODO does not make much sense?
"only" and "more complex"?

[Peiming]

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

343

It has always been absolute, but it does not really matter for now because we only handle 2-variable expressions, the absolute offset equals to the relative offset for the first slice (since it parent is the sparse tensor itself). I changed it just for future reference here.

727–729

Maybe I should rephrase it to "The (size, stride) for each conceptual slice"?

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

328

Yeah, this is wired, but we can handle 2 * d0 + d1 but not 2 * d0. More specifically, we can handle coefficients when "reducing" the index variable (i.e., for slice-driven loop) but not when "resolving" the index variable (i.e., TACO-like co-iteration).

[Peiming]

address comments.

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

 /// and serve as unique identifiers for the corresponding `LatPoint` object.
 using LatPointId = unsigned;
 
 /// `LatSet` identifiers.  These are allocated by `Merger::addSet` (and
 /// by other methods calling that one), and serve as unique identifiers
 /// for the corresponding `SmallVector<LatPointId>` object.
 using LatSetId = unsigned;
 
+/// A pair of level and its corresponding DimLevelType of a tensor.
+using LvlDLTPair = std::pair<Level, DimLevelType>;
+
+/// A pair of loop id and its coefficients. E.g., for affine expression in the
+/// affine map `2 * d0`, loop id = 0, coefficient = 2.
+using LoopCoeffPair = std::pair<LoopId, unsigned>;
+
 /// Tensor expression. Represents an MLIR expression in tensor index notation.
 struct TensorExp final {
   enum class Kind;
 
   /// Child subexpressions for non-leaf expressions.
   struct Children final {
     ExprId e0;
     ExprId e1;
void foreachTensorLoopId(LatPointId p, bool simple,
     }
   }
 
   /// Sets whether the output tensor is sparse or not.
   void setHasSparseOut(bool s) { hasSparseOut = s; }
 
   /// Establishes the two-way map that i <-> <t, lvl, dlt>.
   void setLoopDependentTensorLevel(LoopId i, TensorId t, Level lvl,
-                                   DimLevelType dlt) {
+                                   DimLevelType dlt, unsigned coefficient) {
     assert(isValidLoopId(i) && isValidLevel(t, lvl));
-    assert(!loopToDependencies[i][t].has_value()); // must be the first def
-    loopToDependencies[i][t] = std::make_pair(lvl, dlt);
-    levelToDependentLoop[t][lvl].push_back(i);
+    assert(!loopToUnresolvedLvls[i][t].has_value()); // must be the first def
+    loopToUnresolvedLvls[i][t] = std::make_pair(lvl, dlt);
+    levelToDependentLoop[t][lvl].emplace_back(i, coefficient);
   }
 
   /// Whether the loop has dependent slice.
   bool hasDependentLvl(LoopId i, TensorId t) {
     assert(isValidTensorId(t) && isValidLoopId(i));
-    return loopToDependencies[i][t].has_value();
+    return loopToUnresolvedLvls[i][t].has_value();
   }
 
   /// Returns the list of loop indices which appear in the non-trivial index
   /// expression on t_l, e.g., A[i+j] => {i, j}
-  std::vector<LoopId> &getDependentLoops(TensorId t, Level lvl) {
+  std::vector<LoopCoeffPair> &getDependentLoops(TensorId t, Level lvl) {
     assert(isValidLevel(t, lvl));
     return levelToDependentLoop[t][lvl];
   }
 
   /// Returns the defining [tid, lvl] for the loop.
   std::pair<TensorId, Level> getLoopDefiningLvl(LoopId i) const {
     assert(isValidLoopId(i));
     return loopBounds[i];
   }
 
   /// Checks whether the TensorLoopId represents a tensor level contains
   /// non-trivial index expression.
   bool isLvlWithNonTrivialIdxExp(TensorLoopId b) const {
     const TensorId t = tensor(b);
     const LoopId i = loop(b);
     assert(isValidTensorId(t) && isValidLoopId(i));
-    return loopToDependencies[i][t].has_value();
+    return loopToUnresolvedLvls[i][t].has_value();
   }
 
   /// Checks whether the TensorLoopId represents a sparse tensor level contains
   /// non-trivial index expression.
   bool isSparseLvlWithNonTrivialIdxExp(TensorLoopId b) const {
     if (isLvlWithNonTrivialIdxExp(b)) {
       auto dlt = getLoopDependentLevelType(b);
       return isCompressedDLT(dlt) || isSingletonDLT(dlt);
     }
     return false;
   }
 
   Level getLoopDependentLevel(TensorLoopId b) const {
     assert(isLvlWithNonTrivialIdxExp(b));
-    return loopToDependencies[loop(b)][tensor(b)]->first;
+    return loopToUnresolvedLvls[loop(b)][tensor(b)]->first;
   }
 
   DimLevelType getLoopDependentLevelType(TensorLoopId b) const {
     assert(isLvlWithNonTrivialIdxExp(b));
-    return loopToDependencies[loop(b)][tensor(b)]->second;
+    return loopToUnresolvedLvls[loop(b)][tensor(b)]->second;
   }
 
   /// Convenience getters to immediately access the stored nodes.
   /// These methods return `const&` because the underlying objects must
   /// not be mutated by client code.  The only exception is for mutating
   /// the value associated with an expression, for which there are
   /// dedicated methods below.
   ///
private:
   /// Map that converts pair<TensorId, Level> to the corresponding LoopId.
   std::vector<std::vector<std::optional<LoopId>>> lvlToLoop;
 
   /// Map from a loop to its dependencies if any.
   /// The dependencies of a loop is a set of (tensor, level) pairs.
   /// It is currently only set for non-trivial index expressions.
   /// E.g., A[i+j] => i and j will have dependencies {A0, dlt(A0)} to indicate
   /// that i and j are used in the non-trivial index expression on A0.
-  std::vector<std::vector<std::optional<std::pair<Level, DimLevelType>>>>
-      loopToDependencies;
+  std::vector<std::vector<std::optional<LvlDLTPair>>> loopToUnresolvedLvls;
 
   /// The inverse map of ldxToDependencies from tensor level -> dependent loop
-  /// E.g., A[i+j], we have A0 => {i, j}, to indicate that A0 uses both {i, j}
-  /// to compute its indices.
-  std::vector<std::vector<std::vector<LoopId>>> levelToDependentLoop;
+  /// E.g., A[2i+j], we have A0 => {(2, i), (1, j)}, to indicate that A0 uses
+  /// both {i, j} to compute its indices and the coefficients on the loop id are
+  /// 2 and 1 respectively.
+  std::vector<std::vector<std::vector<LoopCoeffPair>>> levelToDependentLoop;
 
   /// Map from a loop to the [tid, lvl] pair that defines the loop boundary.
   std::vector<std::pair<TensorId, Level>> loopBounds;
 
   llvm::SmallVector<TensorExp> tensorExps;
   llvm::SmallVector<LatPoint> latPoints;
   llvm::SmallVector<SmallVector<LatPointId>> latSets;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_UTILS_MERGER_H_

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

TensorLevel makeTensorLevel(TensorId t, Level l) const {
     // Make sure LoopEmitter, GenericOp, and Merger agree on the number of
     // tensors.
     assert(loopEmitter.getNumManifestTensors() == linalgOp->getNumOperands() &&
            loopEmitter.getNumTensors() == latticeMerger.getNumTensors() &&
            loopEmitter.getOutTensorId() == latticeMerger.getOutTensorID() &&
            loopEmitter.getSynTensorId() == latticeMerger.getSynTensorID());
     return loopEmitter.makeTensorLevel(t, l);
   }
+  TensorLevel makeTensorLevel(std::pair<TensorId, Level> tlPair) const {
+    return makeTensorLevel(tlPair.first, tlPair.second);
+  }
   std::pair<TensorId, Level> unpackTensorLevel(TensorLevel tl) const {
     return loopEmitter.unpackTensorLevel(tl);
   }
   template <class ContainerTy>
   auto unpackTensorLevelRange(ContainerTy &&c) const {
     return loopEmitter.unpackTensorLevelRange(std::forward<ContainerTy>(c));
   }
 

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

 
 /// Returns true if tensor materializes uninitialized into the computation.
 static bool isMaterializing(Value val) {
   return val.getDefiningOp<tensor::EmptyOp>() ||
          val.getDefiningOp<bufferization::AllocTensorOp>();
 }
 
 /// Makes target array's elements sorted according to the `order` array.
-static void sortArrayBasedOnOrder(std::vector<LoopId> &target,
+static void sortArrayBasedOnOrder(std::vector<LoopCoeffPair> &target,
                                   ArrayRef<LoopId> order) {
   std::sort(target.begin(), target.end(),
-            [&order](const LoopId &l, const LoopId &r) {
+            [&order](const LoopCoeffPair &l, const LoopCoeffPair &r) {
               assert(std::addressof(l) == std::addressof(r) || l != r);
               int idxL = -1, idxR = -1;
               for (int i = 0, e = order.size(); i < e; i++) {
-                if (order[i] == l)
+                if (order[i] == l.first)
                   idxL = i;
-                if (order[i] == r)
+                if (order[i] == r.first)
                   idxR = i;
               }
               assert(idxL >= 0 && idxR >= 0);
               return idxL < idxR;
             });
 }
 //===----------------------------------------------------------------------===//
 // Code generation environment constructor and general methods
void CodegenEnv::startEmit() {
   loopEmitter.initialize(
       tensors,
       StringAttr::get(linalgOp.getContext(),
                       linalg::GenericOp::getOperationName()),
       /*hasOutput=*/true,
       /*isSparseOut=*/sparseOut != nullptr, topSort,
       // TODO: compute the map and pass it to loop emitter directly instead of
       // passing in a callback.
-      [this](TensorId t, Level lvl) -> std::vector<std::pair<TensorId, Level>> {
-        // Translates from a list of loop index to a list of [tid, dim] pair.
-        std::vector<LoopId> rLoops = this->merger().getDependentLoops(t, lvl);
-        std::vector<std::pair<TensorId, Level>> ret;
+      /*dependentLvlGetter=*/
+      [this](TensorId t,
+             Level lvl) -> std::vector<std::pair<TensorLevel, unsigned>> {
+        // Translates from a list of loop indices to a list of [tid, lvl] pair.
+        std::vector<LoopCoeffPair> &rLoops = merger().getDependentLoops(t, lvl);
+        std::vector<std::pair<TensorLevel, unsigned>> ret;
         ret.reserve(rLoops.size());
-        for (LoopId l : rLoops)
-          ret.emplace_back(this->merger().getLoopDefiningLvl(l));
+        for (auto [loop, coeff] : rLoops) {
+          TensorLevel tl = makeTensorLevel(merger().getLoopDefiningLvl(loop));
+          ret.emplace_back(tl, coeff);
+        };
         return ret;
       });
 }
 
 std::optional<Operation *> CodegenEnv::genLoopBoundary(
     function_ref<std::optional<Operation *>(MutableArrayRef<Value> parameters)>
         callback) {
   SmallVector<Value> params;

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

public:
   using OutputUpdater = function_ref<Value(OpBuilder &builder, Location loc,
                                            Value memref, Value tensor)>;
 
   /// Optional callback function to set the bound for the synthetic tensor,
   /// which essentially is the dense loop bound.
   using SynTensorBoundSetter =
       function_ref<Value(OpBuilder &builder, Location loc, Level lvl)>;
 
-  // Map from [tid, dim] to a list of dependent [tid, dim] for affine expression
-  // index on sparse tensors.
-  // E.g., for affine index (d0 + d1), it depends on two [tid, dim] that defines
-  // d0 and d1 (for affine expression reduction).
+  // Map from [tid, lvl] to a list of dependent [tidlvl, coeffecient] for
+  // subscript expressions on sparse tensors.
+  //
+  // E.g., for affine index (2 * d0 + d1), it depends on two tidlvls that
+  // defines d0 and d1 (for affine expression reduction) and uses 2 and 1 for
+  // cofficients on d0, d1 respectively.
   // If the list is empty, it means that there is no affine expression on the
-  // input [tid, dim].
+  // input [tid, lvl].
+  //
   // NOTE: The caller is responsible to ensure that the order of the returned
   // list to be consistent with the topological order of the iteration graph,
   // otherwise the loop emitter might reduce a wrong dependent index variable
   // when generating slice-driven loops.
   using DependentLvlGetter =
-      function_ref<std::vector<std::pair<TensorId, Level>>(TensorId, Level)>;
+      function_ref<std::vector<std::pair<TensorLevel, unsigned>>(TensorId,
+                                                                 Level)>;
 
   LoopEmitter() = default;
 
   /// Takes an array of input tensors, which the generated loops will
   /// iterate over.  Each tensor is given a `TensorId` (numerically equal
   /// to the position of that tensor `Value` in the array).  Setting
   /// `isSparseOut` indicates that the sparse output tensor is empty,
   /// so the loop emitter will generate loops over it according to the
SliceInfo(Value minCrd, Value offset, Value isNonEmpty,
           slicedOnLvl(slicedOnLvl), depth(depth) {
       // TODO: use std::optional<pair<Level, minCrd>>
       assert(!slicedOnLvl || minCrd);
     }
 
     // Whether this is the tensor that has not yet been sliced.
     bool isInitialTensor() const { return !slicedOnLvl.has_value(); }
 
-    Value minCrd;                     // the minimum coordinate of the slice.
-    Value offset;                     // the offset of the current slice.
-    Value isNonEmpty;                 // whether the slice is empty.
+    Value minCrd;     // the minimum coordinate of the slice.
+    Value offset;     // the *absolute* offset of the current slice.
+    Value isNonEmpty; // whether the slice is empty.
     std::optional<Level> slicedOnLvl; // the level on which the slice is done
     unsigned depth; // the depth (relative to dependentDimMap[tid][lvl]).
   };
 
   ///
   /// Enums for different kinds of loop conditions.
   ///
 
private:
   /// TODO: We should instead use an SSA value for the index.
   void invalidateSliceIterIdx(OpBuilder &builder, Location loc, TensorId tid,
                               Level lvl);
   /// Generates code to get the first non-empty slice of tid on lvl.
   /// return true if has already been resolved.
   bool genSliceBegin(OpBuilder &builder, Location loc, TensorId tid, Level lvl);
 
   /// Generates code to get the next non-empty slices of tid on lvl.
-  void genSliceNextInduction(OpBuilder &builder, Location loc,
-                             const Operation *whileOp, TensorId tid, Level lvl,
-                             SmallVectorImpl<Value> &operands,
-                             unsigned &retIdx);
+  /// Returns a tuple of values for <NonEmpty, MinCrd, AbsOffset> (see
+  /// SliceInfo) respectively.
+  std::tuple<Value, Value, Value> genSliceNextInduction(OpBuilder &builder,
+                                                        Location loc,
+                                                        TensorId tid,
+                                                        Level lvl);
 
   /// A optional string attribute that should be attached to the loop
   /// generated by loop emitter, it might help following passes to identify
   /// loops that operates on sparse tensors more easily.
   StringAttr loopTag;
   /// Whether the loop emitter needs to treat the last tensor as the output
   /// tensor.
   bool hasOutput;
private:
   //
 
   /// Whether the sparse input is a slice.
   std::vector<bool> isSparseSlices;
   /// Values related to slices.
   std::vector<std::vector<Value>> sliceOffsets;
   std::vector<std::vector<Value>> sliceStrides;
 
-  // Map from [tid, level] to a list of dependent [tid, level].
-  // See comments for `DependentDimGetter`.
-  std::vector<std::vector<std::vector<std::pair<TensorId, Level>>>>
+  // Map from [tid, level] to a list of dependent [tidlevel, coefficient].
+  // See comments for `DependentLvlGetter`.
+  std::vector<std::vector<std::vector<std::pair<TensorLevel, unsigned>>>>
       dependentLvlMap;
 
   // The cached position buffer for the slices, they serve the same purpose as
   // ptrBuffer for compressed dimensions.
   // But they always starts with the first pidx pointing to coord > slice.offset
   // to avoid iteration from the beginning.
   std::vector<std::vector<std::vector<Value>>> slicePosBuffer;
 
-  // The cached size for each slices.
-  std::vector<std::vector<std::vector<Value>>> sliceSizes;
+  // The (size, stride) for each conceptual slice used for index reduction
+  // loops.
+  std::vector<std::vector<std::vector<std::pair<Value, unsigned>>>> sliceMeta;
 
   // The number of reduced dependencies on a tensor level so far.
   std::vector<std::vector<unsigned>> levelReducedDep;
 
   // sliceStack[tid] holds the generated slice stack on tid.
   std::vector<std::vector<SliceInfo>> sliceStack;
 
   /// TODO: not yet used, it should track the current level for each tensor

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

auto whileOp = builder.create<scf::WhileOp>(
         YIELD(nextPos);
       });
   // Return the segment high.
   return whileOp.getResult(0);
 }
 
 Value LoopEmitter::genSparseCrd(OpBuilder &builder, Location loc, TensorId tid,
                                 Level lvl) {
-  Value crd = C_IDX(0);
   // A load on the coordinates array yields the coordinate.
   const Value mem = coordinatesBuffers[tid][lvl];
   /// FIXME: See the [CLARIFY_POSITS_LVL] note in the header.
   const Value pos = posits[tid][lvl];
-  const Value off = genIndexLoad(builder, loc, mem, pos);
-  // Linearized the coordinates within the same collapse reassociation.
-  crd = ADDI(crd, off);
+  const Value crd = genIndexLoad(builder, loc, mem, pos);
   return crd;
 }
 
 LoopEmitter::LoopEmitter(ValueRange tensors, StringAttr loopTag, bool hasOutput,
                          bool isSparseOut, ArrayRef<LoopId> topSort,
                          DependentLvlGetter dimGetter) {
   initialize(tensors, loopTag, hasOutput, isSparseOut, topSort, dimGetter);
 }
void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
   // These zeros will be overwritten below, but we need to initialize
   // them to something since we'll need random-access assignment.
   this->loopIdToOrd.assign(numLoops, 0);
   this->loopStack.reserve(numLoops);
   this->loopSeqStack.reserve(numLoops);
 
   // Index-reduction related fields.
   this->dependentLvlMap.assign(
-      numTensors, std::vector<std::vector<std::pair<TensorId, Level>>>());
+      numTensors, std::vector<std::vector<std::pair<TensorLevel, unsigned>>>());
   this->slicePosBuffer.assign(numTensors, std::vector<std::vector<Value>>());
-  this->sliceSizes.assign(numTensors, std::vector<std::vector<Value>>());
+  this->sliceMeta.assign(
+      numTensors, std::vector<std::vector<std::pair<Value, unsigned>>>());
   this->sliceStack.assign(numTensors, std::vector<SliceInfo>());
   this->levelReducedDep.assign(numTensors, std::vector<unsigned>());
 
   // Initialize nested types of `TensorId`-indexed fields.
   for (TensorId tid = 0; tid < numTensors; tid++) {
     Level lvlRank;
     if (tid == synTensorId) {
       // Synthetic tensor (conceptually) is an all-dense tensor with rank equal
for (TensorId tid = 0; tid < numTensors; tid++) {
     coords[tid].assign(lvlRank, Value());
     positionsBuffers[tid].assign(lvlRank, Value());
     coordinatesBuffers[tid].assign(lvlRank, Value());
     sliceOffsets[tid].assign(lvlRank, Value());
     sliceStrides[tid].assign(lvlRank, Value());
 
     // Slice-driven loops related initialization.
     levelReducedDep[tid].assign(lvlRank, 0);
-    dependentLvlMap[tid].assign(lvlRank,
-                                std::vector<std::pair<TensorId, Level>>());
+    dependentLvlMap[tid].assign(
+        lvlRank, std::vector<std::pair<TensorLevel, unsigned>>());
     slicePosBuffer[tid].assign(lvlRank, std::vector<Value>());
-    sliceSizes[tid].assign(lvlRank, std::vector<Value>());
+    sliceMeta[tid].assign(lvlRank, std::vector<std::pair<Value, unsigned>>());
     sliceStack[tid].emplace_back(/*minCrd=*/Value(),
                                  /*offset=*/Value(), /*isNonEmpty*/ Value(),
                                  std::nullopt, 0);
     if (dimGetter && !isSynTensor(tid)) {
       for (Level l = 0; l < lvlRank; l++) {
         dependentLvlMap[tid][l] = dimGetter(tid, l);
         unsigned depends = dependentLvlMap[tid][l].size();
         if (depends == 0)
           continue;
-        // We need `depends - 1` slices to fully  the affine expression.
-        sliceSizes[tid][l].assign(depends - 1, nullptr);
+        sliceMeta[tid][l].assign(depends, std::make_pair(nullptr, 0));
+        // We need `depends - 1` slices to fully reduce the affine expression.
         slicePosBuffer[tid][l].assign(depends - 1, nullptr);
       }
     }
   }
 
   // Construct the inverse of the `topSort` from the sparsifier.
   // This is needed to map `AffineDimExpr`s back to the `LoopOrd`
   // used in loop emitter.
void LoopEmitter::initializeLoopEmit(
   for (TensorId t = 0, e = tensors.size(); t < e; t++) {
     auto rtp = dyn_cast<RankedTensorType>(tensors[t].getType());
     if (!rtp)
       continue;
 
     Level lvlRank = SparseTensorType(rtp).getLvlRank();
     for (Level lvl = 0; lvl < lvlRank; lvl++) {
       if (!dependentLvlMap[t][lvl].empty()) {
-        ArrayRef<std::pair<TensorId, Level>> depLvls = dependentLvlMap[t][lvl];
+        ArrayRef<std::pair<TensorLevel, unsigned>> depLvls =
+            dependentLvlMap[t][lvl];
         // Needs at least two operands to form a non-trivial affine expression.
-        assert(depLvls.size() > 1);
+        assert(depLvls.size() == sliceMeta[t][lvl].size());
 
         Value size = c0;
-        for (unsigned e = depLvls.size() - 1; e >= 1; e--) {
-          auto [dt, dd] = depLvls[e];
-          size = ADDI(size, lvlSizes[dt][dd]);
-          sliceSizes[t][lvl][e - 1] = size;
+        for (int e = depLvls.size() - 1; e >= 0; e--) {
+          auto [dt, dl] = unpackTensorLevel(depLvls[e].first);
+          unsigned stride = depLvls[e].second;
+          Value stridedSize = lvlSizes[dt][dl];
+          if (stride != 1)
+            stridedSize = MULI(stridedSize, C_IDX(stride));
+          size = ADDI(size, stridedSize);
+          sliceMeta[t][lvl][e] = std::make_pair(size, stride);
         }
       }
     }
   }
   localInsertPos = builder.getInsertionPoint()->getPrevNode();
 }
 
 void LoopEmitter::categorizeLoopCondition(
case LoopCondKind::SparseAffineCond: {
     assert(ivs.size() == 1);
     Value crdHi; // loop upper bound
     {
       OpBuilder::InsertionGuard guard(builder);
       Operation *loop = builder.getInsertionBlock()->getParentOp();
       // crdHi is a loop invariant, hosit the computation outside the loop.
       if (llvm::isa_and_nonnull<scf::WhileOp>(loop))
         builder.setInsertionPoint(loop);
-      crdHi = ADDI(getMostRecentSliceOnLvl(tid, lvl).offset,
-                   sliceSizes[tid][lvl].back());
+      auto [size, stride] = sliceMeta[tid][lvl].back();
+      assert(stride == 1 && "Not yet implemented");
+      crdHi = ADDI(getMostRecentSliceOnLvl(tid, lvl).offset, size);
     }
     assert(crdHi);
     return genSparseReducedAffineCond(builder, loc,
                                       coordinatesBuffers[tid][lvl], crdHi,
                                       ivs[0], highs[tid][lvl]);
   }
   case LoopCondKind::SparseAffineUnRedCond: {
     assert(ivs.size() == 3);
bool LoopEmitter::shouldIteratedByForLoop(ArrayRef<TensorLvlCond> sparseConds,
                                           bool genDedup) {
   assert(llvm::all_of(sparseConds,
                       [](TensorLvlCond c) { return isSparseCond(c.second); }));
 
   // If we need to co-iterate over two sparse tensors, we need a while loop
   if (sparseConds.size() > 1)
     return false;
 
-  // We also need a while loop for levels with affine index expression for
+  // We also need a while loop for levels with affine index expression and
   // non-unique levels when deduplication is required.
   if (sparseConds.size() == 1) {
     auto [tid, lvl] = unpackTensorLevel(sparseConds.back().first);
     return !isAffineIdxCond(sparseConds.back().second) &&
            !(genDedup && !isUniqueDLT(lvlTypes[tid][lvl]));
   }
 
   return true;
if (shouldIteratedByForLoop(spConds, genDedup) && !needsUniv) {
     auto [tid, lvl] = unpackTensorLevel(tlCond.first);
     Value lo = isSparseCond(loopCondKind)
                    ? posits[tid][lvl]           // current offset
                    : loopSeqStack.back().first; // universal index
     Value hi = highs[tid][lvl];
     if (isDenseCond(loopCondKind) && isAffineIdxCond(loopCondKind)) {
       bool unReduc = isAffineIdxUnRedCond(loopCondKind);
       assert(unReduc == !depFullyReduced(tid, lvl));
-      hi = sliceSizes[tid][lvl][sliceStack[tid].back().depth - 1];
+      auto [size, stride] = sliceMeta[tid][lvl][sliceStack[tid].back().depth];
+      assert(stride == 1 && "Not yet implemented");
+      hi = size;
       if (unReduc) {
         // Adjust for loop hi for dense slice-driven loop.
         hi = SUBI(lvlSizes[tid][lvl], hi);
         hi = ADDI(hi, C_IDX(1));
       }
     }
     std::tie(l, iv) = emitForLoopOverTensorAtLvl(builder, loc, tid, lvl, lo, hi,
                                                  reduc, tryParallel);
for (auto [dnTidLvl, denseLoopCond] : dnConds) {
     assert(isDenseDLT(lvlTypes[tid][lvl]));
 
     if (isAffineIdxCond(denseLoopCond)) {
       // Pushes sliced levels to build correct LoopInfo.
       bool unReduc = isAffineIdxUnRedCond(denseLoopCond);
       SliceInfo &info = sliceStack[tid].back();
       // Pushes sliced dense loop info to tell LoopEmitter how to exit it.
       sliceInfo.emplace_back(tid, lvl, /*fullyReduced=*/!unReduc);
+      // FIXME: The offset and position iterator need to be adjusted when the
+      // slice is strided.
       if (unReduc) {
         assert(*info.slicedOnLvl == lvl);
         // Update the slice information as we enter the new loop.
         info.minCrd = info.offset = iv;
         info.isNonEmpty = constantI1(builder, loc, true);
         levelReducedDep[tid][lvl]++;
       } else {
         posits[tid][lvl] =
while (leafLvl < stt.getLvlRank() && !dependentLvlMap[tid][leafLvl].empty()) {
     leafLvl++;
   }
 
   Level curLvl = rootLvl + 1;
   // Prunes all denses subtree.
   while (curLvl < leafLvl && isDenseDLT(lvlTypes[tid][curLvl])) {
     // One step forward in parent level results in forwarding `slice.size` step
     // in child dense level.
-    fcnt = MULI(sliceSizes[tid][curLvl].back(), fcnt);
+    auto [size, stride] = sliceMeta[tid][curLvl].back();
+    assert(stride == 1 && "Not yet implemented");
+    fcnt = MULI(size, fcnt);
     curLvl++;
   }
 
   Value nxPosPtr = nullptr;
   if (curLvl < leafLvl) {
     assert(!isDenseDLT(lvlTypes[tid][curLvl]));
     // The first compressed level, setting up the position pointer for it.
     Value sPosBuf = slicePosBuffer[tid][curLvl].back();
void LoopEmitter::exitWhileLoop(OpBuilder &builder, Location loc,
   for (auto [tid, lvl, resolved] : loopInfo.sliceDrivenInfo) {
     // TODO: handle dense.
     assert(isCompressedDLT(lvlTypes[tid][lvl]));
     levelReducedDep[tid][lvl]--;
     if (!resolved) {
       // TODO: support coiterating multiple slices
       assert(loopInfo.trivialTidLvls.empty() &&
              loopInfo.sliceDrivenInfo.size() == 1);
-      genSliceNextInduction(builder, loc, whileOp, tid, lvl, operands, o);
+      auto [nxNonEmpty, nxMinCrd, nxAbsOffset] =
+          genSliceNextInduction(builder, loc, tid, lvl);
+      // Update while loop induction operands.
+      operands.push_back(nxNonEmpty);
+      operands.push_back(nxMinCrd);
+      operands.push_back(nxAbsOffset);
+
+      // Update the slice stack.
+      SliceInfo &info = sliceStack[tid].back();
+      info.isNonEmpty = whileOp.getResult(o++);
+      info.minCrd = whileOp.getResult(o++);
+      info.offset = whileOp.getResult(o++);
       continue;
     }
 
     Value forwarded = nullptr;
     if (loopInfo.trivialTidLvls.empty() &&
         loopInfo.sliceDrivenInfo.size() == 1) {
       // Forwards the position iterator.
       operands.push_back(ADDI(posits[tid][lvl], one));

 //   body_builder
 //   loopLo ++;
 // }
 std::pair<Operation *, ValueRange> LoopEmitter::genSliceLvlTraverseLoop(
     OpBuilder &builder, Location loc, Value posLo, Value posHi, Value offset,
     Value size, TensorId tid, Level lvl, ValueRange userReduc,
     LoopBodyBuilder bodyBuilder) {
   Value c1 = C_IDX(1);
-  Value sliceHi = ADDI(offset, sliceSizes[tid][lvl].back());
+  auto [sliceSz, stride] = sliceMeta[tid][lvl].back();
+  assert(stride == 1 && "Not yet implemented");
+  Value sliceHi = ADDI(offset, sliceSz);
+
   SmallVector<Value> reduc{posLo}; // loop lower bounds
   const unsigned numMetaReduc = reduc.size();
 
   // Append user required reduction value.
   reduc.append(userReduc.begin(), userReduc.end());
   scf::WhileOp whileOp = builder.create<scf::WhileOp>(
       loc, ValueRange(reduc).getTypes(), reduc,
       /*beforeBuilder=*/
ValueRange LoopEmitter::genUnResolvedSliceTreeTraverse(
   // Wraps body builder and inserts a extra counting instruction at the end.
   auto wrapped = [bodyBuilder](OpBuilder &builder, Location loc, Value iv,
                                MutableArrayRef<Value> reduc) {
     bodyBuilder(builder, loc, iv, reduc.drop_back());
     // Increments the counter.
     reduc.back() = ADDI(reduc.back(), C_IDX(1));
   };
 
+  // FIXME: Need special handling when the previous unresolved slice is strided:
+  // We probably need to filter out coordinates that is not on stride.
   if (firstResLvl.has_value()) {
     // Overwrite position when the first level is fully resolved.
     pos = posits[firstResLvl->first][firstResLvl->second];
     ip = builder.saveInsertionPoint();
   } else {
     const SliceInfo &frontSlice = *unResLvls.back();
     Level firstLvl = *frontSlice.slicedOnLvl;
     if (!lvlFullyResolved(tid, firstLvl)) {
if (!lvlFullyResolved(tid, firstLvl)) {
               // generate traversal for each level.
               Value loopLo = genIndexLoad(builder, loc, sPtrBuf, iv);
               Value loopHi = genIndexLoad(builder, loc, sPtrBuf, ADDI(iv, c1));
               // We need to remember the starting index for next level's
               // position, because slice-driven loop breaks the level into
               // non-consecutive segments.
               builder.create<memref::StoreOp>(loc, iterArgs.back(), sPtrBuf,
                                               ADDI(iv, c2).getResult());
+
+              auto [size, stride] = sliceMeta[tid][firstLvl].back();
+              assert(stride == 1 && "Not yet implemented");
               ValueRange itArgs =
                   genSliceLvlTraverseLoop(
-                      builder, loc, loopLo, loopHi, offset,
-                      sliceSizes[tid][firstLvl].back(), tid, firstLvl, iterArgs,
+                      builder, loc, loopLo, loopHi, offset, size, tid, firstLvl,
+                      iterArgs,
                       [&](OpBuilder &builder, Location, Value iv,
                           MutableArrayRef<Value> reduc) {
                         ip = builder.saveInsertionPoint();
                         pos = iv;
                         innerArgs.assign(reduc.begin(), reduc.end());
                       })
                       .second;
               YIELD(itArgs);
             });
       } else if (isDenseDLT(lvlTypes[tid][firstLvl])) {
         assert(firstLvl == 0); // This must be the first level.
         Value lb = frontSlice.offset;
-        Value sliceSz =
-            sliceSizes[tid][*frontSlice.slicedOnLvl][frontSlice.depth - 1];
+        auto [sliceSz, stride] =
+            sliceMeta[tid][*frontSlice.slicedOnLvl][frontSlice.depth];
+        assert(stride == 1 && "Not yet implemented");
         Value ub = ADDI(lb, sliceSz);
         outerMost = builder.create<scf::ForOp>(
             loc, lb, ub, c1, innerArgs,
             [&](OpBuilder &builder, Location loc, Value iv,
                 ValueRange iterArgs) {
               ip = builder.saveInsertionPoint();
               pos = iv;
               innerArgs.assign(iterArgs.begin(), iterArgs.end());
             });
       }
       // We generated the loop for the first slice above, now remove it.
       unResLvls = unResLvls.drop_back();
     }
   }
   // Reset the insertion point into the loop body.
   builder.restoreInsertionPoint(ip);
   if (!unResLvls.empty()) {
     // Fills in dense slices levels in between.
     SmallVector<Value> lbs, ubs, steps, lvlSzs;
     for (const SliceInfo *slice : llvm::reverse(unResLvls)) {
       Level sliceLvl = *slice->slicedOnLvl;
       assert(isDenseDLT(lvlTypes[tid][sliceLvl]));
       Value offset = slice->offset;
-      Value sliceSz = sliceSizes[tid][sliceLvl][slice->depth - 1];
+      auto [sliceSz, stride] = sliceMeta[tid][sliceLvl][slice->depth];
+      assert(stride == 1 && "Not yet implemented");
       lbs.push_back(offset);
       ubs.push_back(ADDI(offset, sliceSz));
       steps.push_back(c1);
       lvlSzs.push_back(lvlSizes[tid][sliceLvl]);
     }
     auto denseNest =
         scf::buildLoopNest(builder, loc, lbs, ubs, steps, innerArgs,
                            [&innerArgs, &lvlSzs, &pos, bodyBuilder](
Value c0 = C_IDX(0), c1 = C_IDX(1), c2 = C_IDX(2), c3 = C_IDX(3),
         c4 = C_IDX(4);
   if (isDenseDLT(lvlTypes[tid][lvl])) {
     // Dense slice begin is trivial.
     sliceStack[tid].emplace_back(/*minCoord=*/c0, /*offset=*/c0,
                                  /*nonEmpty=*/constantI1(builder, loc, true),
                                  lvl, /*depth=*/1);
     return;
   }
-  Value size = sliceSizes[tid][lvl][0];
+  auto [nxSz, stride] = sliceMeta[tid][lvl][1];
+  assert(stride == 1 && "Not yet implemented");
   Value sPtrBuf = slicePosBuffer[tid][lvl][0];
   Value pHi, pLo;
   if (lvl == 0) {
     pLo = c0;
     pHi = genIndexLoad(builder, loc, positionsBuffers[tid][0], c1);
   } else {
     pLo = genIndexLoad(builder, loc, positionsBuffers[tid][lvl],
                        posits[tid][lvl - 1]);
void LoopEmitter::genResolvedSliceBegin(OpBuilder &builder, Location loc,
   // The minimal coord must be at the first on ordered level.
   // FIXME: Technically we should load the coord only when the slice is
   // nonempty. though we assume that even on empty sparse tensors, a non-empty
   // ptr/idx buffer is allocated for each level so it would not cause OOB to
   // avoid generating a ifOp here.
   Value minCrd = genIndexLoad(builder, loc, coordinatesBuffers[tid][0], c0);
 
   // FIXME: We need the relative offset related to the base slice.
-  Value absOffset = offsetFromMinCoord(builder, loc, minCrd, size, isNonEmpty);
+  Value absOffset = offsetFromMinCoord(builder, loc, minCrd, nxSz, isNonEmpty);
   sliceStack[tid].emplace_back(minCrd, absOffset, isNonEmpty, lvl,
                                /*depth=*/1);
 }
 
 // Fills in the slicePosBuffer before slice-driven loop begin.
 // TODO: it can only handle all compressed tensors.
 //
 // // Loop generated by `genUnResolvedSliceTreeTraverse`

 //     }
 //     loopLo ++;
 //   }
 // }
 void LoopEmitter::genUnResolvedSliceBegin(OpBuilder &builder, Location loc,
                                           TensorId tid, Level lvl) {
   Value c0 = C_IDX(0), c1 = C_IDX(1), c2 = C_IDX(2);
   unsigned depth = levelReducedDep[tid][lvl];
-  Value size = sliceSizes[tid][lvl][depth]; // Dense slice begin is trivial
+  // TODO: handle case when the current slice stride is not one.
+  assert(sliceMeta[tid][lvl][depth].second == 1 && "Not yet implemented");
+
+  // The remaining slice size after reduction.
+  Value remSz = sliceMeta[tid][lvl][depth + 1].first;
+  // Dense slice begin is trivial
   if (isDenseDLT(lvlTypes[tid][lvl])) {
     sliceStack[tid].emplace_back(c0, c0, constantI1(builder, loc, false), lvl,
                                  depth + 1);
     return;
   }
 
   assert(isCompressedDLT(lvlTypes[tid][lvl]));
   // Unhandled Cases:
void LoopEmitter::genUnResolvedSliceBegin(OpBuilder &builder, Location loc,
 
   Value isNonEmpty = result[0];
   Value minCrd = result[1];
   // Two metadata [memSize, idx].
   // TODO: Can use an SSA value for these two metadata
   builder.create<memref::StoreOp>(loc, result[2], sPtrBuf, c0);
   builder.create<memref::StoreOp>(loc, c0, sPtrBuf, c1);
   // FIXME: we need the relative offset related to the base slice.
-  Value absOffset = offsetFromMinCoord(builder, loc, minCrd, size, isNonEmpty);
+  Value absOffset = offsetFromMinCoord(builder, loc, minCrd, remSz, isNonEmpty);
   sliceStack[tid].emplace_back(minCrd, absOffset, isNonEmpty, lvl, depth + 1);
 }
 
 bool LoopEmitter::genSliceBegin(OpBuilder &builder, Location loc, TensorId tid,
                                 Level lvl) {
   Value c1 = C_IDX(1), c2 = C_IDX(2);
 
   if (depFullyReduced(tid, lvl)) {
if (slicePosBuffer[tid][lvl][0] == nullptr && !isDenseDLT(lvlType)) {
     // level. We at most need to a memref<d0xindex>.
     // NOTE: this is apperantly an over-approximation when the previous
     // level is compressed, and we can compute a precise memory size
     // inside the loops. But that would also requires us to allocate/free
     // memorys in loops.
     // TODO: Maybe using allocaScopeOp inside the loop to resolve the issue?
     for (Level curLevel = lvl;
          curLevel >= 1 && !lvlFullyResolved(tid, curLevel - 1); curLevel--) {
-      auto depth = remDepOnLevel(tid, curLevel - 1);
-      assert(sliceSizes[tid][lvl].size() >= depth);
-      Value sz = *(sliceSizes[tid][lvl].rbegin() + depth - 1);
+      // We only handle cases when all the previously unresolved levels are
+      // fully reduced.
+      assert(depFullyReduced(tid, curLevel - 1));
+      assert(!sliceMeta[tid][curLevel - 1].empty());
+      auto [sz, stride] = sliceMeta[tid][curLevel - 1].back();
+      assert(stride == 1 && "Not yet implemented");
       bufSize = MULI(bufSize, sz);
     }
     // For a triple of [pLo, pHi, pPtr]. Note that we can not compress pHi
     // because slice creates segments in the index buffer so that the pHi for
     // the current level is no longer the pLo for the next level.
     bufSize = MULI(bufSize, C_IDX(kSliceIterWidth));
     // Additional two metadata {memSize, idx} at head.
     bufSize = ADDI(bufSize, c2);
void LoopEmitter::invalidateSliceIterIdx(OpBuilder &builder, Location loc,
                                          TensorId tid, Level lvl) {
   for (unsigned i = 0; i <= lvl; i++) {
     if (!isDenseDLT(lvlTypes[tid][i]) && !dependentLvlMap[tid][i].empty()) {
       updateSlicePosPtr(builder, loc, slicePosBuffer[tid][i].back(), C_IDX(0));
     }
   }
 }
 
-void LoopEmitter::genSliceNextInduction(OpBuilder &builder, Location loc,
-                                        const Operation *op, TensorId tid,
-                                        Level lvl,
-                                        SmallVectorImpl<Value> &operands,
-                                        unsigned &retIdx) {
+std::tuple<Value, Value, Value>
+LoopEmitter::genSliceNextInduction(OpBuilder &builder, Location loc,
+                                   TensorId tid, Level lvl) {
   if (!isCompressedDLT(lvlTypes[tid][lvl]))
     llvm_unreachable("TODO");
 
   // else generate code to compute next non empty slice.
   Value c0 = C_IDX(0), c1 = C_IDX(1), c2 = C_IDX(2);
 
-  auto whileOp = llvm::cast<scf::WhileOp>(op);
   SliceInfo &info = sliceStack[tid].back();
   assert(info.slicedOnLvl == lvl);
   //
   // We forward to the next non empty slice by
   // if (minCrd > offset) {
   //   offset += 1
   // } else {
   //    minCrd = nextMinInSlice();
auto forOp = scf::buildLoopNest(
               loc, CMPI(ult, newMin.getResult(0), curMinCrd),
               newMin.getResult(0), curMinCrd);
           return {curMinCrd, isNonEmpty};
         });
 
     builder.setInsertionPointAfter(forOp.loops.front());
     // minOffset = minCrd + 1 >= size ? minCrd + 1 - size : c0
     Value tmp = ADDI(forOp.results.front(), c1);
-    Value minOffset = SUBI(tmp, sliceSizes[tid][lvl][info.depth - 1]);
-    Value p = CMPI(uge, tmp, sliceSizes[tid][lvl][info.depth - 1]);
+    auto [size, stride] = sliceMeta[tid][lvl][info.depth];
+    assert(stride == 1 && "Not yet implemented");
+    Value minOffset = SUBI(tmp, size);
+    Value p = CMPI(uge, tmp, size);
     minOffset = SELECT(p, minOffset, c0);
+
     SmallVector<Value, 3> yields;
     yields.assign(forOp.results.begin(), forOp.results.end());
     yields.push_back(minOffset);
     YIELD(yields);
   }
 
   Value nextMinCrd = ifOp.getResults()[0];
   Value nextNonEmpty = ifOp.getResults()[1];
 
   // The next offset should at least be offset + 1;
   Value minOffset = ifOp.getResults()[2];
   Value nxOffset = ADDI(info.offset, c1);
   Value maxPred = CMPI(ugt, minOffset, nxOffset);
   Value nextAbsOffset = SELECT(maxPred, minOffset, nxOffset);
 
-  Value sliceUB = ADDI(nextAbsOffset, sliceSizes[tid][lvl][info.depth - 1]);
+  auto [size, stride] = sliceMeta[tid][lvl][info.depth];
+  assert(stride == 1 && "Not yet implemented");
+  Value sliceUB = ADDI(nextAbsOffset, size);
 
   // FIXME: this only works if there is only one parent.
   assert(info.depth - 1 == 0);
   // nextNonEmpty = nextNonEmpty && slice upper bound <= parent upperbound.
   nextNonEmpty = ANDI(nextNonEmpty, CMPI(ule, sliceUB, lvlSizes[tid][lvl]));
 
   // FIXME: compute relative offset.
   assert(info.depth - 1 == 0);
   Value nextRelOffset = nextAbsOffset;
   nextRelOffset = SELECT(nextNonEmpty, nextRelOffset, c0);
-
-  operands.push_back(nextNonEmpty);
-  operands.push_back(nextMinCrd);
-  operands.push_back(nextAbsOffset); // we push the absolute offset.
-
-  // Update the slice stack.
-  info.isNonEmpty = whileOp.getResult(retIdx++);
-  info.minCrd = whileOp.getResult(retIdx++);
-  info.offset = whileOp.getResult(retIdx++);
+  return std::make_tuple(nextNonEmpty, nextMinCrd, nextAbsOffset);
 }
 
 #undef CMPI
 #undef C_IDX
 #undef YIELD
 #undef ADDI
 #undef ANDI
 #undef SUBI
 #undef MULI
 #undef SELECT

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

 ///
 /// It rejects cases (returns false)
 /// 1st, when the same index is used more than once, e.g., A[i+j][i]
 /// 2nd, when multiplication is used in the non-trivial index expression.
 /// 3rd, when a constant operand is used in the non-trivial index expression.
 ///
 /// TODO: constant should be easy to handle.
 static bool findDepIdxSet(Merger &merger, TensorId tensor, Level lvl,
-                          AffineExpr a, DimLevelType dlt,
-                          bool isSubExp = false) {
+                          AffineExpr a, DimLevelType dlt, bool isSubExp = false,
+                          int64_t coefficient = 1) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
+    // Only allow positive coefficients on AffineDimExpr.
+    if (coefficient <= 0)
+      return false;
+
     const LoopId ldx = merger.makeLoopId(a.cast<AffineDimExpr>().getPosition());
     if (!isUndefDLT(merger.getLvlType(tensor, ldx)))
       return false; // used more than once, e.g., A[i][i]
 
     // TODO: Generalizes the following two cases. A[i] (with trivial index
     // expression) can be treated as a special affine index expression. We do
     // not necessarily need to differentiate them.
-    if (!isSubExp)
+    if (!isSubExp) {
+      assert(coefficient == 1);
       merger.setLevelAndType(tensor, ldx, lvl, dlt);
+    }
 
     if (isSubExp) {
       // The current loops appears in more than one affine expressions on the
       // same tensor. We can not handle this case. e.g., A[i+j][i+k], `i` is
       // used twice.
       if (merger.hasDependentLvl(ldx, tensor)) {
         // TODO: This can be supported by coiterate slices if the loop idx is
         // appeared on affine index for different tensor, or take slice on
         // mulitple dimensions when it is on the same tensor.
         // E.g.,
         // `d0 + d1` for indexing t0[lvl0] and `d0 + d2` for indexing t1[lvl0]
         // d0_1 = getNextSliceOffset t0 along lvl0
         // d0_2 = getNextSliceOffset t1 along lvl0
         // if d0_1 == d0_2 then d0 = d0_1 = d0_1
         // else increase min(d0_1, d0_2).
         return false;
       }
-      merger.setLoopDependentTensorLevel(ldx, tensor, lvl, dlt);
+      merger.setLoopDependentTensorLevel(ldx, tensor, lvl, dlt, coefficient);
     }
     return true;
   }
   case AffineExprKind::Constant:
-  case AffineExprKind::Mul:
-    // TODO: Support Mul and Constant AffineExp for slice-based codegen
+    // TODO: Support Constant AffineExp for slice-based codegen
+  case AffineExprKind::Mul: {
+    // TODO: Support index expression like `2 * d0`, we now only support more
+    // complicated cases like `2 * d0 + d1`.
+    if (!isSubExp)
       return false;
+    auto binOp = a.cast<AffineBinaryOpExpr>();
+    auto lhs = binOp.getLHS(), rhs = binOp.getRHS();
+    if (rhs.isa<AffineConstantExpr>())
+      std::swap(lhs, rhs);
+    // Must be in form of `constant * d`.
+    assert(lhs.isa<AffineConstantExpr>() && rhs.isa<AffineDimExpr>());
+    int64_t coefficient = lhs.cast<AffineConstantExpr>().getValue();
+    return findDepIdxSet(merger, tensor, lvl, rhs, dlt, isSubExp, coefficient);
+  }
   case AffineExprKind::Add: {
     auto binOp = a.cast<AffineBinaryOpExpr>();
     return findDepIdxSet(merger, tensor, lvl, binOp.getLHS(), dlt, true) &&
            findDepIdxSet(merger, tensor, lvl, binOp.getRHS(), dlt, true);
   }
   default:
     return false;
   }

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

: outTensor(numInputOutputTensors - 1),
       numTensors(numInputOutputTensors + 1), numNativeLoops(numNativeLoops),
       numLoops(numNativeLoops + numFilterLoops), hasSparseOut(false),
       lvlTypes(numTensors,
                std::vector<DimLevelType>(numLoops, DimLevelType::Undef)),
       loopToLvl(numTensors,
                 std::vector<std::optional<Level>>(numLoops, std::nullopt)),
       lvlToLoop(numTensors,
                 std::vector<std::optional<LoopId>>(maxLvlRank, std::nullopt)),
-      loopToDependencies(
-          numLoops, std::vector<std::optional<std::pair<Level, DimLevelType>>>(
+      loopToUnresolvedLvls(numLoops, std::vector<std::optional<LvlDLTPair>>(
                                          numTensors, std::nullopt)),
-      levelToDependentLoop(numTensors, std::vector<std::vector<LoopId>>(
-                                           maxLvlRank, std::vector<LoopId>())),
+      levelToDependentLoop(numTensors,
+                           std::vector<std::vector<LoopCoeffPair>>(
+                               maxLvlRank, std::vector<LoopCoeffPair>())),
       loopBounds(numLoops, std::make_pair(numTensors, numLoops)) {}
 
 //===----------------------------------------------------------------------===//
 // Lattice methods.
 //===----------------------------------------------------------------------===//
 
 ExprId Merger::addTensorExp(TensorId t) {
   assert(isValidTensorId(t));