[Mlir-commits] [mlir] f9d5053 - [mlir][sparse] updating methods to get MLIRContext

[wren romano]

Author: wren romano
Date: 2023-08-04T17:56:09-07:00
New Revision: f9d505316bb40339010a95d1f68dba83252e708d

URL: https://github.com/llvm/llvm-project/commit/f9d505316bb40339010a95d1f68dba83252e708d
DIFF: https://github.com/llvm/llvm-project/commit/f9d505316bb40339010a95d1f68dba83252e708d.diff

LOG: [mlir][sparse] updating methods to get MLIRContext

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`.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D157163

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.h b/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.h
index c39cd9a3e96f75..5552e2fe0fd13d 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.h
+++ b/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.h
@@ -93,7 +93,9 @@ class DimLvlExpr {
     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.
@@ -209,6 +211,8 @@ class DimSpec final {
 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; }
@@ -262,6 +266,12 @@ class LvlSpec final {
 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; }