[mlir][sparse] group tensor id and levels into pairs in loop emitter

[Peiming]

Summary

This addressed some unresolved comments in https://reviews.llvm.org/D142930

[Peiming]

remove unused code.

[wrengr]

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

46

I think it'd be better to define using TensorLevel = unsigned; and compress things like we do for TensorLoopId.

239

I think you should just define the struct directly, rather than using std::tuple. That way client code can use nice names for accessing the fields directly, rather than defining all those wrappers around std::get<N>.

Also, the struct should be named LoopSliceInfo for consistency with LoopInfo, SliceInfo, etc. And, of course, the struct should be final.

If you are concerned about the ability to use structured-binding in the foreach-loops, then take a look at the "Case 3: binding to data members" section of https://en.cppreference.com/w/cpp/language/structured_binding . tl;dr: the syntax/semantics should be exactly the same as for the std::tuple case.

248

I think the name tidLvls would be a lot cleaner here, and consistent with elsewhere.

253–256

I think it'd be best to leave this as "operating on" like it was before. Although this is "iterating" in the sparse-compiler sense of enumerating the stored values of sparse-tensors, I feel that using that term here would cause confusion vs the C++ iterator framework

255

Grammatically, that should be "information" (singular, and lowercase). However, I think it'd be better to leave it as "slice-driven loop conditions" like it was before, since the term "information" doesn't really tell us very much about what exactly that means/contains.

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

954–956

Why keep this defn of ld around but commented out? If it's no longer needed, then should remove it (along with my fixme comment about trying to figure out what it's actually supposed to be)

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

1304–1308

should be "tidLvl" to match the names used elsewhere

1305

