[mlir][sparse] add merger/topo sort support for slice-based affine sparse index codegen

[Peiming]

rebase

[Peiming]

rebase

[Peiming]

rebase.

[Peiming]

fix rebase errors.

[Peiming]

cleanup

[aartbik]

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

550

I think you will need to expand this comment with more details on what an "dependent" slice is exactly.
Also, I realize this is just bookkeeping for now, so is merger the right place (do you plan to do more with this?)

[Peiming]

rebase.

[Peiming]

rebase.

[Peiming]

rebase

[Peiming]

rebase + fix variable name

[Peiming]

add more comments.

[wrengr]

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

444

Please use the TensorId and LoopId types. Although they're just typedefs for now, by the end of the week I'll have a CL that changes them to be actual types.

444

Elsewhere in this file we use the name "i" for loop identifiers, so please stick with that for consistency.

446

Please use the Level type. Again, although it's just a typedef for now, I'll be changing that soon.

447

What does "idx" mean here? I went through a lot of work cleaning up all the ambiguity about several different meanings of "index", so please use a non-ambiguous name that matches how things are named elsewhere

459

should be "i"

466

should be "i"

472

elsewhere in this file we spell out "lvl" in order to avoid any potential confusion with loop identifiers, so you should stick with that here.

477

ditto

483

Based on the body of this method, that should actually be a TensorLoopId

555

Please clarify what exactly this type is supposed to be.

555

I've avoided using the name "ldx" in this file because it's too confusing. Please change this to "loop" or similar, to match the names used by the other fields

560

please rename to "lvlToDependantLoop" or similar. (We never use "tlvl" anywhere. "depended" is misspelled. And I've intentionally avoided using "ldx" elsewhere in this file)

560

Please clarify exactly what this type is supposed to be

[wrengr]

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

764–765

please clarify exactly what all these unsigned are actually supposed to be.

792

Level

794

LoopId

816

LoopId

862

Please clarify exactly what this type is supposed to be. Is it Dimension or something else?

870

ditto

874

ditto

889

ditto

916

ditto

921

ditto

1864

ditto

1865

ditto

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

224

ditto

226

ditto

227

ditto

[aartbik]

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

470

list of loop indices

(you use plural when saying something like set of XXXs)

551

set of (tensor, level) pairs.

552

expressions

554

can you break this like

... that

i and j are used ...

556

dependent

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

113

this comment does not match the implementation
this is just a collector, right, adding it to the dims vector

what you describe seems to match intended use elsewhere?

265

Can you give a small example of that goes in and out here; just giving a basic example will go a long way in helping the reader

824

"relax" is a bit easier to read (I think), same for function name
tryRelaxAffineConstraints

860

\in just in (LaTeX ?)

933

For this revision, can you not change the comment of this method

/ Computes a topologically sorted iteration graph for the linalg operation.
/ Ensures all tensors are visited in natural coordinate order. This is
/// essential for sparse storage formats since these only support access

so that the DIFF will be easier to line up the method with the original code?

[Peiming]

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

447

it means tensor expression like T_{i+j}, that uses a non-trivial expression to index the tensor.

Changed it to isLvlWithNonTrivialIdxExp to be more complete.

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

860

Yeah

[Peiming]

address comments.

[Peiming]

fix error.

[wrengr]

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

444

If you're factoring this out, then you might as well factor out the getLvl(b) and getDimLevelType(b) expressions too...

Originally I was going to call that a minor nit, but on reflection I've noticed that the current implementation allows the first callback to mutate those fields (via setLevelAndType) before the second callback. That's an extremely bug-prone design, so I can't imagine that it's intentional. But if for some reason it is intentional, then it needs to be explicitly documented that this is expected behavior and why it's intended.

446

"level"

448

"level"

458

This should also be "i"

461

Please rename this to loopToDependencies. Again, I do not use "ldx" anywhere in this file.

462

Please rename this to levelToDependentLoops. Again, I've done a lot of work to avoid the ambiguous use of "idx".

465

"loop"

465

ditto my other comment about using "///"

465

Should add a verb here (e.g., "Checks whether").

470

"which appear"

470

Although this matches the name of the getMatchingIndexingMap method, I think it would be a lot clearer to call these "subscript expressions".

470–471

ditto my other comment about using "///"

472

Should be "dependent", here and everywhere else in this file.

476

"lvl"

476

When providing documentation this must be "///" so that it's picked up by the program that generates the documentation.

477

This doesn't match the documentation. If the documentation is correct then this should be LoopId i. Otherwise, the documentation needs updating.

477

"Lvl"

481

This is inconsistent with the rest of the file. A LatPoint contains a whole set of TensorLoopIds.

(If what is currently called TensorLoopId is actually supposed to be a "lattice point", then we need to rename the LatPoint and LatPointId types to something else to avoid confusion. In addition to updating all the documentation, method names, etc.)

481

Ditto, although this matches the name of the getMatchingIndexingMap method, it would be clearer to call these "subscript expressions" throughout this file

483

Should rename this to hasDependencies.

(Because: this has nothing to do with "levels" so far as I can see. Calling it hasDependencies matches the name of the field it's querying. And it's always the case that "dependencies" correspond with "non-trivial subscripting expressions", so there's no need to spell out the latter in the method name.)

550

"identifier"

551

"identifier"

552

"subscript"

554

"subscript"

557

"identifier"

559

"its"

562

"identifier"

562

"lvl"

[Peiming]

fix bug

[Peiming]

fix issue

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

444

OMG, you are SO right, i forget the else!

[Peiming]

address comments.

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

470

I will keep it as TACO paper uses the term index variable.

483

I will keep it for now, as it matches the technique report that I am writing, we can revisit the term later ;-)

[Peiming]

fix names.

[aartbik]

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

438–439

please document the parameters of the expected callback (since the bool is no longer "trivial")

452

in the description of the callback above, use names, and then say

/*name=*/false here and above for true, so that it becomes more readable

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

102

hmm, I actually liked that I was able to hide the topsort completely in this new datastructure, but now it is leaking again
(actually, the slice was already leaking)

but we have plans for an iteration graph class, then it would make sense perhaps to hide it there..

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

277

also, not now but perhaps later, if feels like all the affine analysis could also be in its own class, or perhaps with a new iteration graph class

399

this method also does a lot more analysis than before, perhaps update the method doc for now a bit to reflect this
(having a new affine analyzer could be next)

692

typo

1809

yeah, avoiding too many special cases, but making the general case and special case more uniform would be nice in the longer run

[Peiming]

address comments.

[Peiming]

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

102

I think I can hide it again in the next commit, because I move some code there into CodegenEnv.

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

277

Yeah, that would be better.

[Peiming]

rebase

[aartbik]

Wait to see if Wren is also okay with how her comments were addressed.

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

assert(t < numTensors && i < numLoops && lvl < lvlToLoop[t].size() &&
            isValidDLT(dlt));
     lvlTypes[t][i] = dlt;
     loopToLvl[t][i] = lvl;
     lvlToLoop[t][lvl] = i;
   }
 
   /// Iterates over a set of `TensorLoopId`s, invoking the callback
   /// for each `TensorLoopId` and passing it the corresponding tensor
   /// identifier, level, and level-type.
