[mlir][sparse] Adding `{Var,DimLvlExpr,DimSpec,LvlSpec,DimLvlMap}::str` methods

[wrengr]

Summary

These methods are needed for use with Diagnostic::operator<< etc.

The definitions follow the pattern of Diagnostic::str by simply wrapping the underlying print(raw_ostream) method. Although there is some overhead for constructing the std::string, this seems like the overall most-efficient option: since this overhead only occurs on the error path (under the current intended usage). An alternative approach would be to have one method construct a Twine directly, and then have the print method pass the twine to the stream; however, that would mean introducing the overhead of twine construction on the common/happy path of simply printing things out.

[aartbik]

rebased with main

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

public:
   DimLvlExpr getLHS() const;
   DimLvlExpr getRHS() const;
   std::tuple<DimLvlExpr, AffineExprKind, DimLvlExpr> unpackBinop() const;
 
   /// Checks whether the variables bound/used by this spec are valid
   /// with respect to the given ranks.
   [[nodiscard]] bool isValid(Ranks const &ranks) const;
 
+  std::string str() const;
   void print(llvm::raw_ostream &os) const;
   void print(AsmPrinter &printer) const;
   void dump() const;
 
 protected:
   // Variant of `mlir::AsmPrinter::Impl::BindingStrength`
   enum class BindingStrength : bool { Weak = false, Strong = true };
 
public:
   /// the result of this predicate.
   [[nodiscard]] bool isValid(Ranks const &ranks) const;
 
   // TODO(wrengr): Use it or loose it.
   bool isFunctionOf(Var var) const;
   bool isFunctionOf(VarSet const &vars) const;
   void getFreeVars(VarSet &vars) const;
 
+  std::string str(bool wantElision = true) const;
   void print(llvm::raw_ostream &os, bool wantElision = true) const;
   void print(AsmPrinter &printer, bool wantElision = true) const;
   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<DimSpec>);
 
public:
   // a more sophisticated implementation than `DimSpec::isValid` does.
   [[nodiscard]] bool isValid(Ranks const &ranks) const;
 
   // TODO(wrengr): Use it or loose it.
   bool isFunctionOf(Var var) const;
   bool isFunctionOf(VarSet const &vars) const;
   void getFreeVars(VarSet &vars) const;
 
+  std::string str(bool wantElision = true) const;
   void print(llvm::raw_ostream &os, bool wantElision = true) const;
   void print(AsmPrinter &printer, bool wantElision = true) const;
   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>);
 
public:
 
   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;
 
+  std::string str(bool wantElision = true) 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;

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

DimLvlExpr::unpackBinop() const {
   const DimLvlExpr rhs(kind, binop ? binop.getRHS() : nullptr);
   return {lhs, ak, rhs};
 }
 
 void DimLvlExpr::dump() const {
   print(llvm::errs());
   llvm::errs() << "\n";
 }
+std::string DimLvlExpr::str() const {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  print(os);
+  return os.str();
+}
 void DimLvlExpr::print(AsmPrinter &printer) const {
   print(printer.getStream());
 }
 void DimLvlExpr::print(llvm::raw_ostream &os) const {
   if (!expr)
     os << "<<NULL AFFINE EXPR>>";
   else
     printWeak(os);

 }
 
 void DimSpec::getFreeVars(VarSet &vars) const { vars.add(expr); }
 
 void DimSpec::dump() const {
   print(llvm::errs(), /*wantElision=*/false);
   llvm::errs() << "\n";
 }
+std::string DimSpec::str(bool wantElision) const {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  print(os, wantElision);
+  return os.str();
+}
 void DimSpec::print(AsmPrinter &printer, bool wantElision) const {
   print(printer.getStream(), wantElision);
 }
 void DimSpec::print(llvm::raw_ostream &os, bool wantElision) const {
   os << var;
   if (expr && (!wantElision || !elideExpr))
     os << " = " << expr;
   if (slice) {

 }
 
 void LvlSpec::getFreeVars(VarSet &vars) const { vars.add(expr); }
 
 void LvlSpec::dump() const {
   print(llvm::errs(), /*wantElision=*/false);
   llvm::errs() << "\n";
 }
+std::string LvlSpec::str(bool wantElision) const {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  print(os, wantElision);
+  return os.str();
+}
 void LvlSpec::print(AsmPrinter &printer, bool wantElision) const {
   print(printer.getStream(), wantElision);
 }
 void LvlSpec::print(llvm::raw_ostream &os, bool wantElision) const {
   if (!wantElision || !elideVar)
     os << var << " = ";
   os << expr;
   os << ": " << toMLIRString(type);
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";
 }
+std::string DimLvlMap::str(bool wantElision) const {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  print(os, wantElision);
+  return os.str();
+}
 void DimLvlMap::print(AsmPrinter &printer, bool wantElision) const {
   print(printer.getStream(), wantElision);
 }
 void DimLvlMap::print(llvm::raw_ostream &os, bool wantElision) const {
   // Symbolic identifiers.
   // NOTE: Unlike `AffineMap` we place the SymVar bindings before the DimVar
   // bindings, since the SymVars may occur within DimExprs and thus this
   // ordering helps reduce potential user confusion about the scope of bidings

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

public:
 
   template <typename U>
   constexpr bool isa() const;
   template <typename U>
   constexpr U cast() const;
   template <typename U>
   constexpr std::optional<U> dyn_cast() const;
 
+  std::string str() const;
   void print(llvm::raw_ostream &os) const;
   void print(AsmPrinter &printer) const;
   void dump() const;
 };
 static_assert(IsZeroCostAbstraction<Var>);
 
 class SymVar final : public Var {
   using Var::Var; // inherit `Var(Impl)` ctor for RTTI use.

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

 /// For use in foreach loops.
 static constexpr const VarKind everyVarKind[] = {
     VarKind::Dimension, VarKind::Symbol, VarKind::Level};
 
 //===----------------------------------------------------------------------===//
 // `Var` implementation.
 //===----------------------------------------------------------------------===//
 
+std::string Var::str() const {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  print(os);
+  return os.str();
+}
+
 void Var::print(AsmPrinter &printer) const { print(printer.getStream()); }
 
 void Var::print(llvm::raw_ostream &os) const {
   os << toChar(getKind()) << getNum();
 }
 
 void Var::dump() const {
   print(llvm::errs());