[mlir][sparse] Adding getters/setters to `DimLvlMap`

[wrengr]

Summary

Depends On D156768

mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.h

public:
   void dump() const;
 };
 // Although this class is more than just a newtype/wrapper, we do want
 // to ensure that storing them into `SmallVector` is efficient.
 static_assert(IsZeroCostAbstraction<LvlSpec>);
 
 //===----------------------------------------------------------------------===//
 class DimLvlMap final {
-  // TODO(wrengr): Need to define getters
-  unsigned symRank;
-  SmallVector<DimSpec> dimSpecs;
-  SmallVector<LvlSpec> lvlSpecs;
-  bool mustPrintLvlVars;
-
-  // Checks for integrity of variable-binding structure.
-  // This is already called by the ctor.
-  [[nodiscard]] bool isWF() const;
-
 public:
   DimLvlMap(unsigned symRank, ArrayRef<DimSpec> dimSpecs,
             ArrayRef<LvlSpec> lvlSpecs);
 
   unsigned getSymRank() const { return symRank; }
   unsigned getDimRank() const { return dimSpecs.size(); }
   unsigned getLvlRank() const { return lvlSpecs.size(); }
   unsigned getRank(VarKind vk) const { return getRanks().getRank(vk); }
   Ranks getRanks() const { return {getSymRank(), getDimRank(), getLvlRank()}; }
 
-  DimLevelType getDimLevelType(unsigned i) { return lvlSpecs[i].getType(); }
+  ArrayRef<DimSpec> getDims() const { return dimSpecs; }
+  const DimSpec &getDim(Dimension dim) const { return dimSpecs[dim]; }
+  SparseTensorDimSliceAttr getDimSlice(Dimension dim) const {
+    return getDim(dim).getSlice();
+  }
+
+  ArrayRef<LvlSpec> getLvls() const { return lvlSpecs; }
+  const LvlSpec &getLvl(Level lvl) const { return lvlSpecs[lvl]; }
+  DimLevelType getLvlType(Level lvl) const { return getLvl(lvl).getType(); }
+
+  AffineMap getDimToLvlMap(MLIRContext *context) const;
+  AffineMap getLvlToDimMap(MLIRContext *context) const;
 
   void print(llvm::raw_ostream &os, bool wantElision = true) const;
   void print(AsmPrinter &printer, bool wantElision = true) const;
   void dump() const;
+
+private:
+  /// Checks for integrity of variable-binding structure.
+  /// This is already called by the ctor.
+  [[nodiscard]] bool isWF() const;
+
+  /// Helper function to call `DimSpec::setExpr` while asserting that
+  /// the invariant established by `DimLvlMap:isWF` is maintained.
+  /// This is used by the ctor.
+  void setDimExpr(Dimension dim, DimExpr expr) {
+    assert(expr && getRanks().isValid(expr));
+    dimSpecs[dim].setExpr(expr);
+  }
+
+  // All these fields are const-after-ctor.
+  unsigned symRank;
+  SmallVector<DimSpec> dimSpecs;
+  SmallVector<LvlSpec> lvlSpecs;
+  bool mustPrintLvlVars;
 };
 
 //===----------------------------------------------------------------------===//
 
 } // namespace ir_detail
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_IR_DETAIL_DIMLVLMAP_H

mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp

