[mlir][sparse] updating methods to get MLIRContext

[wrengr]

Summary

Renames the old DimLvlExpr::getContext to tryGetContext to make it clear that the result could be a nullptr. Also adds new delegation methods DimSpec::tryGetContext and LvlSpec::getContext.

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

public:
   constexpr VarKind getAllowedVarKind() const {
     return getVarKindAllowedInExpr(kind);
   }
   constexpr AffineExpr getAffineExpr() const { return expr; }
   AffineExprKind getAffineKind() const {
     assert(expr);
     return expr.getKind();
   }
-  MLIRContext *getContext() const { return expr ? expr.getContext() : nullptr; }
+  MLIRContext *tryGetContext() const {
+    return expr ? expr.getContext() : nullptr;
+  }
 
   //
   // 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.
class DimSpec final {
   /// the `DimLvlMap` ctor will reset it as appropriate.
   bool elideExpr = false;
   /// The dimension-slice; optional, default is null.
   SparseTensorDimSliceAttr slice;
 
 public:
   DimSpec(DimVar var, DimExpr expr, SparseTensorDimSliceAttr slice);
 
+  MLIRContext *tryGetContext() const { return expr.tryGetContext(); }
+
   constexpr DimVar getBoundVar() const { return var; }
   bool hasExpr() const { return static_cast<bool>(expr); }
   constexpr DimExpr getExpr() const { return expr; }
   void setExpr(DimExpr newExpr) {
     assert(!hasExpr());
     expr = newExpr;
   }
   constexpr bool canElideExpr() const { return elideExpr; }
class LvlSpec final {
   // to ensure that the pair is well-formed.
   LvlExpr expr;
   /// The level-type (== level-format + lvl-properties).
   DimLevelType type;
 
 public:
   LvlSpec(LvlVar var, LvlExpr expr, DimLevelType type);
 
+  MLIRContext *getContext() const {
+    MLIRContext *ctx = expr.tryGetContext();
+    assert(ctx);
+    return ctx;
+  }
+
   constexpr LvlVar getBoundVar() const { return var; }
   constexpr bool canElideVar() const { return elideVar; }
   void setElideVar(bool b) { elideVar = b; }
   constexpr LvlExpr getExpr() const { return expr; }
   constexpr DimLevelType getType() const { return type; }
 
   /// Checks whether the variables bound/used by this spec are valid
   /// with respect to the given ranks.