It seems worth adding a comment explaining why we discard tlPair.second and use ldx instead; i.e., a comment saying that isSparse is checking whether any tensor is sparse for the loop ldx, regardless of what level that corresponds to and regardless of what level is stored in tidLvls. (Or, if the levels in tidLvls are supposed to agree with the loop, then say that instead. Assuming we actually maintain that invariant, then we don't need to add assertions to verify that; but we do want a comment explaining that this invariant holds.)

1551

Why condTidLvls instead of just tidLvls? (I'm not saying it should be changed, just curious why the different name)

1654

I think the affineTidLvls and affines should actually be combined into a single SmallVector<TLA> where struct TLA final { TensorId; Level; AffineExpr } (or struct TLA final {TensorLevel; AffineExpr} with the appropriate getters). That helps capture the invariant that translateBitsToTidLvlPairs keeps the the same length, and thus also helps avoid needing to zip them together later on

Of course, you'll need to come up with a better name than my "TLA" ;)

[Peiming]

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

239

Actually I delete the wrappers, I find them useless.

We only use in in structure-bindings in foreach loop, and this is a private type so it should be okay? WDYT?

[Peiming]

address some comments.

[Peiming]

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

1551

Because the callsites was using condTid (to distinguish from affineTid).

but yeah, cond- is not a accurate prefix, since not all tid, lvl are appeared in the loop condition (most of the dense levels are not). I changed it to tidLvls

[wrengr]

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

239

I still think we should define our own struct rather than reusing std::tuple. I feel like the std::pair/std::tuple classes are really only intended to be used for those cases where folks need an ad-hoc struct; hence, whenever things aren't ad-hoc then they should be given their own definitions.

[Peiming]

use single unsigned for tensor lvl pairs

[aartbik]

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

68

please put this empty line back, all

Section
//

comments have whitespace around them

90

empty line

92

agree (since the subject is plural)

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

194

Gets

208

Give this one a doc string too

[Peiming]

address comments.

[Peiming]

address comments.

[Peiming]

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

1654

I want to keep it, because the affineTidLvls are used independently from affines later. See L1662

[wrengr]

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

92

This should be TensorId

92

This should be Level

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

764

It's more appropriate to use push_back here, since the return type of makeTensorLevel is exactly what's being stored (i.e., it's not being passed to some constructor for what's actually stored)

770

ditto, should be push_back

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

203

I think a name like unpackTensorLevel would better match makeTensorLevel.

208

This should use "///" since it's documentation not a comment

208

I get what you're after by saying "Container<>" and "range<>", however since those are not actual names of templates, it'd be better to rephrase this to "Converts a range of TensorLevel to a range of std::pair<TensorId, Level>."

210

I think you'll want to add std::remove_const_t there too, since that's what all the analogous sfinae code in MLIR does for this sort of thing.

214

Mirroring the earlier suggestion, this would then be unpackTensorLevelRange

216

this should be std::forward<ContainerTy>(c)

239

Should be final.

240–242

style-guide says that the fields should be at the end of the defn

253–256

Should be just "tensor"

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

1495

push_back

1565

push_back

1583

push_back

1586

push_back

1624

push_back

1636

push_back

1662–1663

I still think it'd be better to have a single SmallVector<std::pair<TensorLevel, AffineExpr>> since that helps capture the invariant that these should be the same length.

But if you're not going to do that, then you should use llvm::zip_equal so that it checks that they have the same length rather than silently truncating things to the shorter list

1669–1671

If this concat is what's making you reluctant to use SmallVector<std::pair<TensorLevel, AffineExpr>>, then recall that you can use llvm::concat<TensorLevel>(tidLvls, llvm::make_first_range(affineTidLvlExprs)).

[Peiming]

address comments.

[Peiming]

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

240–242

Good to know!

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

1662–1663

Good to know there is a zip_equal.

[Peiming]

Use a pair for affineTidLvls and AffineExpr.

[Peiming]

minor improvement.

[wrengr]

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

210

it should be in the order remove_reference_t<remove_const_t<, at least that's the order I see elsewhere, haven't checked to see if/how it matters

[Peiming]

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

210

I don't think it matters, but yeah, good to be consistent anyway.

[Peiming]

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

210

For the record, the order matters https://stackoverflow.com/questions/30690269/why-does-using-stdremove-reference-and-stdremove-const-in-different-order-pr

[Peiming]

using llvm::remove_cvref_t

[Peiming]

rebase.

[chaitanyav]

@Peiming can you please check, the build is failing. Pasting the link to premerge-check https://buildkite.com/llvm-project/premerge-checks/builds/150384#0187e78d-a047-440e-9018-0eecabff91c8

In file included from /home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h:17,
                 from /home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp:13:
/home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h: In member function ‘constexpr mlir::sparse_tensor::TensorLevel mlir::sparse_tensor::LoopEmitter::makeTensorLevel(mlir::sparse_tensor::TensorId, mlir::sparse_tensor::Level) const’:
/home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h:199:29: error: call to non-‘constexpr’ function ‘unsigned int mlir::sparse_tensor::LoopEmitter::getNumTensors() const’
  199 |     return l * getNumTensors() + t;
      |                ~~~~~~~~~~~~~^~
/home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h:195:12: note: ‘unsigned int mlir::sparse_tensor::LoopEmitter::getNumTensors() const’ declared here
  195 |   unsigned getNumTensors() const { return tensors.size(); }
      |            ^~~~~~~~~~~~~
In file included from /usr/include/c++/11/cassert:44,
                 from /home/scratch/llvm-project/llvm/include/llvm/Support/CommandLine.h:34,
                 from /home/scratch/llvm-project/mlir/include/mlir/Pass/PassOptions.h:21,
                 from /home/scratch/llvm-project/mlir/include/mlir/Pass/PassRegistry.h:17,
                 from /home/scratch/llvm-project/mlir/include/mlir/Pass/Pass.h:13,
                 from /home/scratch/llvm-project/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h:21,
                 from /home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h:21,
                 from /home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp:13:
/home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h: In member function ‘constexpr mlir::sparse_tensor::TensorLevel mlir::sparse_tensor::CodegenEnv::makeTensorLevel(mlir::sparse_tensor::TensorId, mlir::sparse_tensor::Level) const’:
/home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h:95:37: error: call to non-‘constexpr’ function ‘unsigned int mlir::sparse_tensor::LoopEmitter::getNumTensors() const’
   95 |     assert(loopEmitter.getNumTensors() == linalgOp->getNumOperands() &&
      |            ~~~~~~~~~~~~~~~~~~~~~~~~~^~
In file included from /home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h:17

[Peiming]

In D148565#4319356, @chaitanyav wrote:

@Peiming can you please check, the build is failing. Pasting the link to premerge-check https://buildkite.com/llvm-project/premerge-checks/builds/150384#0187e78d-a047-440e-9018-0eecabff91c8

In file included from /home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h:17,
                 from /home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp:13:
/home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h: In member function ‘constexpr mlir::sparse_tensor::TensorLevel mlir::sparse_tensor::LoopEmitter::makeTensorLevel(mlir::sparse_tensor::TensorId, mlir::sparse_tensor::Level) const’:
/home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h:199:29: error: call to non-‘constexpr’ function ‘unsigned int mlir::sparse_tensor::LoopEmitter::getNumTensors() const’
  199 |     return l * getNumTensors() + t;
      |                ~~~~~~~~~~~~~^~
/home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h:195:12: note: ‘unsigned int mlir::sparse_tensor::LoopEmitter::getNumTensors() const’ declared here
  195 |   unsigned getNumTensors() const { return tensors.size(); }
      |            ^~~~~~~~~~~~~
In file included from /usr/include/c++/11/cassert:44,
                 from /home/scratch/llvm-project/llvm/include/llvm/Support/CommandLine.h:34,
                 from /home/scratch/llvm-project/mlir/include/mlir/Pass/PassOptions.h:21,
                 from /home/scratch/llvm-project/mlir/include/mlir/Pass/PassRegistry.h:17,
                 from /home/scratch/llvm-project/mlir/include/mlir/Pass/Pass.h:13,
                 from /home/scratch/llvm-project/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h:21,
                 from /home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h:21,
                 from /home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp:13:
/home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h: In member function ‘constexpr mlir::sparse_tensor::TensorLevel mlir::sparse_tensor::CodegenEnv::makeTensorLevel(mlir::sparse_tensor::TensorId, mlir::sparse_tensor::Level) const’:
/home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h:95:37: error: call to non-‘constexpr’ function ‘unsigned int mlir::sparse_tensor::LoopEmitter::getNumTensors() const’
   95 |     assert(loopEmitter.getNumTensors() == linalgOp->getNumOperands() &&
      |            ~~~~~~~~~~~~~~~~~~~~~~~~~^~
In file included from /home/scratch/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h:17

On it

[vzakhari]

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

92

getNumTensors cannot be called from constexpr function.

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

public:
   std::optional<Operation *>
   genLoopBoundary(function_ref<
                   std::optional<Operation *>(MutableArrayRef<Value> parameters)>
                       callback);
 
   //
   // Merger delegates.
   //
 
   constexpr TensorId makeTensorId(unsigned t) const {
     return latticeMerger.makeTensorId(t);
   }
   constexpr LoopId makeLoopId(unsigned i) const {
     return latticeMerger.makeLoopId(i);
   }
   constexpr TensorLoopId makeTensorLoopId(unsigned t, unsigned i) const {
     return latticeMerger.makeTensorLoopId(t, i);
   }
   const TensorExp &exp(ExprId e) const { return latticeMerger.exp(e); }
   const LatPoint &lat(LatPointId l) const { return latticeMerger.lat(l); }
   ArrayRef<LatPointId> set(LatSetId s) const { return latticeMerger.set(s); }
   DimLevelType dlt(TensorId t, LoopId i) const {
     return latticeMerger.getDimLevelType(t, i);
   }
   DimLevelType dlt(TensorLoopId b) const {
     return latticeMerger.getDimLevelType(b);
   }
 
   //
+  // LoopEmitter delegates.
+  //
+
+  constexpr TensorLevel makeTensorLevel(unsigned t, unsigned l) const {
+    // Make sure LoopEmitter, GenericOp, and Merger agree on the number of
+    // tensors. Merger has one more synthetic tensor for loop invariants.
+    assert(loopEmitter.getNumTensors() == linalgOp->getNumOperands() &&
+           loopEmitter.getNumTensors() == latticeMerger.getNumTensors() - 1);
+    return loopEmitter.makeTensorLevel(t, l);
+  }
+  std::pair<TensorId, Level> toTidLvlPair(TensorLevel tl) const {
+    return loopEmitter.toTidLvlPair(tl);
+  }
+
+  //
   // Code generation environment verify functions.
   //
 
   /// Whether the tensor expression is admissible for codegen.
   /// It also sets the sparseOut if the output tensor is sparse.
   bool isAdmissibleTensorExp(ExprId e);
 
   /// Whether the iteration graph is sorted in admissible topoOrder.

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

 /// In addition `LoopEmitter::genAffine` has `AffineDimExpr::position`
 /// correspond to `LoopId`, however it is unclear what the providence
 /// of those `AffineDimExpr` is.
 //
 // TODO: use a struct/class rather than a typedef, so that we can actually
 // typecheck this to avoid mixups in the code.
 using LoopOrd = unsigned;
 
+// A compressed <tensor id, level> pair.
+using TensorLevel = unsigned;
 //===----------------------------------------------------------------------===//
 // SparseTensorLoopEmiter class, manages sparse tensors and helps to
 // generate loop structure to (co)-iterate sparse tensors.
 //
 // An example usage:
 // To generate the following loops over T1<?x?> and T2<?x?>
 //
 // for i in TENSOR_1_0 {
public:
   ///     break p0
   ///
   ///   // Starts loop from p0
   ///   for (i = p0; i < end; i++)
   ///     ...
   ///   // loop sequence end.
   /// }
   void enterNewLoopSeq(OpBuilder &builder, Location loc,
-                       ArrayRef<TensorId> tids, ArrayRef<Level> lvls);
+                       ArrayRef<TensorLevel> tidLvls);
 
   /// Exits the current loop sequence, this will reset universal index to 0.
   void exitCurrentLoopSeq(OpBuilder &builder, Location loc);
 
   // TODO: Get rid of `lvls` in the argument list? Track the level we
   // are currently at internally. Then it would be enterNextLvlForTensor.
   // Still need a way to specify the lvl for non-annotated tensors though,
   // as those can be accessed out of order.
   //
   /// Emits loop over tensor_tid_lvl, it assumes that loops between
   /// tensor_tid_[0, lvl - 1] have already been generated.
   /// The function will also perform in-place update on the `reduc` vector to
   /// return the reduction variable used inside the generated loop.
   Operation *enterLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
-                                      ArrayRef<TensorId> tids,
-                                      ArrayRef<Level> lvls,
+                                      ArrayRef<TensorLevel> tidLvls,
                                       MutableArrayRef<Value> reduc = {},
                                       bool isParallel = false);
 
   Operation *enterFilterLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
                                             TensorId tid, Level lvl,
                                             AffineExpr affine,
                                             MutableArrayRef<Value> reduc = {});
 
-  void genDenseAffineAddress(OpBuilder &builder, Location loc, TensorId tid,
-                             Level lvl, AffineExpr lvlExpr);
+  void genDenseAffineAddress(OpBuilder &builder, Location loc,
+                             TensorLevel tidLvl, AffineExpr lvlExpr);
 
   /// Emits a co-iteration loop over a set of tensors.
   Operation *enterCoIterationOverTensorsAtLvls(
-      OpBuilder &builder, Location loc, ArrayRef<TensorId> tids,
-      ArrayRef<Level> lvls, bool needsUniv, MutableArrayRef<Value> reduc = {});
+      OpBuilder &builder, Location loc, ArrayRef<TensorLevel> tidLvls,
+      bool needsUniv, MutableArrayRef<Value> reduc = {});
 
   void exitCurrentLoop(RewriterBase &rewriter, Location loc,
                        MutableArrayRef<Value> reduc = {});
 
   /// Fills the out-parameter with the loop induction variables for all
   /// loops in the current loop-stack.  The variables are given in the
   /// same order as the loop-stack, hence `ivs` should be indexed into
   /// by `LoopOrd` (not `LoopId`).
   void getLoopIVs(SmallVectorImpl<Value> &ivs) const {
     ivs.clear();
     ivs.reserve(getCurrentDepth());
     for (auto &l : loopStack)
       ivs.push_back(l.iv);
   }
 
   /// Gets the current depth of the loop-stack.  The result is given
   /// the type `LoopOrd` for the same reason as one-past-the-end iterators.
   LoopOrd getCurrentDepth() const { return loopStack.size(); }
 
   /// Gets loop induction variable for the given `LoopOrd`.
   Value getLoopIV(LoopOrd n) const {
     return n < getCurrentDepth() ? loopStack[n].iv : Value();
   }
 
