[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

417

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

351

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.

351

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

353

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

354

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

361

should be "i"

368

should be "i"

374

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.

379

ditto

385

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

422

Please clarify what exactly this type is supposed to be.

422

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

427

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)

427

Please clarify exactly what this type is supposed to be

[wrengr]

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

507–689

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

535

Level

537

LoopId

559

LoopId

605

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

613

ditto

617

ditto

632

ditto

659

ditto

664

ditto

1757

ditto

1758

ditto

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

215

ditto

217

ditto

218

ditto

[aartbik]

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

372

list of loop indices

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

418

set of (tensor, level) pairs.

419

expressions

421

can you break this like

... that

i and j are used ...

423

dependent

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

87

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?

274

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

567

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

603

\in just in (LaTeX ?)

676

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

354

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

603

Yeah

[Peiming]

address comments.

[Peiming]

fix error.

[wrengr]

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

351

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.

353

"level"

355

"level"

360

This should also be "i"

363

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

364

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

367

"loop"

367

ditto my other comment about using "///"

367

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

372

"which appear"

372

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

372–373

ditto my other comment about using "///"

374

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

378

"lvl"

378

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

379

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

379

"Lvl"

383

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.)

383

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

385

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.)

417

"identifier"

418

"identifier"

419

"subscript"

421

"subscript"

424

"identifier"

426

"its"

429

"identifier"

429

"lvl"

[Peiming]

fix bug

[Peiming]

fix issue

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

351

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

[Peiming]

address comments.

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

372

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

385

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

325–346

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

359

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

286

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

331

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)

611

typo

1706

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

286

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

public:
   void setDimAndDimLevelType(unsigned t, unsigned i, unsigned dim,
                              DimLevelType dlt) {
     assert(isValidDLT(dlt));
     dimTypes[t][i] = dlt;
     loopIdxToDim[t][i] = dim;
     assert(dim < numLoops);
     dimToLoopIdx[t][dim] = i;
   }
 
+  void setLoopDependentSliceDim(unsigned l, unsigned t, unsigned dim) {
+    assert(dim < numLoops);
+    ldxToDependentSlice[t][l] = dim;
+    sliceToRelatedldx[t][dim].push_back(l);
+  }
+
+  // Whether the ldx has dependent slice.
+  bool hasDependentSliceDim(unsigned l) {
+    return llvm::any_of(ldxToDependentSlice,
+                        [l](const std::vector<std::optional<unsigned>> &m) {
+                          return m[l].has_value();
+                        });
+  }
+
+  // Returns the list of loop index appeared in the affine indexing expression
+  // on t[d].
+  std::vector<unsigned> &getRelatedLoops(unsigned t, unsigned d) {
+    return sliceToRelatedldx[t][d];
+  }
+
   // Iterates the bits of a lattice, for each set bit, converts it into the
   // corresponding tensor dimension and invokes the callback.
   void foreachTidDimPairInBits(
       const BitVector &bits,
       function_ref<void(unsigned b, unsigned tid, std::optional<unsigned> dim,
                         DimLevelType dlt)>
           cb) {
     for (unsigned b : bits.set_bits())
       cb(b, tensor(b), getDimNum(b), getDimLevelType(b));
   }
 
   // Has sparse output tensor setter.
   void setHasSparseOut(bool s) { hasSparseOut = s; }
 
   /// 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(unsigned e) { return tensorExps[e]; }
   LatPoint &lat(unsigned l) { return latPoints[l]; }
   SmallVector<unsigned> &set(unsigned s) { return latSets[s]; }
 
 #ifndef NDEBUG
   /// Print methods (for debugging).
   void dumpExp(unsigned e) const;
   void dumpLat(unsigned p) const;
   void dumpSet(unsigned s) const;
   void dumpBits(const BitVector &bits) const;
 #endif
 
   /// Builds the iteration lattices in a bottom-up traversal given the
   /// remaining tensor (sub)expression and the next loop index in the
   /// iteration graph. Returns index of the root expression.
   unsigned buildLattices(unsigned e, unsigned i);
 
   /// Builds a tensor expression from the given Linalg operation.
   /// Returns index of the root expression on success.
   std::optional<unsigned> buildTensorExpFromLinalg(linalg::GenericOp op);
 
   /// Rebuilds SSA format from a tensor expression.
   Value buildExp(RewriterBase &rewriter, Location loc, unsigned e, Value v0,
                  Value v1);
 
 private:
   /// Private helpers.
   bool maybeZero(unsigned e) const;
   bool isInvariant(unsigned e) const;
   Type inferType(unsigned e, Value src);
