[mlir][sparse] Adding `LvlVar` forward-declarations to `DimLvlMap::print`

[wrengr]

Summary

This commit makes DimLvlMap::print consistent with DimLvlMapParser with respect to both the forward-decls per se, as well as the all-or-none constraint on LvlVar-bindings in LvlSpecs.

[wrengr]

git-clang-format

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

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

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

 }
 
 //===----------------------------------------------------------------------===//
 // `DimLvlMap` implementation.
 //===----------------------------------------------------------------------===//
 
 DimLvlMap::DimLvlMap(unsigned symRank, ArrayRef<DimSpec> dimSpecs,
                      ArrayRef<LvlSpec> lvlSpecs)
-    : symRank(symRank), dimSpecs(dimSpecs), lvlSpecs(lvlSpecs) {
+    : symRank(symRank), dimSpecs(dimSpecs), lvlSpecs(lvlSpecs),
+      mustPrintLvlVars(false) {
   // First, check integrity of the variable-binding structure.
   // NOTE: This establishes the invariant that calls to `VarSet::add`
   // below cannot cause OOB errors.
   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.
 
   // 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())
       usedVars.add(dimSpec.getExpr());
-  for (auto &lvlSpec : this->lvlSpecs)
-    lvlSpec.setElideVar(!usedVars.contains(lvlSpec.getBoundVar()));
+  for (auto &lvlSpec : this->lvlSpecs) {
+    // Is this LvlVar used in any overt expression?
+    const bool isUsed = usedVars.contains(lvlSpec.getBoundVar());
+    // This LvlVar can be elided iff it isn't overtly used.
+    lvlSpec.setElideVar(!isUsed);
+    // If any LvlVar cannot be elided, then must forward-declare all LvlVars.
+    mustPrintLvlVars = mustPrintLvlVars || isUsed;
+  }
 }
 
 bool DimLvlMap::isWF() const {
   const auto ranks = getRanks();
   unsigned dimNum = 0;
   for (const auto &dimSpec : dimSpecs)
     if (dimSpec.getBoundVar().getNum() != dimNum++ || !dimSpec.isValid(ranks))
       return false;
void DimLvlMap::print(llvm::raw_ostream &os, bool wantElision) const {
   // whereas only LvlVars have different binding rules).
   if (symRank != 0) {
     os << "[s0";
     for (unsigned i = 1; i < symRank; ++i)
       os << ", s" << i;
     os << ']';
   }
 
+  // LvlVar forward-declarations.
+  if (mustPrintLvlVars) {
+    os << '{';
+    llvm::interleaveComma(
+        lvlSpecs, os, [&](LvlSpec const &spec) { os << spec.getBoundVar(); });
+    os << '}';
+  }
+
   // Dimension specifiers.
   os << '(';
   llvm::interleaveComma(
       dimSpecs, os, [&](DimSpec const &spec) { spec.print(os, wantElision); });
   os << ") -> (";
   // Level specifiers.
+  wantElision = wantElision && !mustPrintLvlVars;
   llvm::interleaveComma(
       lvlSpecs, os, [&](LvlSpec const &spec) { spec.print(os, wantElision); });
   os << ')';
 }
 
 //===----------------------------------------------------------------------===//