[mlir] Affine dim and symbol checkers: support detached regions

Functions checking whether an SSA value is a valid dimension or symbol for
affine operations can be called on values defined in a detached region (a
region that is not yet attached to an operation), for example, during parsing
or operation construction. These functions will attempt to uncondtionally
dereference a pointer to the parent operation of a region, which may be null
(as fixed by the previous commit, uninitialized before that). Since one cannot
know to which operation a region will be attached, conservatively this
operation would not be a valid affine scope and act accordingly, instead of
crashing.

1 files changed, 9 insertions, 5 deletions

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index a07c3fa..1766728 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -88,9 +88,12 @@ Operation *AffineDialect::materializeConstant(OpBuilder &builder,
 /// op with trait `AffineScope`. A value of index type defined at the top
 /// level is always a valid symbol.
 bool mlir::isTopLevelValue(Value value) {
-  if (auto arg = value.dyn_cast<BlockArgument>())
-    return arg.getOwner()->getParentOp()->hasTrait<OpTrait::AffineScope>();
-  return value.getDefiningOp()->getParentOp()->hasTrait<OpTrait::AffineScope>();
+  if (auto arg = value.dyn_cast<BlockArgument>()) {
+    Operation *parentOp = arg.getOwner()->getParentOp();
+    return parentOp && parentOp->hasTrait<OpTrait::AffineScope>();
+  }
+  Operation *parentOp = value.getDefiningOp()->getParentOp();
+  return parentOp && parentOp->hasTrait<OpTrait::AffineScope>();
 }
 
 /// A utility function to check if a value is defined at the top level of
@@ -132,8 +135,9 @@ bool mlir::isValidDim(Value value) {
   // This value has to be a block argument for an op that has the
   // `AffineScope` trait or for an affine.for or affine.parallel.
   auto *parentOp = value.cast<BlockArgument>().getOwner()->getParentOp();
-  return parentOp->hasTrait<OpTrait::AffineScope>() ||
-         isa<AffineForOp>(parentOp) || isa<AffineParallelOp>(parentOp);
+  return parentOp &&
+         (parentOp->hasTrait<OpTrait::AffineScope>() ||
+          isa<AffineForOp>(parentOp) || isa<AffineParallelOp>(parentOp));
 }
 
 // Value can be used as a dimension id iff it meets one of the following