[mlir] Affine symbols: do not expect AffineScope to always exist

In the affine symbol and dimension check, the code currently assumes
`getAffineScope` and its users `isValidDim` and `isValidSymbol` are only called
on values defined in regions that have a parent Op with `AffineScope` trait.
This is not necessarily the case, and these functions may be called on valid IR
that does not satisfy this assumption. Return `nullptr` from `getAffineScope`
if there is no parent op with `AffineScope` trait. Treat this case
conservatively in `isValidSymbol` by only accepting as symbols the values that
are guaranteed to be symbols (constants, and certain operations). No
modifications are necessary to `isValidDim` that delegates most of the work to
`isValidDim`.

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

mlir/docs/Dialects/Affine.md

@@ -63,7 +63,7 @@ The affine dialect imposes certain restrictions on dimension and symbolic
 identifiers to enable powerful analysis and transformation. An SSA value's use
 can be bound to a symbolic identifier if that SSA value is either
 1. a region argument for an op with trait `AffineScope` (eg. `FuncOp`),
-2. a value defined at the top level of a `AffineScope` op (i.e., immediately
+2. a value defined at the top level of an `AffineScope` op (i.e., immediately
 enclosed by the latter),
 3. a value that dominates the `AffineScope` op enclosing the value's use,
 4. the result of a [`constant` operation](Standard.md#constant-operation),
@@ -74,6 +74,8 @@ symbolic identifiers, or
 memref that is an argument to a `AffineScope` op or a memref where the
 corresponding dimension is either static or a dynamic one in turn bound to a
 valid symbol.
+*Note:* if the use of an SSA value is not contained in any op with the
+`AffineScope` trait, only the rules 4-6 can be applied.
 
 Note that as a result of rule (3) above, symbol validity is sensitive to the
 location of the SSA use.  Dimensions may be bound not only to anything that a

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -113,7 +113,7 @@ static bool isTopLevelValue(Value value, Region *region) {
 }
 
 /// Returns the closest region enclosing `op` that is held by an operation with
-/// trait `AffineScope`.
+/// trait `AffineScope`; `nullptr` if there is no such region.
 //  TODO: getAffineScope should be publicly exposed for affine passes/utilities.
 static Region *getAffineScope(Operation *op) {
   auto *curOp = op;
@@ -122,7 +122,7 @@ static Region *getAffineScope(Operation *op) {
       return curOp->getParentRegion();
     curOp = parentOp;
   }
-  llvm_unreachable("op doesn't have an enclosing polyhedral scope");
+  return nullptr;
 }
 
 // A Value can be used as a dimension id iff it meets one of the following
@@ -236,28 +236,31 @@ bool mlir::isValidSymbol(Value value) {
   return false;
 }
 
-// A value can be used as a symbol for `region` iff it meets onf of the the
-// following conditions:
-// *) It is a constant.
-// *) It is defined at the top level of 'region' or is its argument.
-// *) It dominates `region`'s parent op.
-// *) It is the result of an affine apply operation with symbol arguments.
-// *) It is a result of the dim op on a memref whose corresponding size is
-//    a valid symbol.
+/// A value can be used as a symbol for `region` iff it meets onf of the the
+/// following conditions:
+/// *) It is a constant.
+/// *) It is the result of an affine apply operation with symbol arguments.
+/// *) It is a result of the dim op on a memref whose corresponding size is
+///    a valid symbol.
+/// *) It is defined at the top level of 'region' or is its argument.
+/// *) It dominates `region`'s parent op.
+/// If `region` is null, conservatively assume the symbol definition scope does
+/// not exist and only accept the values that would be symbols regardless of
+/// the surrounding region structure, i.e. the first three cases above.
 bool mlir::isValidSymbol(Value value, Region *region) {
   // The value must be an index type.
   if (!value.getType().isIndex())
     return false;
 
   // A top-level value is a valid symbol.
-  if (::isTopLevelValue(value, region))
+  if (region && ::isTopLevelValue(value, region))
     return true;
 
   auto *defOp = value.getDefiningOp();
   if (!defOp) {
     // A block argument that is not a top-level value is a valid symbol if it
     // dominates region's parent op.
-    if (!region->getParentOp()->isKnownIsolatedFromAbove())
+    if (region && !region->getParentOp()->isKnownIsolatedFromAbove())
       if (auto *parentOpRegion = region->getParentOp()->getParentRegion())
         return isValidSymbol(value, parentOpRegion);
     return false;
@@ -277,7 +280,7 @@ bool mlir::isValidSymbol(Value value, Region *region) {
     return isDimOpValidSymbol(dimOp, region);
 
   // Check for values dominating `region`'s parent op.
-  if (!region->getParentOp()->isKnownIsolatedFromAbove())
+  if (region && !region->getParentOp()->isKnownIsolatedFromAbove())
     if (auto *parentRegion = region->getParentOp()->getParentRegion())
       return isValidSymbol(value, parentRegion);