DimLvlMap::DimLvlMap(unsigned symRank, ArrayRef<DimSpec> dimSpecs,
   assert(isWF());
 
   // TODO: Second, we need to infer/validate the `lvlToDim` mapping.
   // Along the way we should set every `DimSpec::elideExpr` according
   // to whether the given expression is inferable or not.  Notably, this
   // needs to happen before the code for setting every `LvlSpec::elideVar`,
   // since if the LvlVar is only used in elided DimExpr, then the
   // LvlVar should also be elided.
-  // NOTE: Whenever we set a new DimExpr, we must make sure to validate it
-  // against our ranks, to restore the invariant established by `isWF` above.
-  // TODO(wrengr): We might should adjust the `DimLvlExpr` ctor to take a
-  // `Ranks` argument and perform the validation then.
+  // NOTE: Be sure to use `DimLvlMap::setDimExpr` for setting the new exprs,
+  // to ensure that we maintain the invariant established by `isWF` above.
 
   // Third, we set every `LvlSpec::elideVar` according to whether that
   // LvlVar occurs in a non-elided DimExpr (TODO: or CountingExpr).
   // NOTE: The invariant established by `isWF` ensures that the following
   // calls to `VarSet::add` cannot raise OOB errors.
   VarSet usedVars(getRanks());
   for (const auto &dimSpec : dimSpecs)
     if (!dimSpec.canElideExpr())
bool DimLvlMap::isWF() const {
   unsigned lvlNum = 0;
   for (const auto &lvlSpec : lvlSpecs)
     if (lvlSpec.getBoundVar().getNum() != lvlNum++ || !lvlSpec.isValid(ranks))
       return false;
   assert(lvlNum == ranks.getLvlRank());
   return true;
 }
 
+AffineMap DimLvlMap::getDimToLvlMap(MLIRContext *context) const {
+  SmallVector<AffineExpr> lvlAffines;
+  lvlAffines.reserve(getLvlRank());
+  for (const auto &lvlSpec : lvlSpecs)
+    lvlAffines.push_back(lvlSpec.getExpr().getAffineExpr());
+  return AffineMap::get(getDimRank(), getSymRank(), lvlAffines, context);
+}
+
+AffineMap DimLvlMap::getLvlToDimMap(MLIRContext *context) const {
+  SmallVector<AffineExpr> dimAffines;
+  dimAffines.reserve(getDimRank());
+  for (const auto &dimSpec : dimSpecs)
+    dimAffines.push_back(dimSpec.getExpr().getAffineExpr());
+  return AffineMap::get(getLvlRank(), getSymRank(), dimAffines, context);
+}
+
 void DimLvlMap::dump() const {
   print(llvm::errs(), /*wantElision=*/false);
   llvm::errs() << "\n";
 }
 void DimLvlMap::print(AsmPrinter &printer, bool wantElision) const {
   print(printer.getStream(), wantElision);
 }
 void DimLvlMap::print(llvm::raw_ostream &os, bool wantElision) const {

mlir/lib/Dialect/SparseTensor/IR/Detail/Var.cpp

case AffineExprKind::CeilDiv: {
     (void)op;
     return occursIn(lhs) || occursIn(rhs);
   }
   }
   llvm_unreachable("unknown AffineExprKind");
 }
 
 void VarSet::add(Var var) {
-  // NOTE: `SmallBitVactor::operator[]` will raise assertion errors for OOB.
+  // NOTE: `SmallBitVector::operator[]` will raise assertion errors for OOB.
   impl[var.getKind()][var.getNum()] = true;
 }
 
 void VarSet::add(VarSet const &other) {
   // NOTE: `SmallBitVector::operator&=` will implicitly resize
   // the bitvector (unlike `BitVector::operator&=`), so we add an
   // assertion against OOB for consistency with the implementation
   // of `VarSet::add(Var)`.

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

if (attrName == "lvlTypes") {
       // examples over to the new surface syntax before making the complete
       // switch once work is completed.
       // TODO: replace everything here with new STEA surface syntax parser
       ir_detail::DimLvlMapParser cParser(parser);
       auto res = cParser.parseDimLvlMap();
       RETURN_ON_FAIL(res);
       // Proof of concept result.
       // TODO: use DimLvlMap directly as storage representation
-      for (unsigned i = 0, e = res->getLvlRank(); i < e; i++)
-        lvlTypes.push_back(res->getDimLevelType(i));
+      for (Level lvl = 0, lvlRank = res->getLvlRank(); lvl < lvlRank; lvl++)
+        lvlTypes.push_back(res->getLvlType(lvl));
     }
 
     // Only the last item can omit the comma
     if (parser.parseOptionalComma().failed())
       break;
   }
 
   RETURN_ON_FAIL(parser.parseRBrace())