-  void
-  foreachTensorLoopId(LatPointId p,
-                      function_ref<void(TensorLoopId, TensorId,
-                                        std::optional<Level>, DimLevelType)>
-                          callback) const {
-    for (const TensorLoopId b : latPoints[p].bits.set_bits())
-      callback(b, tensor(b), getLvl(b), getDimLevelType(b));
+  void foreachTensorLoopId(
+      LatPointId p, function_ref<void(TensorLoopId, TensorId,
+                                      std::optional<Level>, DimLevelType, bool)>
+                        callback) {
+    for (const TensorLoopId b : latPoints[p].bits.set_bits()) {
+      TensorId t = tensor(b);
+      if (isLvlWithNonTrivialIdxExp(b)) {
+        // This must be an undefined dim.
+        assert(!getLvl(b).has_value());
+        // Slice the tid along the dependent dim to iterate current loop.
+        callback(b, t, ldxToDependencies[loop(b)][t], getDimLevelType(b), true);
+      }
+      callback(b, t, getLvl(b), getDimLevelType(b), false);
+    }
   }
 
   /// Sets whether the output tensor is sparse or not.
   void setHasSparseOut(bool s) { hasSparseOut = s; }
 
+  /// Establishes the two-way map that l <-> <t, lvl>.
+  void setLoopDependentTensorLevel(LoopId i, TensorId t, Level lvl) {
+    assert(lvl < numLoops);
+    ldxToDependencies[i][t] = lvl;
+    lvlToDependentIdx[t][lvl].push_back(i);
+  }
+
+  // Whether the ldx has dependent slice.
+  bool hasDependentLvl(LoopId i, TensorId tid) {
+    return ldxToDependencies[i][tid].has_value();
+  }
+
+  // Returns the list of loop indices appeared in the non-trivial indexing
+  // expression on t_l, e.g., A[i+j] => {i, j}
+  std::vector<LoopId> &getDependedLoops(TensorId t, Level lvl) {
+    return lvlToDependentIdx[t][lvl];
+  }
+
+  // Return the defining [tid, dim] for the loop.
+  std::pair<TensorId, Level> getLoopDefiningDim(Level lvl) const {
+    return loopBounds[lvl];
+  }
+
+  /// Whether the lattice point represents a tensor level with non-trivial index
+  /// expression on it.
+  bool isLvlWithNonTrivialIdxExp(TensorLoopId b) const {
+    return ldxToDependencies[loop(b)][tensor(b)].has_value();
+  }
+
   /// Convenience getters to immediately access the stored nodes.
   /// Typically it is inadvisible to keep the reference around, as in
   /// `TensorExpr &te = merger.exp(e)`, since insertions into the merger
   /// may cause data movement and invalidate the underlying memory address.
   TensorExp &exp(ExprId e) { return tensorExps[e]; }
   LatPoint &lat(LatPointId p) { return latPoints[p]; }
   SmallVector<LatPointId> &set(LatSetId s) { return latSets[s]; }
 
private:
 
   // Map that converts pair<TensorId, LoopId> to the corresponding
   // level.
   std::vector<std::vector<std::optional<Level>>> loopToLvl;
 
   // Map that converts pair<TensorId, Level> to the corresponding LoopId.
   std::vector<std::vector<std::optional<LoopId>>> lvlToLoop;
 
+  // Map from a loop idx to its dependencies if any.
+  // The dependencies of a loop idx 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} to indicate that
+  // i and j are used in the non-trivial index expression on A0.
+  std::vector<std::vector<std::optional<Level>>> ldxToDependencies;
+  // The inverse map of ldxToDependencies from tensor level -> dependent loop
+  // index.
+  // E.g., A[i+j], we have A0 => {i, j}, to indicate that A0 uses both {i, j}
+  // to compute it indices.
+  std::vector<std::vector<std::vector<LoopId>>> lvlToDependentIdx;
+
+  // Map from loop index to the [tid, dim] 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

public:
   LoopOrd topSortSize() const { return topSort.size(); }
   LoopId topSortAt(LoopOrd n) const { return topSort.at(n); }
   void topSortPushBack(LoopId i) { topSort.push_back(i); }
   void topSortClear(size_t capacity = 0) {
     topSort.clear();
     topSort.reserve(capacity);
   }
 
+  ArrayRef<LoopId> getTopSort() const { return topSort; };
   ArrayRef<LoopId> getTopSortSlice(LoopOrd n, LoopOrd m) const;
   ArrayRef<LoopId> getLoopStackUpTo(LoopOrd n) const;
   ArrayRef<LoopId> getCurrentLoopStack() const;
   /// Returns the induction-variable for the loop identified by the given
   /// `LoopId`.  This method handles application of the topological sort
   /// in order to convert the `LoopId` into the corresponding `LoopOrd`.
   Value getLoopVar(LoopId i) const;
 

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

private:
   /// doesn't have a default ctor.
   AffineExpr pickedDim;
   /// The iterator type that we want.
   utils::IteratorType pickIterType;
   /// The mapping between dim=>iterator type.
   SmallVector<utils::IteratorType> iterTypes;
 };
 