private:
 
   // Map that converts pair<tensor id, loop id> to the corresponding
   // dimension.
   std::vector<std::vector<std::optional<unsigned>>> loopIdxToDim;
 
   // Map that converts pair<tensor id, dim> to the corresponding loop id.
   std::vector<std::vector<std::optional<unsigned>>> dimToLoopIdx;
 
+  // Map from loop idx to the dependent slice [tid, dim] pair (if any).
+  std::vector<std::vector<std::optional<unsigned>>> ldxToDependentSlice;
+  // Map from the dependent slice [tid, dim] pair to a list of loop idx.
+  std::vector<std::vector<std::vector<unsigned>>> sliceToRelatedldx;
+
   llvm::SmallVector<TensorExp> tensorExps;
   llvm::SmallVector<LatPoint> latPoints;
   llvm::SmallVector<SmallVector<unsigned>> latSets;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_UTILS_MERGER_H_

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

public:
   size_t topSortSize() const { return topSort.size(); }
   unsigned topSortAt(unsigned i) const { return topSort.at(i); }
   void topSortPushBack(unsigned i) { topSort.push_back(i); }
   void topSortClear(unsigned capacity = 0) {
     topSort.clear();
     topSort.reserve(capacity);
   }
 
+  ArrayRef<unsigned> getTopSort() const { return ArrayRef<unsigned>(topSort); }
   ArrayRef<unsigned> getTopSortSlice(size_t n, size_t m) const;
   ArrayRef<unsigned> getLoopCurStack() const;
   Value getLoopIdxValue(size_t loopIdx) const;
 
   //
   // Sparse tensor output and expansion methods.
   //
 

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

