[mlir][sparse] Add more helper methods for converting DimLvlExpr to Var

[wrengr]

Summary

These new methods help clean up some code for doing LvlExpr-analysis during DimExpr-inference.

[aartbik]

rebased with main

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

public:
   //
   // Getters for handling `AffineExpr` subclasses.
   //
   // TODO(wrengr): is there any way to make these typesafe without too much
   // templating?
   // TODO(wrengr): Most if not all of these don't actually need to be
   // methods, they could be free-functions instead.
   //
+  Var castAnyVar() const;
+  std::optional<Var> dyn_castAnyVar() const;
   SymVar castSymVar() const;
+  std::optional<SymVar> dyn_castSymVar() const;
   Var castDimLvlVar() const;
+  std::optional<Var> dyn_castDimLvlVar() const;
   int64_t castConstantValue() const;
-  std::optional<int64_t> tryGetConstantValue() const;
+  std::optional<int64_t> dyn_castConstantValue() const;
   bool hasConstantValue(int64_t val) const;
   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;
class DimExpr final : public DimLvlExpr {
   constexpr explicit DimExpr(DimLvlExpr expr) : DimLvlExpr(expr) {}
 
 public:
   static constexpr ExprKind Kind = ExprKind::Dimension;
   static constexpr bool classof(DimLvlExpr const *expr) {
     return expr->getExprKind() == Kind;
   }
   constexpr explicit DimExpr(AffineExpr expr) : DimLvlExpr(Kind, expr) {}
+
+  LvlVar castLvlVar() const { return castDimLvlVar().cast<LvlVar>(); }
+  std::optional<LvlVar> dyn_castLvlVar() const {
+    const auto var = dyn_castDimLvlVar();
+    return var ? std::make_optional(var->cast<LvlVar>()) : std::nullopt;
+  }
 };
 static_assert(IsZeroCostAbstraction<DimExpr>);
 
 class LvlExpr final : public DimLvlExpr {
   // FIXME(wrengr): These two are needed for the current RTTI implementation.
   friend class DimLvlExpr;
   constexpr explicit LvlExpr(DimLvlExpr expr) : DimLvlExpr(expr) {}
 
 public:
   static constexpr ExprKind Kind = ExprKind::Level;
   static constexpr bool classof(DimLvlExpr const *expr) {
     return expr->getExprKind() == Kind;
   }
   constexpr explicit LvlExpr(AffineExpr expr) : DimLvlExpr(Kind, expr) {}
+
+  DimVar castDimVar() const { return castDimLvlVar().cast<DimVar>(); }
+  std::optional<DimVar> dyn_castDimVar() const {
+    const auto var = dyn_castDimLvlVar();
+    return var ? std::make_optional(var->cast<DimVar>()) : std::nullopt;
+  }
 };
 static_assert(IsZeroCostAbstraction<LvlExpr>);
 
 // FIXME(wrengr): See comments elsewhere re RTTI implementation issues/questions
 template <typename U>
 constexpr bool DimLvlExpr::isa() const {
   if constexpr (std::is_same_v<U, DimExpr>)
     return getExprKind() == ExprKind::Dimension;

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

 using namespace mlir;
 using namespace mlir::sparse_tensor;
 using namespace mlir::sparse_tensor::ir_detail;
 
 //===----------------------------------------------------------------------===//
 // `DimLvlExpr` implementation.
 //===----------------------------------------------------------------------===//
 
+Var DimLvlExpr::castAnyVar() const {
+  assert(expr && "uninitialized DimLvlExpr");
+  const auto var = dyn_castAnyVar();
+  assert(var && "expected DimLvlExpr to be a Var");
+  return *var;
+}
+
+std::optional<Var> DimLvlExpr::dyn_castAnyVar() const {
+  if (const auto s = expr.dyn_cast_or_null<AffineSymbolExpr>())
+    return SymVar(s);
+  if (const auto x = expr.dyn_cast_or_null<AffineDimExpr>())
+    return Var(getAllowedVarKind(), x);
+  return std::nullopt;
+}
+
 SymVar DimLvlExpr::castSymVar() const {
   return SymVar(expr.cast<AffineSymbolExpr>());
 }
 
+std::optional<SymVar> DimLvlExpr::dyn_castSymVar() const {
+  if (const auto s = expr.dyn_cast_or_null<AffineSymbolExpr>())
+    return SymVar(s);
+  return std::nullopt;
+}
+
 Var DimLvlExpr::castDimLvlVar() const {
   return Var(getAllowedVarKind(), expr.cast<AffineDimExpr>());
 }
 
+std::optional<Var> DimLvlExpr::dyn_castDimLvlVar() const {
+  if (const auto x = expr.dyn_cast_or_null<AffineDimExpr>())
+    return Var(getAllowedVarKind(), x);
+  return std::nullopt;
+}
+
 int64_t DimLvlExpr::castConstantValue() const {
   return expr.cast<AffineConstantExpr>().getValue();
 }
 
-std::optional<int64_t> DimLvlExpr::tryGetConstantValue() const {
+std::optional<int64_t> DimLvlExpr::dyn_castConstantValue() const {
   const auto k = expr.dyn_cast_or_null<AffineConstantExpr>();
   return k ? std::make_optional(k.getValue()) : std::nullopt;
 }
 
 // This helper method is akin to `AffineExpr::operator==(int64_t)`
 // except it uses a different implementation, namely the implementation
 // used within `AsmPrinter::Impl::printAffineExprInternal`.
 //

 static std::optional<MatchNeg> matchNeg(DimLvlExpr expr) {
   const auto [lhs, op, rhs] = expr.unpackBinop();
   if (op == AffineExprKind::Constant) {
     const auto val = expr.castConstantValue();
     if (val < 0)
       return MatchNeg{DimLvlExpr{expr.getExprKind(), AffineExpr()}, val};
   }
   if (op == AffineExprKind::Mul)
-    if (const auto rval = rhs.tryGetConstantValue(); rval && *rval < 0)
+    if (const auto rval = rhs.dyn_castConstantValue(); rval && *rval < 0)
       return MatchNeg{lhs, *rval};
   return std::nullopt;
 }
 
 // A heavily revised version of `AsmPrinter::Impl::printAffineExprInternal`.
 void DimLvlExpr::printAffineExprInternal(
     llvm::raw_ostream &os, BindingStrength enclosingTightness) const {
   const char *binopSpelling = nullptr;