+  /// Gets the total number of tensors that loopEmitter is operating on.
+  unsigned getNumTensors() const { return tensors.size(); }
+
+  /// Compresses a TensorId and Level into a TensorLevel.
+  constexpr TensorLevel makeTensorLevel(TensorId t, Level l) const {
+    return l * getNumTensors() + t;
+  }
+
+  /// De-compresses a TensorLevel back to a pair of TensorId and Level.
+  std::pair<TensorId, Level> toTidLvlPair(TensorLevel tidLvl) const {
+    unsigned nt = getNumTensors();
+    return std::make_pair(tidLvl % nt, tidLvl / nt);
+  }
+
+  // Maps a Container<TensorLevel> into a range<std::pair<TensorId, Level>>.
+  template <class ContainerTy,
+            std::enable_if_t<std::is_same_v<typename std::remove_reference_t<
+                                                ContainerTy>::value_type,
+                                            TensorLevel>,
+                             bool> = true>
+  auto toTidLvlPairRange(ContainerTy &&c) const {
+    return llvm::map_range(
+        c, [this](TensorLevel tl) { return this->toTidLvlPair(tl); });
+  }
+
   ///
   /// Getters.
   ///
   const std::vector<std::vector<Value>> &getPosits() const { return posits; };
   const std::vector<std::vector<Value>> &getCoords() const { return coords; };
   const std::vector<std::vector<Value>> &getHighs() const { return highs; };
   const std::vector<std::vector<Value>> &getPositionBuffers() const {
     return positionsBuffers;
   };
   const std::vector<std::vector<Value>> &getCoordinateBuffers() const {
     return coordinatesBuffers;
   };
   const std::vector<Value> &getValBuffer() const { return valBuffer; };
 
   constexpr static llvm::StringLiteral getLoopEmitterLoopAttrName() {
     return llvm::StringLiteral("Emitted from");
   }
 
 private:
+  // A tuple that stored the slice-driven loop information.
+  using SliceLoopTuple = std::tuple<TensorId, Level, bool /*reduced*/>;
+
   // LoopInfo stores information of a loop generated by LoopEmitter. E.g.,
   // the set of tensors levels that the loop is iterating over.
   struct LoopInfo final {
-    LoopInfo(ArrayRef<TensorId> tids, ArrayRef<Level> lvls,
-             ArrayRef<TensorId> slicedTids, ArrayRef<Level> slicedLvls,
-             ArrayRef<bool> sliceReduced, Operation *loop, Block *userBlock,
-             Value iv, StringAttr loopTag)
-        : tids(tids), lvls(lvls), slicedTids(slicedTids),
-          slicedLvls(slicedLvls), sliceReduced(sliceReduced), loop(loop),
+    LoopInfo(ArrayRef<TensorLevel> tidLvls,
+             ArrayRef<SliceLoopTuple> sliceDrivenInfo, Operation *loop,
+             Block *userBlock, Value iv, StringAttr loopTag)
+        : tidLvls(tidLvls), sliceDrivenInfo(sliceDrivenInfo), loop(loop),
           userCodeBlock(userBlock), iv(iv) {
       // Attached a special tag to loop emitter generated loop.
       if (loopTag)
         loop->setAttr(LoopEmitter::getLoopEmitterLoopAttrName(), loopTag);
     }
-    // TODO: maybe use a vector<pair> for tid and lvl?
-    //       (Or compress them together with a `TensorLoopId`.)
-    // The set of tensors that the loop is operating on
-    const llvm::SmallVector<TensorId> tids;
-    // The corresponding levels for the tensors
-    const llvm::SmallVector<Level> lvls;
-    // The set of tensors for slice-driven loop conditions.
-    const llvm::SmallVector<TensorId> slicedTids;
-    // The corresponding level for slice-driven tensors.
-    const llvm::SmallVector<Level> slicedLvls;
-    // Whether the tensor is fully reduced (e.g., i + j => j).
-    const llvm::SmallVector<bool> sliceReduced;
+    // The set of <tensors, lvl> that the loop is operating on
+    const llvm::SmallVector<TensorLevel> tidLvls;
+    // Slice-driven loop conditions.
+    const llvm::SmallVector<SliceLoopTuple> sliceDrivenInfo;
     const Operation *loop;      // the loop operation
     Block *const userCodeBlock; // the block holding users' generated code.
     const Value iv;             // the induction variable for the loop
   };
 
   // SliceInfo stores information of an extracted slice for slice-driven loop.
   // E.g., the in-scope SSA values for the minimum coordinates and offset for
   // the slice, etc.
private:
 
   /// Generates a predicate to determine whether the tranformed coordinates are
   /// in the given slice.
   /// Returns std::pair<Transformed coordinates, Predicate>
   std::pair<Value, Value> genSliceLegitPredicate(OpBuilder &builder,
                                                  Location loc, Value crd,
                                                  TensorId tid, Level lvl);
 
-  unsigned getNumTensors() const { return tensors.size(); }
-
   bool isOutputTensor(TensorId tid) const {
     return hasOutput && tid == getNumTensors() - 1;
   }
 
   bool isSparseOutput(TensorId tid) const {
     return isOutputTensor(tid) && isSparseOut;
   }
 
   bool isValidLevel(TensorId tid, Level lvl) const {
     return tid < lvlTypes.size() && lvl < lvlTypes[tid].size();
   }
 
   /// Prepares loop for iterating over `tensor[lvl]`, under the assumption
   /// that `tensor[0...lvl-1]` loops have already been set up.
   void prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
                                   TensorId tid, Level lvl);
 
   /// Emits extra locals, since the locals might not be in simplified lattices
   /// point used to generate the loops, but are still required to generate
   /// expressions.
   void emitExtraLocalsForTensorsAtDenseLvls(OpBuilder &builder, Location loc,
-                                            ArrayRef<TensorId> tids,
-                                            ArrayRef<Level> lvls);
+                                            ArrayRef<TensorLevel> tidLvls);
 
   /// Emits a for loop to iterate over a tensor level with the provided lower
   /// bound `lo` and upper bound `hi`.
   /// Apart from iterating just single tensor level, for loops can be used for
   /// slice-driven loop on dense level too.
   Operation *emitForLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
                                         TensorId tid, Level lvl, Value lo,
                                         Value hi, MutableArrayRef<Value> reduc,

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