+// Flattens an affine expression into a list of AffineDimExprs.
+struct AffineDimCollector : public AffineExprVisitor<AffineDimFinder> {
+  void visitDimExpr(AffineDimExpr expr) { dims.push_back(expr); }
+  SmallVector<AffineDimExpr> dims;
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler analysis methods.
 //===----------------------------------------------------------------------===//
 
 // TODO: the "idx"-vs-"ldx" naming convention is not self-explanatory,
 // and those letters are too easy to confuse visually.  We should switch
case AffineExprKind::Constant: {
     // Falls through when it is a constant Affine
     return true;
   }
   default:
     return false;
   }
 }
 
+/// Helper method to inspect affine expressions for index variable reduction
+/// based codegen. It finds the dependent index set for all tensor levels in the
+/// current expression we are generating.
+///
+/// For example, when handling A[i+j][j+k], we build the two way mapping in
+/// merger between (tensor, level) pairs and their dependent index variable set:
+/// A_0 <=> [i, j] and A_1 <=> [j, k]
+///
+/// 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) {
+  switch (a.getKind()) {
+  case AffineExprKind::DimId: {
+    LoopId ldx = a.cast<AffineDimExpr>().getPosition();
+    if (!isUndefDLT(merger.getDimLevelType(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)
+      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);
+    }
+    return true;
+  }
+  case AffineExprKind::Constant:
+  case AffineExprKind::Mul:
+    // TODO: Support Mul and Constant AffineExp for slice-based codegen
+    return false;
+  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;
+  }
+}
+
 /// Get the total number of compound affine expressions in the
 /// `getMatchingIndexingMap` for the given tensor.  For the following inputs:
 ///
 /// map = (d0, d1, d2) => (d0 + d1, d2)
 /// lvlTypes = ["compressed", "compressed"]
 ///
 /// Returns 1 (because the first level is compressed and its corresponding
 /// indexing-expression is `d0 + d1`)