bool CodegenEnv::isAdmissibleTensorExp(unsigned exp) {
 
   // Accept "truly dynamic" if the output tensor materializes uninitialized
   // into the computation and insertions occur in lexicographic index order.
   sparseOut = lhs;
   return isMaterializing(lhs->get());
 }
 
 bool CodegenEnv::isAdmissibleTopoOrder() {
-  if (!hasSparseOutput())
-    return true;
-
-  OpOperand *lhs = linalgOp.getDpsInitOperand(0);
+  if (hasSparseOutput()) {
+    OpOperand *lhs = op().getDpsInitOperand(0);
     // Accept "truly dynamic" if the output tensor materializes uninitialized
     // into the computation and insertions occur in lexicographic index order.
     unsigned nest = 0;
-  auto iteratorTypes = linalgOp.getIteratorTypesArray();
-  for (unsigned i = 0, e = latticeMerger.getNumLoops(); i < e; i++) {
-    if (!latticeMerger.isFilterLoop(topSortAt(i))) {
+    auto iteratorTypes = op().getIteratorTypesArray();
+    for (unsigned i = 0, e = merger().getNumLoops(); i < e; i++) {
+      if (!merger().isFilterLoop(topSortAt(i))) {
         // We only count non-filter loops as filter loops should be considered
         // as a special type of parallel loops.
         if (linalg::isReductionIterator(iteratorTypes[topSortAt(i)]))
           break; // terminate at first reduction
         nest++;
       }
     }
     // Determine admissible dynamic insertion situations:
     // (1) fully injective, since there are no reductions,
     // (2) admissible 1-d expansion in innermost dimension.
-  if (nest >= linalgOp.getRank(lhs) - 1) {
+    if (nest >= op().getRank(lhs) - 1)
       outerParNest = nest;
-    return true;
-  }
+    else
       return false;
   }
 
+  return true;
+}
+
 //===----------------------------------------------------------------------===//
 // Code generation environment topological sort methods
 //===----------------------------------------------------------------------===//
 
 ArrayRef<unsigned> CodegenEnv::getTopSortSlice(size_t n, size_t m) const {
   return ArrayRef<unsigned>(topSort).slice(n, m);
 }
 

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

private:
   /// The picked AffineDimExpr after visit.
   AffineExpr pickedDim;
   /// The iterator type that we want.
   utils::IteratorType pickIterType;
   /// The mapping between dim=>iterator type.
   SmallVector<utils::IteratorType> iterTypes;
 };
 
+/// A helper class that visits an affine expression and tries to find an
+/// AffineDimExpr to which the corresponding iterator from a GenericOp matches
+/// the desired iterator type.
+struct AffineDimCollector : public AffineExprVisitor<AffineDimFinder> {
+  void visitDimExpr(AffineDimExpr expr) { dims.push_back(expr); }
+  SmallVector<AffineDimExpr> dims;
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler analysis methods.
 //===----------------------------------------------------------------------===//
 
 /// Determines if affine expression is invariant.
 static bool isInvariantAffine(AffineExpr a, ArrayRef<unsigned> loopStack,
while (worklist.any() && loopDepth <= env.topSortSize()) {
     // Tries to entering the next loop level.
     loopDepth += 1;
   }
 
   assert(perm.size() == numResults);
   return AffineMap::getPermutationMap(perm, env.op().getContext());
 }
 
+/// Get the total number of compound affine expressions in affineMap that are
+/// attached to the given tensor. For the following inputs:
+///
+/// affineMap = (d0, d1, d2) => (d0 + d1, d2)
+/// tensor = ["compressed", "compressed"]
+///
+/// Returns 1 (because the first level is compressed and its corresponding
+/// affineMap is d0 + d1)
+static unsigned getNumCompoundAffineOnSparseDims(AffineMap affineMap,
+                                                 Value tensor) {
+  unsigned num = 0;
+  const auto enc = getSparseTensorEncoding(tensor.getType());
+  if (enc) {
+    const ArrayRef<AffineExpr> exps = affineMap.getResults();
+    const Level lvlRank = enc.getLvlRank();
+    assert(static_cast<Level>(exps.size()) == lvlRank);
+    for (Level l = 0; l < lvlRank; l++) {
+      // FIXME: `toOrigDim` is deprecated.
+      const Dimension d = toOrigDim(enc, l);
+      // FIXME: there's some dim/lvl confusion here; since `d` isn't
+      // guaranteed to be in bounds (for non-permutations).
+      if (!exps[d].isa<AffineDimExpr>() && !enc.isDenseLvl(l))
+        num++;
+    }
+  }
+  return num;
+}
+
+/// Get the total number of compound affine expressions attached on a sparse
+/// level in the given GenericOp.
+static unsigned getNumCompoundAffineOnSparseDims(linalg::GenericOp op) {
+  unsigned num = 0;
+  for (OpOperand &t : op->getOpOperands())
+    num += getNumCompoundAffineOnSparseDims(op.getMatchingIndexingMap(&t),
+                                            t.get());
+  return num;
+}
+
+static bool hasCompoundAffineOnSparseOut(linalg::GenericOp op) {
+  OpOperand *out = op.getDpsInitOperand(0);
+  if (getSparseTensorType(out->get()).isAllDense())
+    return false;
+  return getNumCompoundAffineOnSparseDims(op.getMatchingIndexingMap(out),
+                                          out->get());
+}
+
 /// Helper method to inspect affine expressions. Rejects cases where the
 /// same index is used more than once. Also rejects compound affine
 /// expressions in sparse dimensions.
 /// filterIdx stores the current filter loop idx should be used for the next
 /// compound affine sparse level, and it will be incremented by one when
 /// used.
 static bool findAffine(Merger &merger, unsigned tensor, unsigned dim,
                        AffineExpr a, DimLevelType dlt, unsigned &filterLdx,
case AffineExprKind::Constant: {
     // Falls through when it is a constant Affine
     return true;
   }
   default:
     return false;
   }
 }
 
-/// Get the total number of compound affine expressions in affineMap that are
-/// attached to the given tensor. For the following inputs:
-///
-/// affineMap = (d0, d1, d2) => (d0 + d1, d2)
-/// tensor = ["compressed", "compressed"]
-///
-/// Returns 1 (because the first level is compressed and its corresponding
-/// affineMap is d0 + d1)
-static unsigned getNumCompoundAffineOnSparseDims(AffineMap affineMap,
-                                                 Value tensor) {
-  unsigned num = 0;
-  const auto enc = getSparseTensorEncoding(tensor.getType());
-  if (enc) {
-    const ArrayRef<AffineExpr> exps = affineMap.getResults();
-    const Level lvlRank = enc.getLvlRank();
-    assert(static_cast<Level>(exps.size()) == lvlRank);
-    for (Level l = 0; l < lvlRank; l++) {
-      // FIXME: `toOrigDim` is deprecated.
-      const Dimension d = toOrigDim(enc, l);
-      // FIXME: there's some dim/lvl confusion here; since `d` isn't
-      // guaranteed to be in bounds (for non-permutations).
-      if (!exps[d].isa<AffineDimExpr>() && !enc.isDenseLvl(l))
-        num++;
+/// Helper method to inspect affine expressions for slice based codegen. Rejects
+/// cases where the same index is used more than once.
+static bool findSliceBasedAffine(Merger &merger, unsigned tensor, unsigned dim,
+                                 AffineExpr a, DimLevelType dlt, bool needSlice,
+                                 bool isSubExp = false) {
+  switch (a.getKind()) {
+  case AffineExprKind::DimId: {
+    unsigned idx = a.cast<AffineDimExpr>().getPosition();
+    if (!isUndefDLT(merger.getDimLevelType(tensor, idx)))
+      return false; // used more than once
+    if (!isSubExp)
+      merger.setDimAndDimLevelType(tensor, idx, dim, dlt);
+
+    if (isSubExp && needSlice) {
+      // The current loops appears in more than one affine expressions.
+      // Now, for simplicity, we do not handle this case.
+      if (merger.hasDependentSliceDim(idx)) {
+        // 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.setLoopDependentSliceDim(idx, tensor, dim);
     }
-  return num;
+    return true;
   }
-
-/// Get the total number of compound affine expressions attached on a sparse
-/// level in the given GenericOp.
-static unsigned getNumCompoundAffineOnSparseDims(linalg::GenericOp op) {
-  unsigned num = 0;
-  for (OpOperand &t : op->getOpOperands())
-    num += getNumCompoundAffineOnSparseDims(op.getMatchingIndexingMap(&t),
-                                            t.get());
-  return num;
+  case AffineExprKind::Constant:
+  case AffineExprKind::Mul:
+    // TODO: Support Mul and Constant AffineExp for slice-based codegen
+    if (needSlice)
+      return false;
+    [[fallthrough]];
+  case AffineExprKind::Add: {
+    auto binOp = a.cast<AffineBinaryOpExpr>();
+    // We do not set dim level format for affine expresssion like d0 + d1 on
+    // either loop index at d0 or d1.
+    // We continue the recursion merely to check whether current affine is
+    // admissible or not.
+    return findSliceBasedAffine(merger, tensor, dim, binOp.getLHS(), dlt,
+                                needSlice, true) &&
+           findSliceBasedAffine(merger, tensor, dim, binOp.getRHS(), dlt,
+                                needSlice, true);
   }
-
-static bool hasCompoundAffineOnSparseOut(linalg::GenericOp op) {
-  OpOperand *out = op.getDpsInitOperand(0);
-  if (getSparseTensorType(out->get()).isAllDense())
+  default:
     return false;
-  return getNumCompoundAffineOnSparseDims(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 sliceBased) {
   bool annotated = false;
   unsigned filterLdx = env.merger().getFilterLoopStartingIdx();
   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);
-    for (Level l = 0; l < lvlRank; l++) {
-      const unsigned tensor = t.getOperandNumber();
-      // FIXME: `toOrigDim` is deprecated.
-      const AffineExpr a = map.getResult(toOrigDim(enc, l));
-      if (!findAffine(env.merger(), tensor, l, a, enc.getLvlType(l), filterLdx))
+    assert(map.getNumResults() == env.op().getRank(&t));
+    // If then current tensor being inspected requires affine index, it need
+    // to be sliced.
+    bool needSlice = enc && getNumCompoundAffineOnSparseDims(map, t.get()) != 0;
+    for (unsigned l = 0; l < lvlRank; l++) {
+      unsigned tensor = t.getOperandNumber();
+      AffineExpr a = map.getResult(toOrigDim(enc, l));
+      DimLevelType dlt = enc.getLvlType(l);
+      if (sliceBased) {
+        if (!findSliceBasedAffine(env.merger(), tensor, l, a, dlt, needSlice))
+          return false; // inadmissible affine expression
+      } else {
+        if (!findAffine(env.merger(), tensor, 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.
 /// The sorted result will put the first Reduction iterator to the
 /// latest possible index.
 static bool topSortOptimal(CodegenEnv &env, unsigned n,
                            ArrayRef<utils::IteratorType> iteratorTypes,
-                           std::vector<unsigned> &inDegree,
-                           std::vector<std::vector<bool>> &adjM) {
+                           std::vector<unsigned> inDegree,
+                           const std::vector<std::vector<bool>> &adjM) {
   std::vector<unsigned> redIt;    // reduce iterator with 0 degree
   std::vector<unsigned> parIt;    // parallel iterator with 0 degree
   std::vector<unsigned> filterIt; // filter loop with 0 degree
   for (unsigned i = 0; i < n; i++) {
     if (inDegree[i] == 0) {
       if (env.merger().isFilterLoop(i))
         filterIt.push_back(i);
       else if (linalg::isReductionIterator(iteratorTypes[i]))
while (!redIt.empty() || !parIt.empty() || !filterIt.empty()) {
     // for (1 to N)
     //  for (1 to K)
     //    if (xxx)
     //      for (1 to M)
     //        O(X) computation  => O(NK+NMX) time complexity
     auto &it = !filterIt.empty() ? filterIt : (!parIt.empty() ? parIt : redIt);
     auto src = it.back();
     env.topSortPushBack(src);
+    // llvm::outs() << src;
     it.pop_back();
     // Update in-degree, and push 0-degree node into worklist.
     for (unsigned dst = 0; dst < n; dst++) {
       if (adjM[src][dst] && --inDegree[dst] == 0) {
         if (env.merger().isFilterLoop(dst))
           filterIt.push_back(dst);
         else if (linalg::isReductionIterator(iteratorTypes[dst]))
           redIt.push_back(dst);
case AffineExprKind::Mul: {
     }
     break;
   }
   default:
     break;
   }
 }
 
-static void tryLoosenAffineDenseConstraints(linalg::GenericOp op,
-                                            std::optional<unsigned> &fldx,
-                                            AffineExpr &fa,
-                                            std::optional<unsigned> &tldx,
-                                            AffineExpr &ta) {
+static void tryLoosenAffineConstraints(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 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, unsigned 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 unsigned 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() + getNumCompoundAffineOnSparseDims(env.op()) == n);
-    // Skip dense tensor constraints when not requested.
-    if (!(mask & SortMask::kIncludeDense) && !enc)
-      continue;
-    // 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.
-      AffineExpr ta = map.getResult(toOrigDim(enc, l));
-      std::optional<unsigned> tldx =
-          env.merger().getLoopIdx(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 sortArrayBasedOnOrder(std::vector<unsigned> &target,
+                                  ArrayRef<unsigned> order) {
+  std::sort(target.begin(), target.end(), [&order](unsigned l, unsigned 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;
+  });
+}
+
+static void addFilterLoopBasedConstraints(CodegenEnv &env, OpOperand &t,
+                                          OpOperand *skip, unsigned mask,
+                                          std::vector<std::vector<bool>> &adjM,
+                                          std::vector<unsigned> &inDegree) {
+  auto map = env.op().getMatchingIndexingMap(&t);
+  auto enc = getSparseTensorEncoding(t.get().getType());
+
+    return;
+  for (unsigned 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.
-        AffineExpr fa = map.getResult(toOrigDim(enc, l - 1));
-        std::optional<unsigned> fldx =
-            env.merger().getLoopIdx(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 (!(mask & SortMask::kIncludeDense))
-          tryLoosenAffineDenseConstraints(env.op(), fldx, fa, tldx, ta);
+      // We also loosen the affine constraint when use slice-based algorithm
+      // as there is no filter loop for affine index on sparse dimension.
+        tryLoosenAffineConstraints(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, unsigned 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());
+
+  // Skip dense tensor constraints when not requested.
+  if (!(mask & SortMask::kIncludeDense) && !enc)
+    return;
+
+  // No special treatment for simple indices.
+  if (getNumCompoundAffineOnSparseDims(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 (unsigned d = 1, rank = map.getNumResults(); d < rank; d++) {
+    AffineExpr fa = map.getResult(toOrigDim(enc, d - 1));
+    AffineExpr ta = map.getResult(toOrigDim(enc, d));
+
+    // 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();
+    unsigned fldx = fexp.getPosition();
+
+    finder.walkPostOrder(ta);
+    AffineDimExpr texp = finder.getDimExpr();
+    unsigned 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) {
+      unsigned f = fd.getPosition();
+      if (f == fldx)
+        continue;
+      if (!adjM[f][fldx]) {
+        adjM[f][fldx] = true;
+        inDegree[fldx]++;
+      }
+    }
+    for (auto td : tCollector.dims) {
+      unsigned 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) {
+      unsigned f = fd.getPosition();
+      if (f == fldx) // skip d_x
+        continue;
+
+      for (auto td : tCollector.dims) {
+        unsigned 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, unsigned 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 unsigned 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()) {
     // 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);
+
     if (mask & SortMask::kIncludeUndef) {
       unsigned tensor = t.getOperandNumber();
       for (unsigned i = 0; i < n; i++) {
         if (isCompressedDLT(env.dlt(tensor, i)) ||
             isSingletonDLT(env.dlt(tensor, i))) {
           for (unsigned j = 0; j < n; j++)
             if (isUndefDLT(env.dlt(tensor, j))) {
               adjM[i][j] = true;
public:
 
   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))
       return failure();
 
-    if (options.enableSliceBasedAffine)
-      llvm_unreachable("not yet implemented");
-
     // Sets up a code generation environment.
     unsigned numTensors = op->getNumOperands();
     unsigned numLoops = op.getNumLoops();
-    unsigned numFilterLoops = getNumCompoundAffineOnSparseDims(op);
-    CodegenEnv env(op, options, numTensors, numLoops, numFilterLoops);
+    unsigned numAffineOnSparse = getNumCompoundAffineOnSparseDims(op);
+    // TODO: we should probably always use slice-based codegen whenever
+    // possible, we can even intermix slice-based and filter-loop based codegen.
+    bool sliceBased = options.enableSliceBasedAffine && numAffineOnSparse != 0;
+
+    // If we uses slice based algorithm for affine index, we do not need filter
+    // loop.
+    CodegenEnv env(op, options, numTensors, numLoops,
+                   /*numFilterLoops=*/sliceBased ? 0 : numAffineOnSparse);
 
     // Detects sparse annotations and translates the per-dimension sparsity
     // information for all tensors to loop indices in the kernel.
-    if (!findSparseAnnotations(env))
+    if (!findSparseAnnotations(env, sliceBased))
       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.
     const auto allMask = {SortMask::kIncludeAll, SortMask::kIncludeUndef,
                           SortMask::kIncludeDense, SortMask::kSparseOnly};
     for (auto mask : allMask) {
-      if (computeIterationGraph(env, mask)) {
+      if (computeIterationGraph(env, mask, nullptr, sliceBased)) {
         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 sliceBased
+                 ? 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()) {
+      unsigned rank = env.op().getMatchingIndexingMap(&t).getNumResults();
+      for (unsigned i = 0; i < rank; i++) {
+        sortArrayBasedOnOrder(
+            env.merger().getRelatedLoops(t.getOperandNumber(), i),
+            env.getTopSort());
+      }
+    }
+
     // Recursively generates code if admissible.
     env.startEmit();
     genBuffers(env, rewriter);
     genInitConstantDenseAddress(env, rewriter);
     genStmt(env, rewriter, env.getTensorExp(), 0);
     genResult(env, rewriter);
     return success();
   }

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

 LatPoint::LatPoint(const BitVector &b, unsigned e) : bits(b), exp(e) {}
 
 Merger::Merger(unsigned t, unsigned l, unsigned fl)
     : outTensor(t - 1), syntheticTensor(t), numTensors(t + 1),
       numNativeLoops(l), numLoops(l + fl), hasSparseOut(false),
       dimTypes(numTensors,
                std::vector<DimLevelType>(numLoops, DimLevelType::Undef)),
       loopIdxToDim(numTensors, std::vector<std::optional<unsigned>>(
                                    numLoops, std::nullopt)),
       dimToLoopIdx(numTensors, std::vector<std::optional<unsigned>>(
-                                   numLoops, std::nullopt)) {}
+                                   numLoops, std::nullopt)),
+      ldxToDependentSlice(numTensors, std::vector<std::optional<unsigned>>(
+                                          numLoops, std::nullopt)),
+      sliceToRelatedldx(numTensors, std::vector<std::vector<unsigned>>(
+                                        numLoops, std::vector<unsigned>())) {}
 
 //===----------------------------------------------------------------------===//
 // Lattice methods.
 //===----------------------------------------------------------------------===//
 
 unsigned Merger::addExp(Kind k, unsigned e0, unsigned e1, Value v,
                         Operation *op) {
   unsigned e = tensorExps.size();

 
 void Merger::dumpBits(const BitVector &bits) const {
   for (unsigned b = 0, be = bits.size(); b < be; b++) {
     if (bits[b]) {
       unsigned t = tensor(b);
       unsigned i = index(b);
       DimLevelType dlt = dimTypes[t][i];
       llvm::dbgs() << " i_" << t << "_" << i << "_" << toMLIRString(dlt);
+      if (ldxToDependentSlice[t][i])
+        llvm::dbgs() << "_D_" << *ldxToDependentSlice[t][i];
     }
   }
 }
 
 #endif // NDEBUG
 
 //===----------------------------------------------------------------------===//
 // Builder methods.