for (Level lvl = 0; lvl < lvlRank; lvl++) {
         }
       }
     }
   }
   localInsertPos = builder.getInsertionPoint()->getPrevNode();
 }
 
 void LoopEmitter::enterNewLoopSeq(OpBuilder &builder, Location loc,
-                                  ArrayRef<TensorId> tids,
-                                  ArrayRef<Level> lvls) {
+                                  ArrayRef<TensorLevel> tidLvls) {
   // TODO: sort
   assert(loopSeqStack.size() == loopStack.size());
   // Prepares for all the tensors used in the current loop sequence.
   std::vector<std::tuple<TensorId, Level, bool>> slicedTids;
-  for (auto [tid, lvl] : llvm::zip(tids, lvls)) {
+  for (auto [tid, lvl] : toTidLvlPairRange(tidLvls)) {
     if (!dependentLvlMap[tid][lvl].empty()) {
       bool fullyRed = genSliceBegin(builder, loc, tid, lvl);
       slicedTids.emplace_back(tid, lvl, fullyRed);
     } else {
       prepareLoopOverTensorAtLvl(builder, loc, tid, lvl);
     }
   }
 
Operation *loop =
             coords[tid][lvl] = insertPoint->getResult(0);
           })
           .first;
   // Sets the insertionn pointer inside loop body.
   builder.setInsertionPointAfter(insertPoint);
   return loop;
 }
 