-static unsigned getNumCompoundAffineOnSparseLvls(AffineMap map, Value tensor) {
+static unsigned getNumNonTrivialIdxExpOnSparseLvls(AffineMap map,
+                                                   Value tensor) {
   // The `tensor` is not guaranted to have `RankedTensorType`, therefore
   // we can't use `getRankedTensorType`/`getSparseTensorType` here.
   // However, we don't need to handle `StorageSpecifierType`, so we
   // can use `SparseTensorType` once we guard against non-tensors.
   const auto rtp = tensor.getType().dyn_cast<RankedTensorType>();
   if (!rtp)
     return 0;
   const SparseTensorType stt(rtp);
for (Level l = 0; l < lvlRank; l++) {
     if (!exprs[d].isa<AffineDimExpr>() && !stt.isDenseLvl(l))
       num++;
   }
   return num;
 }
 
 /// Get the total number of sparse levels with compound affine
 /// expressions, summed over all operands of the `GenericOp`.
-static unsigned getNumCompoundAffineOnSparseLvls(linalg::GenericOp op) {
+static unsigned getNumNonTrivialIdxExpOnSparseLvls(linalg::GenericOp op) {
   unsigned num = 0;
   for (OpOperand &t : op->getOpOperands())
-    num += getNumCompoundAffineOnSparseLvls(op.getMatchingIndexingMap(&t),
+    num += getNumNonTrivialIdxExpOnSparseLvls(op.getMatchingIndexingMap(&t),
                                               t.get());
   return num;
 }
 
-static bool hasCompoundAffineOnSparseOut(linalg::GenericOp op) {
+static bool hasNonTrivialAffineOnSparseOut(linalg::GenericOp op) {
   OpOperand *out = op.getDpsInitOperand(0);
   if (getSparseTensorType(out->get()).isAllDense())
     return false;
-  return getNumCompoundAffineOnSparseLvls(op.getMatchingIndexingMap(out),
+  return getNumNonTrivialIdxExpOnSparseLvls(op.getMatchingIndexingMap(out),
                                             out->get());
 }
 
 /// Helper method to inspect sparse encodings in the tensor types.
 /// Fills the per-dimension sparsity information for all tensors.
 /// Returns true if the sparse annotations and affine subscript
 /// expressions of all tensors are admissible. Returns false if
 /// no annotations are found or inadmissible constructs occur.
-static bool findSparseAnnotations(CodegenEnv &env) {
+static bool findSparseAnnotations(CodegenEnv &env, bool idxReducBased) {
   bool annotated = false;
   // `filterLdx` may be mutated by `findAffine`.
   LoopId filterLdx = env.merger().getStartingFilterLoopId();
   for (OpOperand &t : env.op()->getOpOperands()) {
     const auto map = env.op().getMatchingIndexingMap(&t);
     const auto enc = getSparseTensorEncoding(t.get().getType());
     if (enc)
       annotated = true;
+
     const Level lvlRank = map.getNumResults();
     assert(!enc || lvlRank == enc.getLvlRank());
     assert(static_cast<Level>(env.op().getRank(&t)) == lvlRank);
+
+    // We only need to do index reduction if there is at least one non-trivial
+    // index expression on sparse levels.
+    // If all non-trivial index expression is on dense levels, we can
+    // efficiently rely on the random access to locate the element.
+    bool needIdxReduc =
+        enc && getNumNonTrivialIdxExpOnSparseLvls(map, t.get()) != 0;
+    // If then current tensor being inspected requires affine index, it need
+    // to be sliced.
     for (Level l = 0; l < lvlRank; l++) {
       const TensorId tid = t.getOperandNumber();
-      // FIXME: `toOrigDim` is deprecated.
-      // FIXME: above we asserted that there are `lvlRank` many results,
-      // but this is assuming there are in fact `dimRank` many results instead.
-      const AffineExpr a = map.getResult(toOrigDim(enc, l));
-      if (!findAffine(env.merger(), tid, l, a, enc.getLvlType(l), filterLdx))
+      AffineExpr a = map.getResult(toOrigDim(enc, l));
+      DimLevelType dlt = enc.getLvlType(l);
+      if (idxReducBased && needIdxReduc) {
+        if (!findDepIdxSet(env.merger(), tid, l, a, dlt))
           return false; // inadmissible affine expression
+      } else {
+        if (!findAffine(env.merger(), tid, l, a, dlt, filterLdx))
+          return false; // inadmissible affine expression
+      }
     }
   }
   assert(filterLdx == env.merger().getNumLoops());
   return annotated;
 }
 
 /// A helper to compute a topological sort. O(n^2) time complexity
 /// as we use adj matrix for the graph.
case AffineExprKind::Mul: {
     }
     break;
   }
   default:
     break;
   }
 }
 
-static void tryLoosenAffineDenseConstraints(linalg::GenericOp op,
-                                            std::optional<LoopId> &fldx,
-                                            AffineExpr &fa,
-                                            std::optional<LoopId> &tldx,
-                                            AffineExpr &ta) {
+static void tryRelaxAffineConstraints(linalg::GenericOp op,
   // We use a heuristic here to only pick one dim expression from each
   // compound affine expression to establish the order between two dense
   // dimensions.
   if (!tldx) {
     AffineDimFinder finder(op);
     // NOTE: The ordering can only be loosen when the destination level is
     // dense (when !tldx), for [dense, sparse] -> (d0 + d1, d2), we still
     // require both d0 < d2 and d1 < d2 to ensure correct ordering (i.e.,
     // no ordering like d0->d2->d1).
     // TODO: this is obviously a sub optimal solution.
     if (!fldx && !fa.isa<AffineConstantExpr>()) {
       // Heuristic: we prefer parallel loop for lhs to reduce the chance
       // we add reduce < parallel ordering.
       finder.setPickedIterType(utils::IteratorType::parallel);
       finder.walkPostOrder(fa);
       fa = finder.getDimExpr();
       fldx = finder.getDimExpr().getPosition();
     }
     if (!ta.isa<AffineConstantExpr>()) {
       // Heuristic: we prefer reduction loop for rhs to reduce the chance
-      // addint reduce < parallel ordering.
+      // adding reduce < parallel ordering.
       finder.setPickedIterType(utils::IteratorType::reduction);
       finder.walkPostOrder(ta);
       ta = finder.getDimExpr();
       tldx = finder.getDimExpr().getPosition();
     }
   }
 }
 
-/// Computes a topologically sorted iteration graph for the linalg operation.
-/// Ensures all tensors are visited in natural coordinate order.  This is
-/// essential for sparse storage formats since these only support access
-/// along fixed levels.  Even for dense storage formats, however, the natural
-/// coordinate order yields innermost unit-stride access with better spatial
-/// locality.
-static bool computeIterationGraph(CodegenEnv &env, SortMask mask,
-                                  OpOperand *skip = nullptr) {
-  // Set up an n x n from/to adjacency matrix of the iteration graph
-  // for the implicit loop indices i_0 .. i_n-1.
-  const LoopId n = env.merger().getNumLoops();
-  std::vector<std::vector<bool>> adjM(n, std::vector<bool>(n, false));
-  std::vector<unsigned> inDegree(n, 0); // in-degree of each node.
-  const auto iteratorTypes = env.op().getIteratorTypesArray();
-  // Iterate over the indexing maps of every tensor in the tensor expression.
-  for (OpOperand &t : env.op()->getOpOperands()) {
-    // Get map and encoding.
-    const auto map = env.op().getMatchingIndexingMap(&t);
-    const auto enc = getSparseTensorEncoding(t.get().getType());
-    assert(map.getNumDims() + getNumCompoundAffineOnSparseLvls(env.op()) == n);
-
-    // Skips dense inputs/outputs when not requested.
-    const bool isDenseInput = !enc && env.op().isDpsInput(&t);
-    const bool isDenseOutput = !enc && !isDenseInput;
-    if ((isDenseInput && !includesDenseInput(mask)) ||
-        (isDenseOutput && !includesDenseOutput(mask)))
-      continue;
-
-    // Push unrelated loops into sparse iteration space, so these
-    // will be skipped more often.
-    // TODO: Do we really need this?
-    if (includesUndef(mask)) {
-      const TensorId tensor = t.getOperandNumber();
-      for (LoopId i = 0; i < n; i++) {
-        const auto dltI = env.dlt(tensor, i);
-        if (isCompressedDLT(dltI) || isSingletonDLT(dltI)) {
-          for (LoopId j = 0; j < n; j++)
-            if (isUndefDLT(env.dlt(tensor, j))) {
-              adjM[i][j] = true;
-              inDegree[j]++;
-            }
-        } else {
-          assert(isDenseDLT(dltI) || isUndefDLT(dltI));
-        }
-      }
-    }
+/// Makes target array's elements appear in the same order as the `order` array.
+static void sortArrayBasedOnOrder(std::vector<LoopId> &target,
+                                  ArrayRef<LoopId> order) {
+  std::sort(target.begin(), target.end(), [&order](LoopId l, LoopId r) {
+    assert(l != r);
+    int idxL = -1, idxR = -1;
+    for (int i = 0, e = order.size(); i < e; i++) {
+      if (order[i] == l)
+        idxL = i;
+      if (order[i] == r)
+        idxR = i;
+    assert(idxL >= 0 && idxR >= 0);
+    return idxL < idxR;
+  });
 
-    // Each tensor expression and optional dimension ordering (row-major
-    // by default) puts an ordering constraint on the loop indices. For
-    // example, the tensor expresion A_ijk forces the ordering i < j < k
-    // on the loop indices if no explicit dimension ordering is given.
-    const Level lvlRank = map.getNumResults();
-    assert(!enc || lvlRank == enc.getLvlRank());
-    for (Level l = 0; l < lvlRank; l++) {
-      // FIXME: `toOrigDim` is deprecated.
-      // FIXME: above we asserted that there are `lvlRank` many results,
-      // but this is assuming there are in fact `dimRank` many results instead.
-      AffineExpr ta = map.getResult(toOrigDim(enc, l));
-      std::optional<LoopId> tldx =
-          env.merger().getLoopId(t.getOperandNumber(), l);
-
-      // Filter loops should be constructed after all the dependent loops,
-      // i.e., d0 + d1 < filter_loop(d0 + d1)
-      if (tldx && env.merger().isFilterLoop(*tldx)) {
-        assert(!ta.isa<AffineDimExpr>() && !isDenseDLT(enc.getLvlType(l)));
-        addAffineOrderings(adjM, inDegree, ta, AffineExpr(), std::nullopt,
-                           tldx);
-        // Now that the ordering of affine expression is captured by filter
-        // loop idx, we only need to ensure the affine ordering against filter
-        // loop. Thus, we reset the affine express to nil here to mark it as
-        // resolved.
-        ta = AffineExpr();
-      }
+static void addFilterLoopBasedConstraints(CodegenEnv &env, OpOperand &t,
+                                          OpOperand *skip, SortMask mask,
+                                          std::vector<std::vector<bool>> &adjM,
+                                          std::vector<unsigned> &inDegree) {
+  // Get map and encoding.
+  auto map = env.op().getMatchingIndexingMap(&t);
+  auto enc = getSparseTensorEncoding(t.get().getType());
+
+  for (Level l = 0, rank = map.getNumResults(); l < rank; l++) {
+      enc.isAllDense();
+      assert(!ta.isa<AffineDimExpr>() && !isDenseDLT(enc.getDimLevelType()[l]));
+      addAffineOrderings(adjM, inDegree, ta, AffineExpr(), std::nullopt, tldx);
 
-      // Skip tensor during cycle resolution, though order between filter loop
-      // and dependent loops need to be guaranteed unconditionally.
-      if (&t == skip)
-        continue;
 
-      if (l > 0) {
-        // FIXME: `toOrigDim` is deprecated.
-        // FIXME: above we asserted that there are `lvlRank` many results,
-        // but this is assuming there are in fact `dimRank` many results.
-        AffineExpr fa = map.getResult(toOrigDim(enc, l - 1));
-        std::optional<LoopId> fldx =
-            env.merger().getLoopId(t.getOperandNumber(), l - 1);
 
-        // Applying order constraints on every pair of dimExpr between two
-        // compound affine expressions can sometime too strict:
-        // E.g, for [dense, dense] -> (d0 + d1, d2 + d3).
-        // It is totally fine to have loop sequence d0->d2->d1->d3 instead of
-        // requiring d0 < d2, d1 < d2, d0 < d3, d1 < d3.
-        if (!includesDense(mask))
-          tryLoosenAffineDenseConstraints(env.op(), fldx, fa, tldx, ta);
+      // We also relax the affine constraint when use slice-based algorithm
+      // as there is no filter loop for affine index on sparse dimension.
+      // TODO: do we really need the condition?
+        tryRelaxAffineConstraints(env.op(), fldx, fa, tldx, ta);
 
-        // (d0 + d1) < (d2 + d3), or
-        // filter_loop_d-1 < (d2 + d3), or
-        // (d0 + d1) < filter_loop_d, or
-        // filter_loop_d-1 < filter_loop_d depending on whether fa/ta is reset
-        // above.
-        addAffineOrderings(adjM, inDegree, fa, ta, fldx, tldx);
-      }
-    }
-  }
+
+static void addSliceBasedConstraints(CodegenEnv &env, OpOperand &t,
+                                     OpOperand *skip, SortMask mask,
+                                     std::vector<std::vector<bool>> &adjM,
+                                     std::vector<unsigned> &inDegree) {
+  // Get map and encoding.
+  auto map = env.op().getMatchingIndexingMap(&t);
+  auto enc = getSparseTensorEncoding(t.get().getType());
+
+  // No special treatment for simple indices.
+  if (getNumNonTrivialIdxExpOnSparseLvls(map, t.get()) == 0)
+    return addFilterLoopBasedConstraints(env, t, skip, mask, adjM, inDegree);
+
+  // Skip tensor during cycle resolution, though order between filter loop
+  // and dependent loops need to be guaranteed unconditionally.
+  if (&t == skip)
+    return;
+
+  AffineDimFinder finder(env.op());
+  finder.setPickedIterType(utils::IteratorType::reduction);
+  // To compute iteration graph for tensor[d0 + d1 + d3, d4 + d5 + d6],
+  // we requires there exist d_x \in {d0, d1, d3} and d_y \in {d4, d5, d6},
+  // and d_x > d_y && {d0, d1, d3} - d_x > {d4, d5, d6} - d_y
+  for (Level lvl = 1, rank = map.getNumResults(); lvl < rank; lvl++) {
+    AffineExpr fa = map.getResult(toOrigDim(enc, lvl - 1));
+    AffineExpr ta = map.getResult(toOrigDim(enc, lvl));
+
+    // This is a heurisitic, we pick an abitrary reduction loop from lhs and
+    // rhs and use them as d_x and d_y.
+    finder.walkPostOrder(fa);
+    AffineDimExpr fexp = finder.getDimExpr();
+    LoopId fldx = fexp.getPosition();
+
+    finder.walkPostOrder(ta);
+    AffineDimExpr texp = finder.getDimExpr();
+    LoopId tldx = texp.getPosition();
+
+    // d_x > d_y
+    if (!adjM[fldx][tldx]) {
+      adjM[fldx][tldx] = true;
+      inDegree[tldx]++;
+    }
+
+    AffineDimCollector fCollector;
+    fCollector.walkPostOrder(fa);
+    AffineDimCollector tCollector;
+    tCollector.walkPostOrder(ta);
+
+    // make sure dx and dy is the last;
+    for (auto fd : fCollector.dims) {
+      LoopId f = fd.getPosition();
+      if (f == fldx)
+        continue;
+      if (!adjM[f][fldx]) {
+        adjM[f][fldx] = true;
+        inDegree[fldx]++;
+      }
+    }
+    for (auto td : tCollector.dims) {
+      LoopId t = td.getPosition();
+      if (t == tldx)
+        continue;
+      if (!adjM[t][tldx]) {
+        adjM[t][tldx] = true;
+        inDegree[tldx]++;
+      }
+    }
+    // Since we only support affine addition, the order between two dim
+    // expression does not really matters.
+    // {d0, d1, d3} - d_x > {d4, d5, d6} - d_y
+    // This is to ensure that the affine expressions are reduced in sparse
+    // tensor level ordering.
+    // TODO: this ordering could probably be loosen if we support out-of-order
+    // reduction.
+    // TODO: the evaluation order need to be ensure to
+    // support affine multiplication.
+    for (auto fd : fCollector.dims) {
+      LoopId f = fd.getPosition();
+      if (f == fldx) // skip d_x
+        continue;
+
+      for (auto td : tCollector.dims) {
+        LoopId t = td.getPosition();
+        if (t == tldx) // skip d_y
+          continue;
+        if (!adjM[f][t]) {
+          adjM[f][t] = true;
+          inDegree[t]++;
+        }
+      }
+    }
+  }
+}
+
+/// Computes a topologically sorted iteration graph for the linalg operation.
+/// Ensures all tensors are visited in natural index order. This is
+/// essential for sparse storage formats since these only support access
+/// along fixed dimensions. Even for dense storage formats, however, the natural
+/// index order yields innermost unit-stride access with better spatial
+/// locality.
+static bool computeIterationGraph(CodegenEnv &env, SortMask mask,
+                                  OpOperand *skip, bool useSlice = false) {
+  // Set up an n x n from/to adjacency matrix of the iteration graph
+  // for the implicit loop indices i_0 .. i_n-1.
+  const LoopId n = env.merger().getNumLoops();
+  std::vector<std::vector<bool>> adjM(n, std::vector<bool>(n, false));
+  std::vector<unsigned> inDegree(n, 0); // in-degree of each node.
+  const auto iteratorTypes = env.op().getIteratorTypesArray();
+  // Iterate over the indexing maps of every tensor in the tensor expression.
+  for (OpOperand &t : env.op()->getOpOperands()) {
+    // Get map and encoding.
+    const auto map = env.op().getMatchingIndexingMap(&t);
+    const auto enc = getSparseTensorEncoding(t.get().getType());
+    assert(map.getNumDims() + getNumNonTrivialIdxExpOnSparseLvls(env.op()) ==
+           n);
+
+    // Skips dense inputs/outputs when not requested.
+    const bool isDenseInput = !enc && env.op().isDpsInput(&t);
+    const bool isDenseOutput = !enc && !isDenseInput;
+    if ((isDenseInput && !includesDenseInput(mask)) ||
+        (isDenseOutput && !includesDenseOutput(mask)))
+      continue;
+
+    // Push unrelated loops into sparse iteration space, so these
+    // will be skipped more often.
+    // TODO: Do we really need this?
+    if (includesUndef(mask)) {
+      const TensorId tensor = t.getOperandNumber();
+      for (LoopId i = 0; i < n; i++) {
+        const auto dltI = env.dlt(tensor, i);
+        if (isCompressedDLT(dltI) || isSingletonDLT(dltI)) {
+          for (LoopId j = 0; j < n; j++)
+            if (isUndefDLT(env.dlt(tensor, j))) {
+              adjM[i][j] = true;
+              inDegree[j]++;
+            }
+        } else {
+          assert(isDenseDLT(dltI) || isUndefDLT(dltI));
+        }
+      }
+    }
+    // Push unrelated loops into sparse iteration space, so these
+    // will be skipped more often.
+    if (useSlice)
+      addSliceBasedConstraints(env, t, skip, mask, adjM, inDegree);
+    else
+      addFilterLoopBasedConstraints(env, t, skip, mask, adjM, inDegree);
+  }
   // Topologically sort the iteration graph to determine loop order.
   // Report failure for a cyclic iteration graph.
   env.topSortClear(n);
   return topSortOptimal(env, n, iteratorTypes, inDegree, adjM);
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler synthesis methods (statements and expressions).
 //===----------------------------------------------------------------------===//
 
 /// Local bufferization of all dense and sparse data structures.
 static void genBuffers(CodegenEnv &env, OpBuilder &builder) {
   linalg::GenericOp op = env.op();
   Location loc = op.getLoc();
   assert(op.getNumOperands() == op.getNumDpsInputs() + 1);
 
   env.emitter().initializeLoopEmit(
       builder, loc,
       /// Generates buffer for the output tensor.
       /// Note that all sparse kernels assume that when all elements are written
env.emitter().initializeLoopEmit(
         // TODO: use better analysis to avoid zeroing out the buffer?
         bool isInit = op.isInitTensor(lhs);
         Value init = memref;
         if (!isInit) {
           Value zero = constantZero(builder, loc,
                                     getElementTypeOrSelf(tensor.getType()));
           builder.create<linalg::FillOp>(loc, ValueRange{zero},
                                          ValueRange{init});
         }
         return init;
       });
 }
 
 /// Generates index for load/store on sparse tensor.
 // FIXME: It's not entirely clear what "index" means here (i.e., is it
 // a "coordinate", or "Ldx", or what).  So the function should be renamed
 // and/or the documentation expanded in order to clarify.
 static Value genIndex(CodegenEnv &env, OpOperand *t) {
   auto map = env.op().getMatchingIndexingMap(t);
   const auto stt = getSparseTensorType(t->get());
   const Level lvlRank = stt.getLvlRank();
   assert(static_cast<Level>(map.getNumResults()) == lvlRank);
   // FIXME: `toOrigDim` is deprecated.
   // FIXME: above we asserted that there are `lvlRank` many results,
   // but this is assuming there are in fact `dimRank` many results instead.
   AffineExpr a = map.getResult(toOrigDim(stt.getEncoding(), lvlRank - 1));
   assert(a.getKind() == AffineExprKind::DimId);
   const LoopId idx = a.cast<AffineDimExpr>().getPosition();
   return env.getLoopVar(idx);
 }
 
 /// Generates subscript for load/store on a dense or sparse tensor.
 static Value genSubscript(CodegenEnv &env, OpBuilder &builder, OpOperand *t,
                           SmallVectorImpl<Value> &args) {
   const Location loc = env.op().getLoc();
   const TensorId tid = t->getOperandNumber();
   const auto map = env.op().getMatchingIndexingMap(t);
   const auto stt = getSparseTensorType(t->get());
   if (stt.hasEncoding()) {
     // For sparse tensors we only push the last-level's position onto `args`.
     const auto pos = env.emitter().getPosits()[tid].back();
     assert(pos);
     args.push_back(pos);
   } else {
     // For dense tensors we push all level's coordinates onto `args`.
     const Level lvlRank = stt.getLvlRank();

 /// Generates insertion code to implement dynamic tensor load.
 static Value genInsertionLoad(CodegenEnv &env, OpBuilder &builder,
                               OpOperand *t) {
   linalg::GenericOp op = env.op();
   Location loc = op.getLoc();
   // Direct lexicographic coordinate order, tensor loads as zero.
   if (!env.isExpand()) {
     Type tp = getElementTypeOrSelf(t->get().getType());
     return constantZero(builder, loc, tp);
   }
   // Load from expanded access pattern.
   Value index = genIndex(env, t);
   return builder.create<memref::LoadOp>(loc, env.getExpandValues(), index);
 }
 
 /// Generates insertion code to implement dynamic tensor load for reduction.
 static Value genInsertionLoadReduce(CodegenEnv &env, OpBuilder &builder,
                                     OpOperand *t) {
   linalg::GenericOp op = env.op();
   Location loc = op.getLoc();
   Value identity = env.getCustomRedId();
   // Direct lexicographic coordinate order, tensor loads as identity.
   if (!env.isExpand())
     return identity;
   // Load from expanded access pattern if filled, identity otherwise.
   Value values = env.getExpandValues();
   Value filled = env.getExpandFilled();
   Value index = genIndex(env, t);
   Value isFilled = builder.create<memref::LoadOp>(loc, filled, index);
   Value valAtIndex = builder.create<memref::LoadOp>(loc, values, index);
   return builder.create<arith::SelectOp>(loc, isFilled, valAtIndex, identity);
 }
 
 /// Generates insertion code to implement dynamic tensor store.
static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp,
   // Emit further intitialization at this loop sequence level.
   const LatPointId l0 = env.set(lts)[0];
   bool needsUniv = false;
 
   SmallVector<TensorId> tids;
   SmallVector<Level> lvls;
   env.merger().foreachTensorLoopId(l0, [&](TensorLoopId b, TensorId tid,
                                            std::optional<Level> lvl,
-                                           DimLevelType dlt) {
+                                           DimLevelType dlt, bool) {
     assert(env.merger().loop(b) == idx);
     if (isDenseDLT(dlt) || isUndefDLT(dlt)) {
       needsUniv = true;
     } else {
       // sparse/singleton levels.
       tids.push_back(tid);
       lvls.push_back(*lvl);
     }
static bool translateBitsToTidLvlPairs(
   const BitVector &simple = env.lat(li).simple;
   const TensorId outTid = env.merger().getOutTensorID();
   const std::optional<Level> outLvl = env.merger().getLvl(outTid, ldx);
 
   unsigned numloopCond = 0;
   bool hasNonUnique = false;
   env.merger().foreachTensorLoopId(li, [&, ldx](TensorLoopId b, TensorId tid,
                                                 std::optional<Level> lvl,
-                                                DimLevelType dlt) {
+                                                DimLevelType dlt, bool) {
     if (simple.test(b)) {
       if (isUndefDLT(dlt)) {
         // An undefined dlt in the lattices, we probably mean to
         // iterate based on the level of output tensor.  E.g., this
         // could be a synthetic tensor (for invariants and sparse
         // output tensor).
         // out[i][j] = invariant; or a broadcast
         // out[i][j] = in[i] (j is undef for input)

 public:
   GenericOpSparsifier(MLIRContext *context, SparsificationOptions o)
       : OpRewritePattern<linalg::GenericOp>(context), options(o) {}
 
   LogicalResult matchAndRewrite(linalg::GenericOp op,
                                 PatternRewriter &rewriter) const override {
     // Only accept single output operations without affine index on sparse
     // output.
-    if (op.getNumDpsInits() != 1 || hasCompoundAffineOnSparseOut(op))
+    if (op.getNumDpsInits() != 1 || hasNonTrivialAffineOnSparseOut(op))
       return failure();
 
-    if (options.enableIndexReduction)
-      llvm_unreachable("not yet implemented");
-
     // Sets up a code generation environment.
     const unsigned numTensors = op->getNumOperands();
     const unsigned numLoops = op.getNumLoops();
-    const unsigned numFilterLoops = getNumCompoundAffineOnSparseLvls(op);
-    CodegenEnv env(op, options, numTensors, numLoops, numFilterLoops);
+    const unsigned numFilterLoops = getNumNonTrivialIdxExpOnSparseLvls(op);
+    // TODO: we should probably always use slice-based codegen whenever
+    // possible, we can even intermix slice-based and filter-loop based codegen.
+    bool idxReducBased = options.enableIndexReduction && numFilterLoops != 0;
+
+    // If we uses slice based algorithm for affine index, we do not need filter
+    // loop.
+    CodegenEnv env(op, options, numTensors, numLoops,
+                   /*numFilterLoops=*/idxReducBased ? 0 : numFilterLoops);
 
     // Detects sparse annotations and translates the per-level sparsity
     // information for all tensors to loop indices in the kernel.
-    if (!findSparseAnnotations(env))
+    if (!findSparseAnnotations(env, idxReducBased))
       return failure();
 
     // Constructs the tensor expressions tree from `op`, returns failure if the
     // tree can not be built or the tensor expression is inadmissible.
     if (failed(env.initTensorExp()))
       return failure();
 
     // Computes a topologically sorted iteration graph to ensure tensors
     // are visited in natural index order. Gradually relaxes the considered
     // constraints until an acyclic iteration graph results, such that sparse
     // code generation can proceed. As a last resort, an attempt is made
     // to resolve cycles by inserting a conversion.
     bool isAdmissible = false;
     bool hasCycle = true;
 
     // An const list of all masks that we used for interation graph
     // computation. Must be ordered from more strict to less strict.
     // Ideally (though might not be guaranteed), the eariler a constraint mask
     // can be satisfied, the faster the generated kernel will be.
     const auto allMasks = {
         SortMask::kIncludeAll,        SortMask::kIncludeDense,
         SortMask::kIncludeDenseInput, SortMask::kIncludeDenseOutput,
         SortMask::kIncludeUndef,      SortMask::kSparseOnly};
     for (const SortMask mask : allMasks) {
-      if (computeIterationGraph(env, mask)) {
+      if (computeIterationGraph(env, mask, nullptr, idxReducBased)) {
         hasCycle = false;
         if (env.isAdmissibleTopoOrder()) {
           isAdmissible = true;
           break;
         }
         // else try a set of less strict constraints.
       }
     }
-    if (hasCycle)
-      return resolveCycle(env, rewriter); // one last shot
+    if (hasCycle) {
+      return idxReducBased
+                 ? failure() // TODO: should cycle be resolved differently?
+                 : resolveCycle(env, rewriter); // one last shot
+    }
+
     if (!isAdmissible)
       return failure(); // inadmissible expression, reject
 
+    for (OpOperand &t : env.op()->getOpOperands()) {
+      Level rank = env.op().getMatchingIndexingMap(&t).getNumResults();
+      for (Level lvl = 0; lvl < rank; lvl++) {
+        sortArrayBasedOnOrder(
+            env.merger().getDependedLoops(t.getOperandNumber(), lvl),
+            env.getTopSort());
+      }
+    }
+
     // Recursively generates code if admissible.
     env.startEmit();
     genBuffers(env, rewriter);
     genInitConstantDenseAddress(env, rewriter);
     genStmt(env, rewriter, env.getExprId(), 0);
     genResult(env, rewriter);
     return success();
   }

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

: outTensor(numInputOutputTensors - 1),
       syntheticTensor(numInputOutputTensors),
       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>>(numLoops, std::nullopt)) {}
+                std::vector<std::optional<LoopId>>(numLoops, std::nullopt)),
+      ldxToDependencies(numLoops, std::vector<std::optional<Level>>(
+                                      numTensors, std::nullopt)),
+      lvlToDependentIdx(numTensors, std::vector<std::vector<LoopId>>(
+                                        numLoops, std::vector<LoopId>())),
+      loopBounds(numLoops, std::make_pair(numTensors, numLoops)) {}
 
 //===----------------------------------------------------------------------===//
 // Lattice methods.
 //===----------------------------------------------------------------------===//
 
 ExprId Merger::addExp(Kind k, unsigned x, ExprId y, Value v, Operation *op) {
   const ExprId e = tensorExps.size();
   assert((k != kTensor || x < numTensors) && (k != kLoopVar || x < numLoops));

 }
 
 void Merger::dumpBits(const BitVector &bits) const {
   for (TensorLoopId b = 0, be = bits.size(); b < be; b++) {
     if (bits[b]) {
       const TensorId t = tensor(b);
       const LoopId i = loop(b);
       const auto dlt = lvlTypes[t][i];
+      if (isLvlWithNonTrivialIdxExp(b))
+        llvm::dbgs() << " DEP_" << t << "_" << i;
+      else
         llvm::dbgs() << " i_" << t << "_" << i << "_" << toMLIRString(dlt);
     }
   }
 }
 
 #endif // NDEBUG
 
 //===----------------------------------------------------------------------===//
 // Builder methods.