[mlir][sparse] Cleaning up names in {Merger,LoopEmitter,CodegenEnv}.{…

…h,cpp}

This change does a bunch of renaming to clear up confusions in these files.  In particular, this change:

* Renames variables and methods to clarify the "dim"/"lvl" distinction, and changes them to use the `Dimension`/`Level` types as appropriate.
* Introduces new typedefs
  * `ExprId`, `LatPointId`, `LatSetId`: to clarify the interning design of the Merger.
  * `LoopId`, `LoopOrd`: to clarify the distinction between arbitrary names for loop-variables, vs numeric identifiers based on the actual order of loop generation.
  * `TensorId`
  * (Future CLs will change these from typedefs to structs/classes, so that the typechecker can help avoid mixups.)
* Updates documentation to match the new terminology
* Adds additional assertions
* Adds `const` to local variables along the way

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D145756

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorType.h

@@ -61,6 +61,7 @@ class SparseTensorType {
   // Copy-assignment would be implicitly deleted (because our fields
   // are const), so we explicitly delete it for clarity.
   SparseTensorType &operator=(const SparseTensorType &) = delete;
+  // So we must explicitly define the copy-ctor to silence -Wdeprecated-copy.
   SparseTensorType(const SparseTensorType &) = default;
 
   /// Constructs a new `SparseTensorType` with the same dimension-shape

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

@@ -38,12 +38,12 @@ CodegenEnv::CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts,
     : linalgOp(linop), sparseOptions(opts),
       latticeMerger(numTensors, numLoops, numFilterLoops), loopEmitter(),
       topSort(), sparseOut(nullptr), outerParNest(-1u), insChain(), expValues(),
-      expFilled(), expAdded(), expCount(), redVal(), redExp(-1u),
-      redCustom(-1u), redValidLexInsert() {}
+      expFilled(), expAdded(), expCount(), redVal(), redExp(kInvalidId),
+      redCustom(kInvalidId), redValidLexInsert() {}
 
 LogicalResult CodegenEnv::initTensorExp() {
   // Builds the tensor expression for the Linalg operation in SSA form.
-  std::optional<unsigned> optExp = latticeMerger.buildTensorExpFromLinalg(op());
+  std::optional<ExprId> optExp = latticeMerger.buildTensorExpFromLinalg(op());
   if (!optExp || !isAdmissibleTensorExp(*optExp))
     return failure();
 
@@ -101,7 +101,7 @@ std::optional<Operation *> CodegenEnv::genLoopBoundary(
 // Code generation environment verify functions.
 //===----------------------------------------------------------------------===//
 
-bool CodegenEnv::isAdmissibleTensorExp(unsigned exp) {
+bool CodegenEnv::isAdmissibleTensorExp(ExprId exp) {
   // We reject any expression that makes a reduction from `-outTensor`, as those
   // expressions create a dependency between the current iteration (i) and the
   // previous iteration (i-1). It would require iterating over the whole
@@ -115,7 +115,10 @@ bool CodegenEnv::isAdmissibleTensorExp(unsigned exp) {
   }
 
   OpOperand *lhs = linalgOp.getDpsInitOperand(0);
-  unsigned tensor = lhs->getOperandNumber();
+  // That the operand number is a valid `TensorId` will be verified
+  // by the call to `isSingleCondition` below; though we may want to add
+  // assertions to check it here, in order to give better error messages.
+  const TensorId tensor = lhs->getOperandNumber();
   // An non-annotated output tensor is assumed dense, and becomes a random
   // access n-dim memref. Admissible since insertions cannot occur.
   if (getSparseTensorType(lhs->get()).isAllDense())
@@ -140,21 +143,22 @@ bool CodegenEnv::isAdmissibleTopoOrder() {
   OpOperand *lhs = linalgOp.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))) {
+  LoopOrd nest = 0;
+  const auto iteratorTypes = linalgOp.getIteratorTypesArray();
+  assert(topSortSize() == latticeMerger.getNumLoops());
+  for (const LoopId i : topSort) {
+    if (!latticeMerger.isFilterLoop(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)]))
+      // a special type of parallel loops.
+      if (linalg::isReductionIterator(iteratorTypes[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 (static_cast<int64_t>(nest) >= linalgOp.getRank(lhs) - 1) {
     outerParNest = nest;
     return true;
   }
@@ -165,19 +169,26 @@ bool CodegenEnv::isAdmissibleTopoOrder() {
 // Code generation environment topological sort methods
 //===----------------------------------------------------------------------===//
 
-ArrayRef<unsigned> CodegenEnv::getTopSortSlice(size_t n, size_t m) const {
-  return ArrayRef<unsigned>(topSort).slice(n, m);
+ArrayRef<LoopId> CodegenEnv::getTopSortSlice(LoopOrd n, LoopOrd m) const {
+  return ArrayRef<LoopId>(topSort).slice(n, m);
 }
 
-ArrayRef<unsigned> CodegenEnv::getLoopCurStack() const {
-  return getTopSortSlice(0, loopEmitter.getCurrentDepth());
+ArrayRef<LoopId> CodegenEnv::getLoopStackUpTo(LoopOrd n) const {
+  return ArrayRef<LoopId>(topSort).take_front(n);
 }
 
-Value CodegenEnv::getLoopIdxValue(size_t loopIdx) const {
-  for (unsigned lv = 0, lve = topSort.size(); lv < lve; lv++)
-    if (topSort[lv] == loopIdx)
-      return loopEmitter.getLoopIV(lv);
-  llvm_unreachable("invalid loop index");
+ArrayRef<LoopId> CodegenEnv::getCurrentLoopStack() const {
+  return getLoopStackUpTo(loopEmitter.getCurrentDepth());
+}
+
+Value CodegenEnv::getLoopVar(LoopId i) const {
+  // TODO: this class should store the inverse of `topSort` so that
+  // it can do this conversion directly, instead of searching through
+  // `topSort` every time.  (Or else, `LoopEmitter` should handle this.)
+  for (LoopOrd n = 0, numLoops = topSortSize(); n < numLoops; n++)
+    if (topSort[n] == i)
+      return loopEmitter.getLoopIV(n);
+  llvm_unreachable("invalid loop identifier");
 }
 
 //===----------------------------------------------------------------------===//
@@ -189,8 +200,10 @@ void CodegenEnv::updateInsertionChain(Value chain) {
   insChain = chain;
 }
 
-bool CodegenEnv::atExpandLevel(OpOperand *o, unsigned rank, unsigned lv) const {
-  return sparseOut == o && outerParNest == rank - 1 && outerParNest == lv;
+// FIXME: clarify what this "rank" is really supposed to mean/be.
+bool CodegenEnv::atExpandLevel(OpOperand *o, unsigned rank, LoopOrd n) const {
+  return sparseOut == o && outerParNest == static_cast<LoopOrd>(rank - 1) &&
+         outerParNest == n;
 }
 
 void CodegenEnv::startExpand(Value values, Value filled, Value added,
@@ -216,21 +229,21 @@ void CodegenEnv::endExpand() {
 // Code generation environment reduction methods
 //===----------------------------------------------------------------------===//
 
-void CodegenEnv::startReduc(unsigned exp, Value val) {
-  assert(redExp == -1u && exp != -1u);
+void CodegenEnv::startReduc(ExprId exp, Value val) {
+  assert(!isReduc() && exp != kInvalidId);
   redExp = exp;
   updateReduc(val);
 }
 
 void CodegenEnv::updateReduc(Value val) {
-  assert(redExp != -1u);
+  assert(isReduc());
   redVal = exp(redExp).val = val;
 }
 
 Value CodegenEnv::endReduc() {
   Value val = redVal;
   updateReduc(Value());
-  redExp = -1u;
+  redExp = kInvalidId;
   return val;
 }
 
@@ -244,17 +257,17 @@ void CodegenEnv::clearValidLexInsert() {
   redValidLexInsert = Value();
 }
 
-void CodegenEnv::startCustomReduc(unsigned exp) {
-  assert(redCustom == -1u && exp != -1u);
+void CodegenEnv::startCustomReduc(ExprId exp) {
+  assert(!isCustomReduc() && exp != kInvalidId);
   redCustom = exp;
 }
 
 Value CodegenEnv::getCustomRedId() {
-  assert(redCustom != -1u);
+  assert(isCustomReduc());
   return dyn_cast<sparse_tensor::ReduceOp>(exp(redCustom).op).getIdentity();
 }
 
 void CodegenEnv::endCustomReduc() {
-  assert(redCustom != -1u);
-  redCustom = -1u;
+  assert(isCustomReduc());
+  redCustom = kInvalidId;
 }

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

@@ -45,7 +45,7 @@ class CodegenEnv {
   //
 
   LogicalResult initTensorExp();
-  unsigned getTensorExp() const { return tensorExp; }
+  ExprId getExprId() const { return tensorExp; }
 
   linalg::GenericOp op() const { return linalgOp; }
   const SparsificationOptions &options() const { return sparseOptions; }
@@ -65,13 +65,13 @@ class CodegenEnv {
   // Merger delegates.
   //
 
-  TensorExp &exp(unsigned e) { return latticeMerger.exp(e); }
-  LatPoint &lat(unsigned l) { return latticeMerger.lat(l); }
-  SmallVector<unsigned> &set(unsigned s) { return latticeMerger.set(s); }
-  DimLevelType dlt(unsigned t, unsigned i) const {
+  TensorExp &exp(ExprId e) { return latticeMerger.exp(e); }
+  LatPoint &lat(LatPointId l) { return latticeMerger.lat(l); }
+  SmallVector<LatPointId> &set(LatSetId s) { return latticeMerger.set(s); }
+  DimLevelType dlt(TensorId t, LoopId i) const {
     return latticeMerger.getDimLevelType(t, i);
   }
-  DimLevelType dlt(unsigned b) const {
+  DimLevelType dlt(TensorLoopId b) const {
     return latticeMerger.getDimLevelType(b);
   }
 
@@ -81,7 +81,7 @@ class CodegenEnv {
 
   /// Whether the tensor expression is admissible for codegen.
   /// It also sets the sparseOut if the output tensor is sparse.
-  bool isAdmissibleTensorExp(unsigned exp);
+  bool isAdmissibleTensorExp(ExprId e);
 
   /// Whether the iteration graph is sorted in admissible topoOrder.
   /// Sets outerParNest on success with sparse output
@@ -91,17 +91,21 @@ class CodegenEnv {
   // Topological delegate and sort methods.
   //
 
-  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) {
+  LoopOrd topSortSize() const { return topSort.size(); }
+  LoopId topSortAt(LoopOrd n) const { return topSort.at(n); }
+  void topSortPushBack(LoopId i) { topSort.push_back(i); }
+  void topSortClear(size_t capacity = 0) {
     topSort.clear();
     topSort.reserve(capacity);
   }
 
-  ArrayRef<unsigned> getTopSortSlice(size_t n, size_t m) const;
-  ArrayRef<unsigned> getLoopCurStack() const;
-  Value getLoopIdxValue(size_t loopIdx) const;
+  ArrayRef<LoopId> getTopSortSlice(LoopOrd n, LoopOrd m) const;
+  ArrayRef<LoopId> getLoopStackUpTo(LoopOrd n) const;
+  ArrayRef<LoopId> getCurrentLoopStack() const;
+  /// Returns the induction-variable for the loop identified by the given
+  /// `LoopId`.  This method handles application of the topological sort
+  /// in order to convert the `LoopId` into the corresponding `LoopOrd`.
+  Value getLoopVar(LoopId i) const;
 
   //
   // Sparse tensor output and expansion methods.
@@ -113,7 +117,8 @@ class CodegenEnv {
   Value getInsertionChain() const { return insChain; }
   void updateInsertionChain(Value chain);
 
-  bool atExpandLevel(OpOperand *o, unsigned rank, unsigned lv) const;
+  // FIXME: clarify what this "rank" is really supposed to mean/be.
+  bool atExpandLevel(OpOperand *o, unsigned rank, LoopOrd n) const;
   void startExpand(Value values, Value filled, Value added, Value count);
   bool isExpand() const { return expValues != nullptr; }
   void updateExpandCount(Value count);
@@ -127,17 +132,17 @@ class CodegenEnv {
   // Reduction methods.
   //
 
-  void startReduc(unsigned exp, Value val);
-  bool isReduc() const { return redExp != -1u; }
+  void startReduc(ExprId exp, Value val);
+  bool isReduc() const { return redExp != kInvalidId; }
   void updateReduc(Value val);
   Value getReduc() const { return redVal; }
   Value endReduc();
   void setValidLexInsert(Value val);
   void clearValidLexInsert();
   Value getValidLexInsert() const { return redValidLexInsert; }
 
-  void startCustomReduc(unsigned exp);
-  bool isCustomReduc() const { return redCustom != -1u; }
+  void startCustomReduc(ExprId exp);
+  bool isCustomReduc() const { return redCustom != kInvalidId; }
   Value getCustomRedId();
   void endCustomReduc();
 
@@ -154,14 +159,16 @@ class CodegenEnv {
   // Loop emitter helper class.
   LoopEmitter loopEmitter;
 
-  // Topological sort.
-  std::vector<unsigned> topSort;
+  // Topological sort.  This serves as a mapping from `LoopOrd` to `LoopId`
+  // (cf., `getLoopVar` and `topSortAt`).
+  std::vector<LoopId> topSort;
 
   // Sparse tensor as output. Implemented either through direct injective
   // insertion in lexicographic index order or through access pattern
   // expansion in the innermost loop nest (`expValues` through `expCount`).
   OpOperand *sparseOut;
-  unsigned outerParNest;
+  // The count of outer non-filter loops, as defined by `isAdmissibleTopoOrder`.
+  LoopOrd outerParNest;
   Value insChain;
   Value expValues;
   Value expFilled;
@@ -172,16 +179,16 @@ class CodegenEnv {
   // into the merger's expression tree. When the indices of a tensor reduction
   // expression are exhausted, all inner loops can use a scalarized reduction.
   Value redVal;
-  unsigned redExp;
-  unsigned redCustom;
+  ExprId redExp;
+  ExprId redCustom;
 
   // Bookkeeping for lex insertion during reductions. Holds the runtime boolean
   // value of whether any reduction occurred. This is only set during a
   // reduction and cleared once the reduction is finished.
   Value redValidLexInsert;
 
   // The root tensor expression of the kernel.
-  unsigned tensorExp;
+  ExprId tensorExp;
 };
 
 } // namespace sparse_tensor