-Operation *LoopEmitter::enterLoopOverTensorAtLvl(
-    OpBuilder &builder, Location loc, ArrayRef<TensorId> tids,
-    ArrayRef<Level> lvls, MutableArrayRef<Value> reduc, bool isParallel) {
+Operation *LoopEmitter::enterLoopOverTensorAtLvl(OpBuilder &builder,
+                                                 Location loc,
+                                                 ArrayRef<TensorLevel> tidLvls,
+                                                 MutableArrayRef<Value> reduc,
+                                                 bool isParallel) {
   // TODO: support multiple return on parallel for?
   assert(!isParallel || reduc.size() <= 1);
   bool isSparseCond = false, isSparseSliceCond = false;
-  size_t tid = tids.front(), lvl = lvls.front();
+  auto [tid, lvl] = toTidLvlPair(tidLvls.front());
 
   // Finds out the tensor level that we should use to generate loops. Amongs all
   // the tensor levels, there is at most one sparse tensor level.
-  for (auto [t, l] : llvm::zip(tids, lvls)) {
+  for (auto [t, l] : toTidLvlPairRange(tidLvls)) {
     assert(lvlTypes[t].size() > l);         // Must be a valid tid, dim pair
     assert(!coords[t][l] ||                 // We cannot re-enter the same level
            !dependentLvlMap[t][l].empty()); // unless it is a slice-driver loop
     auto lvlType = lvlTypes[t][l];
     // Must be a recognizable DLT.
     assert(isDenseDLT(lvlType) || isCompressedDLT(lvlType) ||
            isCompressedWithHiDLT(lvlType) || isSingletonDLT(lvlType));
 
Operation *LoopEmitter::enterLoopOverTensorAtLvl(OpBuilder &builder,
   bool isDenseSliceCond =
       isDenseDLT(lvlType) && !dependentLvlMap[tid][lvl].empty();
   // if the slice is fully reduced, we can now use TACO-based algorithm to
   // iterate it.
 
   Operation *l = nullptr;
 
   // At most one tensor used as condition in for loop;
-  SmallVector<TensorId, 1> condTid;
-  SmallVector<Level, 1> condLvl;
+  SmallVector<TensorLevel, 1> condTidLvl;
   // There Might be multiple dense slice driven tensor.
-  SmallVector<TensorId> sliceTids;
-  SmallVector<Level> sliceLvls;
-  SmallVector<bool> sliceReduc;
+  SmallVector<SliceLoopTuple> sliceDrivenTuples;
 
   // Generates loops differently depending on whether we need a slice-driven
   // loop or a simple level traversal loop.
   if (isSparseSliceCond) {
     bool fullyReduced = depFullyReduced(tid, lvl);
     if (!fullyReduced) {
       l = emitSliceDrivenLoopOverTensorAtLvl(builder, loc, tid, lvl, reduc);
     } else {
       // If the slice is fully reduced, we can now use TACO-based algorithm to
       // iterate it.
       l = emitWhileLoopOverSliceAtSparseLvl(
           builder, loc, posits[tid][lvl], highs[tid][lvl],
           getFinalSliceOnLvl(tid, lvl).offset, sliceSizes[tid][lvl].back(), tid,
           lvl, reduc);
     }
     levelReducedDep[tid][lvl]++;
-    sliceTids.push_back(tid);
-    sliceLvls.push_back(lvl);
-    sliceReduc.push_back(fullyReduced);
+    sliceDrivenTuples.emplace_back(tid, lvl, fullyReduced);
   } else {
     Value lo = isSparseCond ? posits[tid][lvl]           // current offset
                             : loopSeqStack.back().first; // universal index
     Value hi = highs[tid][lvl];
     if (isDenseSliceCond) {
       bool fullyReduced = depFullyReduced(tid, lvl);
       Value sliceSz = sliceSizes[tid][lvl][sliceStack[tid].back().depth - 1];
       // Adjust for loop hi for dense slice-driven loop.
       if (fullyReduced) {
         hi = sliceSz;
-        condTid.push_back(tid);
-        condLvl.push_back(lvl);
+        condTidLvl.emplace_back(makeTensorLevel(tid, lvl));
       } else {
         hi = SUBI(lvlSizes[tid][lvl], sliceSz);
         hi = ADDI(hi, C_IDX(1));
       }
     } else {
-      condTid.push_back(tid);
-      condLvl.push_back(lvl);
+      condTidLvl.emplace_back(makeTensorLevel(tid, lvl));
     }
     l = emitForLoopOverTensorAtLvl(builder, loc, tid, lvl, lo, hi, reduc,
                                    isParallel);
   }
   Value iv = coords[tid][lvl];
-  for (auto [t, l] : llvm::zip(tids, lvls)) {
+  for (auto [t, l] : toTidLvlPairRange(tidLvls)) {
     // We only need to handle slice-driven loops on dense level here.
     // If it is a slice-driven loop on sparse level, it needs a while loop to
     // insert break statements, and it must have been handled correctly in L692.
     if (!dependentLvlMap[t][l].empty() && isDenseDLT(lvlTypes[t][l])) {
       // Pushes sliced levels to build correct LoopInfo.
       bool fullyReduc = depFullyReduced(t, l);
       SliceInfo &info = sliceStack[t].back();
       if (fullyReduc) {
         posits[t][l] = genAddress(builder, loc, t, l, ADDI(info.offset, iv));
       } else {
         // Puts sliced dense loop into LoopInfo so that LoopEmitter knows how to
         // exit it.
-        sliceTids.push_back(t);
-        sliceLvls.push_back(l);
-        sliceReduc.push_back(fullyReduc);
+        sliceDrivenTuples.emplace_back(t, l, fullyReduc);
         // Update the slice information as we enter the new loop.
         assert(*info.slicedOnLvl == l);
         info.minCrd = info.offset = iv;
         info.isNonEmpty = constantI1(builder, loc, true);
         levelReducedDep[t][l]++;
       }
     }
   }
   // NOTE: we can also prepare for next dim here in advance
   // Pushes the loop into stack.
-  loopStack.emplace_back(condTid, condLvl, sliceTids, sliceLvls, sliceReduc, l,
+  loopStack.emplace_back(condTidLvl, sliceDrivenTuples, l,
                          builder.getInsertionBlock(), iv, loopTag);
   // Emit extra locals.
-  emitExtraLocalsForTensorsAtDenseLvls(builder, loc, tids, lvls);
+  emitExtraLocalsForTensorsAtDenseLvls(builder, loc, tidLvls);
   return l;
 }
 
 Operation *LoopEmitter::enterFilterLoopOverTensorAtLvl(
     OpBuilder &builder, Location loc, TensorId tid, Level lvl,
     AffineExpr affine, MutableArrayRef<Value> reduc) {
   assert(isValidLevel(tid, lvl));
   assert(!affine.isa<AffineDimExpr>() && !isDenseDLT(lvlTypes[tid][lvl]));
if (hasReduc) {
     // On mismatch.
     YIELD(reduc);
   }
   // Set the insert point to matched branch.
   builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
 
   // NOTE: we can also prepare for next lvl here in advance
   // Push the loop into stack
-  loopStack.emplace_back(ArrayRef<TensorId>(tid), ArrayRef<Level>(lvl),
-                         ArrayRef<TensorId>(), ArrayRef<Level>(),
-                         ArrayRef<bool>(), forOp, builder.getInsertionBlock(),
-                         coords[tid][lvl], nullptr);
+  loopStack.emplace_back(ArrayRef<TensorLevel>(makeTensorLevel(tid, lvl)),
+                         ArrayRef<SliceLoopTuple>(), forOp,
+                         builder.getInsertionBlock(), coords[tid][lvl],
+                         nullptr);
   return forOp;
 }
 
 void LoopEmitter::genDenseAffineAddress(OpBuilder &builder, Location loc,
-                                        TensorId tid, Level lvl,
+                                        TensorLevel tidLvl,
                                         AffineExpr lvlExpr) {
+  auto [tid, lvl] = toTidLvlPair(tidLvl);
   assert(isDenseDLT(lvlTypes[tid][lvl]));
   // For dense levels, the level-coordinate also serves as the position.
   Value lvlCrd = genAffine(builder, loc, lvlExpr);
   posits[tid][lvl] = genAddress(builder, loc, tid, lvl, lvlCrd);
 }
 
 Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
-    OpBuilder &builder, Location loc, ArrayRef<TensorId> tids,
-    ArrayRef<Level> lvls, bool needsUniv, MutableArrayRef<Value> reduc) {
+    OpBuilder &builder, Location loc, ArrayRef<TensorLevel> tidLvls,
+    bool needsUniv, MutableArrayRef<Value> reduc) {
   // NOTE: the slice driven tensor-related reduction variable must
   // appear before normal tensors.
-  assert(tids.size() == lvls.size());
   SmallVector<Type> types;
   SmallVector<Value> operands;
   // Construct the while-loop with a parameter for each coordinate.
   const Type indexType = builder.getIndexType();
-  for (auto [tid, lvl] : llvm::zip(tids, lvls)) {
+  for (auto [tid, lvl] : toTidLvlPairRange(tidLvls)) {
     // TODO: support coiteration with slice driven tensors.
     const auto lvlTp = lvlTypes[tid][lvl];
     assert(dependentLvlMap[tid][lvl].empty() && "TODO: not yet implemented");
     if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
         isCompressedWithHiDLT(lvlTp)) {
       const auto reassoc = getCollapseReassociation(tid, lvl);
       for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
         if (!isUniqueDLT(lvlTypes[tid][reassoc[i]])) {
Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
   Block *before = builder.createBlock(&whileOp.getBefore(), {}, types, locs);
   Block *after = builder.createBlock(&whileOp.getAfter(), {}, types, locs);
 
   // Build the "before" region, which effectively consists
   // of a conjunction of "i < upper" tests on all induction.
   builder.setInsertionPointToStart(&whileOp.getBefore().front());
   Value cond;
   unsigned o = 0;
-  for (auto [t, lvl] : llvm::zip(tids, lvls)) {
+  for (auto [t, lvl] : toTidLvlPairRange(tidLvls)) {
     const TensorId tid = t; // Why `t` can not be captured by lambda?
     const auto lvlTp = lvlTypes[tid][lvl];
     if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
         isCompressedWithHiDLT(lvlTp)) {
       const auto reassoc = getCollapseReassociation(tid, lvl);
       assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
       for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
         if (!isUniqueDLT(lvlTypes[tid][reassoc[i]])) {
Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
   }
   builder.create<scf::ConditionOp>(loc, cond, before->getArguments());
 
   // Generates while body.
   builder.setInsertionPointToStart(&whileOp.getAfter().front());
 
   SmallVector<std::pair<Value, unsigned>> slicesPreds;
   unsigned i = 0;
-  for (auto [tid, lvl] : llvm::zip(tids, lvls)) {
+  for (auto [tid, lvl] : toTidLvlPairRange(tidLvls)) {
     // Prepares for next level.
     const auto lvlTp = lvlTypes[tid][lvl];
     if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
         isCompressedWithHiDLT(lvlTp)) {
       coords[tid][lvl] = genSparseCrd(builder, loc, tid, lvl);
       if (isSparseSlices[tid]) {
         auto [trans, pred] =
             genSliceLegitPredicate(builder, loc, coords[tid][lvl], tid, lvl);
if (!slicesPreds.empty()) {
 
     // If all slices are legit, start the user generated code.
     builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
   }
 
   Value min;
   // Finds the minimum coordinate
   if (!needsUniv) {
-    for (auto [tid, lvl] : llvm::zip(tids, lvls)) {
+    for (auto [tid, lvl] : toTidLvlPairRange(tidLvls)) {
       const auto lvlTp = lvlTypes[tid][lvl];
       if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
           isCompressedWithHiDLT(lvlTp)) {
         const auto crd = coords[tid][lvl];
         if (min) {
           Value cmp = CMPI(ult, coords[tid][lvl], min);
           min = SELECT(cmp, coords[tid][lvl], min);
         } else {
           min = crd;
         }
       }
     }
   } else {
     assert(!min);
     // Otherwise, universal index is the minimal pos.
     min = after->getArguments().back();
   }
 
   // Sets up the loop stack.
-  loopStack.emplace_back(tids, lvls, ArrayRef<TensorId>(), ArrayRef<Level>(),
-                         ArrayRef<bool>(), whileOp, builder.getInsertionBlock(),
-                         min, loopTag);
+  loopStack.emplace_back(tidLvls, ArrayRef<SliceLoopTuple>(), whileOp,
+                         builder.getInsertionBlock(), min, loopTag);
   assert(loopStack.size() == loopSeqStack.size());
 
-  for (auto [tid, dstLvl] : llvm::zip(tids, lvls)) {
+  for (auto [tid, dstLvl] : toTidLvlPairRange(tidLvls)) {
     const auto reassoc = getCollapseReassociation(tid, dstLvl);
     assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
     // TODO: Refactors this into smaller functions.
     // NOTE: For all the collapsed level (except for the last one, that is why
     // the loop ends with `reassoc.size() - 1`), as each iteration is advanced
     // by the segment size of the last level, which does not always invalidate
     // the segment size for the previous levels, thus we need to propagate the
     // segment sizes across loop iterations and only forward if needed.
for (auto [tid, dstLvl] : toTidLvlPairRange(tidLvls)) {
     const auto srcLvl = reassoc.back();
     if (!isUniqueDLT(lvlTypes[tid][srcLvl])) {
       segHi[tid][srcLvl] = genSegmentHigh(
           builder, loc, tid, srcLvl, posits[tid][srcLvl], highs[tid][srcLvl]);
     }
   }
 
   // Emits extra locals
-  emitExtraLocalsForTensorsAtDenseLvls(builder, loc, tids, lvls);
+  emitExtraLocalsForTensorsAtDenseLvls(builder, loc, tidLvls);
 
   // Updates reduction variables
   assert(after->getNumArguments() == o + reduc.size() + (needsUniv ? 1 : 0));
   // In-place update on reduction variable.
   for (unsigned i = 0, e = reduc.size(); i < e; i++)
     reduc[i] = after->getArgument(o + i);
 
   return whileOp;
if (isSingletonDLT(lvlTp)) {
               : ADDI(pLo, c1);
       return;
     }
   }
 
   llvm_unreachable("Unrecognized level-type!");
 }
 
-void LoopEmitter::emitExtraLocalsForTensorsAtDenseLvls(OpBuilder &builder,
-                                                       Location loc,
-                                                       ArrayRef<TensorId> tids,
-                                                       ArrayRef<Level> lvls) {
+void LoopEmitter::emitExtraLocalsForTensorsAtDenseLvls(
+    OpBuilder &builder, Location loc, ArrayRef<TensorLevel> tidLvls) {
   // Initialize dense positions. Note that we generate dense coordinates of the
   // output tensor unconditionally, since they may not appear in the lattice,
   // but may be needed for linearized codegen.
-  assert(tids.size() == lvls.size());
-  for (auto [tid, lvl] : llvm::zip(tids, lvls)) {
+  for (auto [tid, lvl] : toTidLvlPairRange(tidLvls)) {
     if (isDenseDLT(lvlTypes[tid][lvl])) {
       // Slice-driven dense level should have be handled already.
       if (!dependentLvlMap[tid][lvl].empty())
         continue;
 
       auto enc = getSparseTensorEncoding(tensors[tid].getType());
       if (enc && !isSparseOutput(tid)) {
         bool validPos = lvl == 0 || posits[tid][lvl - 1];
for (auto [tid, lvl] : toTidLvlPairRange(tidLvls)) {
     }
   }
 }
 
 void LoopEmitter::exitForLoop(RewriterBase &rewriter, Location loc,
                               MutableArrayRef<Value> reduc) {
   const LoopInfo &loopInfo = loopStack.back();
   rewriter.setInsertionPointToEnd(loopInfo.userCodeBlock);
-  for (auto [tid, lvl, reduced] : llvm::zip(
-           loopInfo.slicedTids, loopInfo.slicedLvls, loopInfo.sliceReduced)) {
+  for (auto [tid, lvl, reduced] : loopInfo.sliceDrivenInfo) {
     SliceInfo &info = sliceStack[tid].back();
     assert(isDenseDLT(lvlTypes[tid][lvl]));
     assert(*info.slicedOnLvl == lvl && !reduced);
     (void)reduced;
     // Resets slices pointers as the resolved slices are invalidated after we
     // moves forward to the next slice.
     invalidateSliceIterIdx(rewriter, loc, tid, lvl);
     info.minCrd = info.offset = info.isNonEmpty = Value();
#endif // NDEBUG
     // In-place update reduction variables.
     for (unsigned i = 0, e = parOp.getResults().size(); i < e; i++)
       reduc[i] = parOp.getResult(i);
   }
 
   // Finished iterating a tensor, clean up
   // We only do the clean up on for loop as while loops do not necessarily
   // finish the iteration on a sparse tensor
-  for (auto [tid, lvl] : llvm::zip(loopInfo.tids, loopInfo.lvls)) {
+  for (auto [tid, lvl] : toTidLvlPairRange(loopInfo.tidLvls)) {
     // Reset to null.
     coords[tid][lvl] = Value();
     posits[tid][lvl] = Value();
     // Dense level, high is fixed.
     if (!isDenseDLT(lvlTypes[tid][lvl]))
       highs[tid][lvl] = Value();
   }
 }
 
 void LoopEmitter::exitWhileLoop(OpBuilder &builder, Location loc,
                                 MutableArrayRef<Value> reduc) {
   const LoopInfo &loopInfo = loopStack.back();
   auto whileOp = llvm::cast<scf::WhileOp>(loopInfo.loop);
   builder.setInsertionPointToEnd(loopInfo.userCodeBlock);
   Value iv = loopInfo.iv;
 
   // Finalize the induction. Note that the induction could be performed
   // in the individual if-branches to avoid re-evaluating the conditions.
   // However, that would result in a rather elaborate forest of yield
   // instructions during code generation. Moreover, performing the induction
   // after the if-statements more closely resembles code generated by TACO.
   unsigned o = 0;
   SmallVector<Value> operands;
   unsigned delta = 0;
-  for (auto [tid, lvl, resolved] : llvm::zip(
-           loopInfo.slicedTids, loopInfo.slicedLvls, loopInfo.sliceReduced)) {
+  for (auto [tid, lvl, resolved] : loopInfo.sliceDrivenInfo) {
     // TODO: handle dense.
     assert(isCompressedDLT(lvlTypes[tid][lvl]));
     levelReducedDep[tid][lvl]--;
     if (!resolved) {
       genSliceNextInduction(builder, loc, whileOp, tid, lvl, operands, o);
       continue;
     }
     // TODO: We need to distinguish coiterate loop with slice-driven loop and
     // fully reduced while op for iterating one slices.
     // FIXME: since we didn't implement coiteration, this must be iteration
     // just on fully resolved slice.
-    assert(loopInfo.slicedTids.size() == 1 && loopInfo.tids.empty());
+    assert(loopInfo.sliceDrivenInfo.size() == 1 && loopInfo.tidLvls.empty());
     // The if guard to filter out out-range coordinates.
     assert(llvm::isa<scf::IfOp>(builder.getInsertionBlock()->getParentOp()));
     posits[tid][lvl] = whileOp->getResult(o++);
     // FIXME: we are not using continue here since we do not support
     // coiteration on slices. But it need to be treated similarly as the
     // universal index.
     o++; // skip continue flag.
     // Since we did not push two results from whileOp. The size of the
     // operands vector is smaller than the actual number of return values from
     // the whileOp.
     // It is because we are actually generating yield in the IfOp inside the
     // whileOp to only iterates over inbound coordinates within the slices.
     delta += 2;
   };
 
   Value one = C_IDX(1);
-  for (auto [tid, dstLvl] : llvm::zip(loopInfo.tids, loopInfo.lvls)) {
+  for (auto [tid, dstLvl] : toTidLvlPairRange(loopInfo.tidLvls)) {
     const auto lvlTp = lvlTypes[tid][dstLvl];
     if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
         isCompressedWithHiDLT(lvlTp)) {
       const auto reassoc = getCollapseReassociation(tid, dstLvl);
       assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
       for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
         const Level srcLvl = reassoc[i];
         if (!isUniqueDLT(lvlTypes[tid][srcLvl])) {
void LoopEmitter::exitWhileLoop(OpBuilder &builder, Location loc,
   builder.setInsertionPointAfter(whileOp);
 }
 
 void LoopEmitter::exitCurrentLoop(RewriterBase &rewriter, Location loc,
                                   MutableArrayRef<Value> reduc) {
   // Clean up the values, it would help use to discover potential bug at a
   // earlier stage (instead of silently using a wrong value).
   const LoopInfo &loopInfo = loopStack.back();
-  assert(loopInfo.tids.size() == loopInfo.lvls.size());
   SmallVector<Value> red;
   if (llvm::isa<scf::WhileOp>(loopInfo.loop)) {
     exitWhileLoop(rewriter, loc, reduc);
   } else {
     exitForLoop(rewriter, loc, reduc);
   }
 
   assert(loopStack.size() == loopSeqStack.size());

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

LogicalResult matchAndRewrite(ForeachOp op,
     // 1. Generates loop for the sparse input.
     LoopEmitter loopEmitter(
         ValueRange{input},
         StringAttr::get(getContext(), ForeachOp::getOperationName()));
     loopEmitter.initializeLoopEmit(rewriter, loc);
     for (Level l = 0; l < lvlRank; l++) {
       // TODO: provide utility function for loop sequences that only contains
       // one for loop?
-      // FIXME(wrengr): what is this "ld" supposed to be really?
-      const Level ld = op.getOrder() ? op.getOrder()->getDimPosition(l) : l;
-      const SmallVector<TensorId, 1> tids{0};
-      loopEmitter.enterNewLoopSeq(rewriter, loc, tids, ld);
+      const SmallVector<TensorLevel, 1> tidLvls{
+          loopEmitter.makeTensorLevel(0, l)};
+      loopEmitter.enterNewLoopSeq(rewriter, loc, tidLvls);
       // Note that reduc will be taken care of by loop emitter and get updated
       // in place.
-
-      loopEmitter.enterLoopOverTensorAtLvl(rewriter, loc, tids, l, reduc);
+      loopEmitter.enterLoopOverTensorAtLvl(rewriter, loc, tidLvls, reduc);
     }
 
     SmallVector<Value> lcvs;
     lcvs.reserve(lvlRank);
     loopEmitter.getLoopIVs(lcvs);
 
     if (op.getOrder()) {
       // FIXME: There is some dim/lvl confusion here since `dimRank != lvlRank`

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

static bool isParallelFor(CodegenEnv &env, bool isOuter, bool isSparse) {
   case SparseParallelizationStrategy::kAnyStorageAnyLoop:
     return true;
   }
   llvm_unreachable("unexpected parallelization strategy");
 }
 
 /// Generates a for-loop on a single index.
 static Operation *genFor(CodegenEnv &env, OpBuilder &builder, bool isOuter,
-                         bool isInner, LoopId ldx, ArrayRef<TensorId> tids,
-                         ArrayRef<Level> lvls) {
+                         bool isInner, LoopId ldx,
+                         ArrayRef<TensorLevel> tidLvls) {
   linalg::GenericOp op = env.op();
   Location loc = op.getLoc();
   auto iteratorTypes = op.getIteratorTypesArray();
-  bool isSparse = llvm::any_of(tids, [ldx, &env](TensorId tid) {
-    const auto dlt = env.dlt(tid, ldx);
+  bool isSparse = llvm::any_of(tidLvls, [ldx, &env](TensorLevel tidLvl) {
+    // Queries the DLT based on the tensor id and loop idx, as requested by
+    // `CodegenEnv::dlt(TensorId, LoopIdx)`. The returned DLT from CodegenEnv
+    // should be consistent with the DLT indexed by <TensorId, Level>.
+    const auto dlt = env.dlt(env.toTidLvlPair(tidLvl).first, ldx);
     return isCompressedDLT(dlt) || isSingletonDLT(dlt);
   });
 
   bool isParallel = isParallelFor(env, isOuter, isSparse);
 
   Operation *loop = *env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
     if (env.merger().isFilterLoop(ldx)) {
-      const TensorId tid = tids.front();
-      const Level lvl = lvls.front();
+      const auto [tid, lvl] = env.toTidLvlPair(tidLvls.front());
       // tids/lvls must only have one value because filter loops only
       // corresponding to the one and only sparse tensor level.
-      assert(isSparse && tids.size() == 1 && lvls.size() == 1);
+      assert(isSparse && tidLvls.size() == 1);
       OpOperand *t = &op->getOpOperand(tid);
       auto enc = getSparseTensorEncoding(t->get().getType());
       // Retrieves the affine expression for the filter loop.
       // FIXME: `toOrigDim` is deprecated.
       AffineExpr a =
           op.getMatchingIndexingMap(t).getResult(toOrigDim(enc, lvl));
       return env.emitter().enterFilterLoopOverTensorAtLvl(builder, loc, tid,
                                                           lvl, a, reduc);
     }
-    return env.emitter().enterLoopOverTensorAtLvl(builder, loc, tids, lvls,
-                                                  reduc, isParallel);
+    return env.emitter().enterLoopOverTensorAtLvl(builder, loc, tidLvls, reduc,
+                                                  isParallel);
   });
   assert(loop);
   return loop;
 }
 
 /// Emit a while-loop for co-iteration over multiple indices.
 static Operation *genWhile(CodegenEnv &env, OpBuilder &builder, LoopId idx,
-                           bool needsUniv, ArrayRef<TensorId> tids,
-                           ArrayRef<Level> lvls) {
+                           bool needsUniv, ArrayRef<TensorLevel> tidLvls) {
   Operation *loop = *env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
     // Construct the while-loop with a parameter for each
     // index.
     return env.emitter().enterCoIterationOverTensorsAtLvls(
-        builder, env.op().getLoc(), tids, lvls, needsUniv, reduc);
+        builder, env.op().getLoc(), tidLvls, needsUniv, reduc);
   });
   assert(loop);
   return loop;
 }
 
 /// Generates a for-loop or a while-loop, depending on whether it implements
 /// singleton iteration or co-iteration over the given conjunction.
 static Operation *genLoop(CodegenEnv &env, OpBuilder &builder, LoopOrd at,
-                          bool needsUniv, ArrayRef<TensorId> tids,
-                          ArrayRef<Level> lvls, bool isFor) {
-  assert(tids.size() == lvls.size());
+                          bool needsUniv, ArrayRef<TensorLevel> tidLvls,
+                          bool isFor) {
   const LoopId idx = env.topSortAt(at);
   if (isFor) {
     bool isOuter = at == 0;
     bool isInner = at == env.topSortSize() - 1;
-    return genFor(env, builder, isOuter, isInner, idx, tids, lvls);
+    return genFor(env, builder, isOuter, isInner, idx, tidLvls);
   }
-  return genWhile(env, builder, idx, needsUniv, tids, lvls);
+  return genWhile(env, builder, idx, needsUniv, tidLvls);
 }
 
 /// Generates the induction structure for a while-loop.
 static void finalizeWhileOp(CodegenEnv &env, OpBuilder &builder, LoopId idx,
                             bool needsUniv, scf::WhileOp whileOp) {
   Location loc = env.op().getLoc();
   // Finalize each else branch of all if statements.
   if (env.isReduc() || env.isExpand() || env.getInsertionChain()) {
static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp,
   // Emit invariants at this loop sequence level.
   genInvariants(env, builder, exp, ldx, /*atStart=*/true);
   // Emit access pattern expansion for sparse tensor output.
   genExpand(env, builder, at, /*atStart=*/true);
   // Emit further intitialization at this loop sequence level.
   const LatPointId l0 = env.set(lts)[0];
   bool needsUniv = false;
 
-  SmallVector<TensorId> tids;
-  SmallVector<Level> lvls;
+  SmallVector<TensorLevel> tidLvls;
   env.merger().foreachTensorLoopId(l0, [&](TensorLoopId b, TensorId tid,
                                            std::optional<Level> lvl,
                                            DimLevelType dlt, bool isIdxReduc) {
     assert(env.merger().loop(b) == idx);
     if (isDenseDLT(dlt) || isUndefDLT(dlt))
       needsUniv = true;
     if (isCompressedDLT(dlt) || isSingletonDLT(dlt) ||
         isCompressedWithHiDLT(dlt) || isIdxReduc) {
       // Only when this is a index reduction loop, can the dlt be undefined.
       assert(!isUndefDLT(dlt) || isIdxReduc);
       // sparse/singleton levels, or a dense/sparse index reduction loop.
-      tids.push_back(tid);
-      lvls.push_back(*lvl);
+      tidLvls.emplace_back(env.makeTensorLevel(tid, *lvl));
     }
   });
 
-  env.emitter().enterNewLoopSeq(builder, env.op().getLoc(), tids, lvls);
+  env.emitter().enterNewLoopSeq(builder, env.op().getLoc(), tidLvls);
 
   // Maintain the universal index only if it is actually
   // consumed by a subsequent lattice point.
   if (needsUniv) {
     for (const LatPointId li : env.set(lts).drop_front())
       if (!env.merger().hasAnySparse(env.lat(li).simple))
         return true;
   }
if (enc) {
     const TensorId tid = env.makeTensorId(input->getOperandNumber());
     const Level lvlRank = enc.getLvlRank();
     assert(lvlExprs.size() == static_cast<size_t>(lvlRank));
     // FIXME: there is dim/lvl confusion here
     for (Level l = startLvl; l < lvlRank; l++) {
       // FIXME: `toOrigDim` is deprecated.
       AffineExpr lvlExpr = lvlExprs[toOrigDim(enc, l)];
       if (enc.isDenseLvl(l) && lvlExpr.isa<AffineConstantExpr>())
-        env.emitter().genDenseAffineAddress(builder, loc, tid, l, lvlExpr);
+        env.emitter().genDenseAffineAddress(
+            builder, loc, env.makeTensorLevel(tid, l), lvlExpr);
       else
         return; // break on first non-dense non-constant level
     }
   }
 }
 
 static void genInitConstantDenseAddress(CodegenEnv &env,
                                         RewriterBase &rewriter) {
   // We can generate address for constant affine expression before any loops
   // starting from the first level as they do not depend on any thing.
   // E.g., [Dense, Dense, Sparse] -> (1, 2, d0), the addresses for the first two
   // levels can be determined before loops.
   for (TensorId tid = 0, e = env.op().getNumDpsInputs(); tid < e; tid++)
     genConstantDenseAddressFromLevel(env, rewriter, tid, 0);
 }
 
 /// Return true if the lattices bit can be iterated by a for loop.
-static bool translateBitsToTidLvlPairs(
-    CodegenEnv &env, LatPointId li, LoopId ldx, SmallVectorImpl<TensorId> &tids,
-    SmallVectorImpl<Level> &lvls, SmallVectorImpl<TensorId> &affineTids,
-    SmallVectorImpl<Level> &affineLvls, SmallVectorImpl<AffineExpr> &exps) {
+static bool
+translateBitsToTidLvlPairs(CodegenEnv &env, LatPointId li, LoopId ldx,
+                           SmallVectorImpl<TensorLevel> &tidLvls,
+                           SmallVectorImpl<TensorLevel> &affineTidLvls,
+                           SmallVectorImpl<AffineExpr> &exps) {
   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, bool isIdxReduc) {
         if (simple[b]) {
           if (isIdxReduc) {
-            tids.push_back(tid);
-            lvls.push_back(*lvl);
+            tidLvls.emplace_back(env.makeTensorLevel(tid, *lvl));
             numloopCond++;
             return;
           }
           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)
             tid = outTid;
             lvl = outLvl;
             // Skips invalid lvl (e.g., when this is a zero ranked tensor).
             if (!lvl)
               return;
           }
           hasNonUnique = !isUniqueDLT(dlt) || hasNonUnique;
-          tids.push_back(tid);
-          lvls.push_back(*lvl);
+          tidLvls.emplace_back(env.makeTensorLevel(tid, *lvl));
           numloopCond++;
         } else if (isDenseDLT(dlt) || isIdxReduc) {
-          tids.push_back(tid);
-          lvls.push_back(*lvl);
+          tidLvls.emplace_back(env.makeTensorLevel(tid, *lvl));
         } else {
           assert(isUndefDLT(dlt));
           linalg::GenericOp op = env.op();
           if (tid >= op.getNumDpsInputs())
             // We only handle affine expression on input tensors (for now).
             return;
           OpOperand *operand = &op->getOpOperand(tid);
           const auto stt = getSparseTensorType(operand->get());
env.merger().foreachTensorLoopId(
                 // level. We need to generate the address according to the
                 // affine expression. This is also the best place we can do it
                 // to avoid putting it inside inner loops.
                 // NOTE: It assumes that the levels of the input tensor are
                 // initialized in order (and it is also currently guaranteed by
                 // computeIterationGraph), another more admissible approach
                 // might be accepting out-of-order access between consecutive
                 // dense levels.
-                affineTids.push_back(tid);
-                affineLvls.push_back(l);
+                affineTidLvls.emplace_back(env.makeTensorLevel(tid, l));
                 exps.push_back(exp);
               }
             }
           }
         }
       });
 
   if (isDenseDLT(env.dlt(outTid, ldx))) {
     // Note that we generate dense indices of the output tensor
     // unconditionally, since they may not appear in the lattice, but may be
     // needed for linearized env.
-    tids.push_back(outTid);
-    lvls.push_back(*outLvl);
+    tidLvls.emplace_back(env.makeTensorLevel(outTid, *outLvl));
   }
 
   assert(numloopCond > 0);
   // If we just need to one loop conditions and the conditions is not imposed on
   // non-unique level, the loop can be generated by a for loop.
   return numloopCond == 1 && !hasNonUnique;
 }
 
 /// Starts a single loop in current sequence.
 static std::pair<Operation *, bool> startLoop(CodegenEnv &env,
                                               OpBuilder &builder, LoopOrd at,
                                               LatPointId li, bool needsUniv) {
   // The set of tensors + lvls to generate loops on
-  SmallVector<TensorId> tids, affineTids;
-  SmallVector<Level> lvls, affineLvls;
+  SmallVector<TensorLevel> tidLvls;
   // The set of dense tensors with non-trivial affine expression that just
   // becomes invariant and the address shall now be generated at the current
   // level.
+  SmallVector<TensorLevel> affineTidLvls;
   SmallVector<AffineExpr> affines;
   bool isSingleCond = translateBitsToTidLvlPairs(
-      env, li, env.topSortAt(at), tids, lvls, affineTids, affineLvls, affines);
+      env, li, env.topSortAt(at), tidLvls, affineTidLvls, affines);
 
   // Emit the for/while-loop control.
-  Operation *loop =
-      genLoop(env, builder, at, needsUniv, tids, lvls, isSingleCond);
+  Operation *loop = genLoop(env, builder, at, needsUniv, tidLvls, isSingleCond);
   Location loc = env.op().getLoc();
-  for (auto [tid, lvl, exp] : llvm::zip(affineTids, affineLvls, affines)) {
-    env.emitter().genDenseAffineAddress(builder, loc, tid, lvl, exp);
+  for (auto [tidLvl, exp] : llvm::zip(affineTidLvls, affines)) {
+    env.emitter().genDenseAffineAddress(builder, loc, tidLvl, exp);
   }
 
   // Until now, we have entered every <tid, lvl> pair in {cond, extra,
   // affine}Tids/Lvls. The addresses of the upcoming levels which are dependent
   // on constant affines expression may now be determined.
-  auto allTids = llvm::concat<TensorId>(tids, affineTids);
-  auto allLvls = llvm::concat<Level>(lvls, affineLvls);
-  for (auto [tid, lvl] : llvm::zip(allTids, allLvls)) {
+  auto allTidLvls = llvm::concat<TensorLevel>(tidLvls, affineTidLvls);
+  for (TensorLevel tidLvl : allTidLvls) {
+    auto [tid, lvl] = env.toTidLvlPair(tidLvl);
     if (tid != env.merger().getOutTensorID())
       genConstantDenseAddressFromLevel(env, builder, tid, lvl + 1);
   }
 
   return std::make_pair(loop, isSingleCond);
 }
 
 /// Ends a single loop in current sequence. Returns new values for needsUniv.