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

[ftynse]

Summary

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.

[bondhugula]

There was an assumption that valid IR would always live inside a ModuleOp, which has the AffineScope trait. This isn't true any more? Are you referring to unlinked blocks as valid IR?

[mehdi_amini]

In D81753#2090541, @bondhugula wrote:

There was an assumption that valid IR would always live inside a ModuleOp, which has the AffineScope trait. This isn't true any more? Are you referring to unlinked blocks as valid IR?

We don't enforce a top-level module at the moment (you can construct such IR and it'll pass the verifier).

There is such an example in-tree with the SPIRVDeserializer which produces a SPIRVModuleOp that isn't wrapped in a ModuleOp.

[rriddle]

In D81753#2091045, @mehdi_amini wrote:

In D81753#2090541, @bondhugula wrote:

There was an assumption that valid IR would always live inside a ModuleOp, which has the AffineScope trait. This isn't true any more? Are you referring to unlinked blocks as valid IR?

We don't enforce a top-level module at the moment (you can construct such IR and it'll pass the verifier).

Verifier is a bit overloaded, because it isn't clear what you mean. The only reason we have a "top-level" verify method is because dominance was previously out of IR/. Given that you can't run passes without a Module, and many other parts of the infra explicitly assume a top-level module. I'd say it is a relatively fair assumption to make.

There is such an example in-tree with the SPIRVDeserializer which produces a SPIRVModuleOp that isn't wrapped in a ModuleOp.

Realistically, that could just a take a builder instead of returning a SPIRVModuleOp.

[ftynse]

I've hit both cases:

• upstream, we call into these functions when building IR in a detached region;
• in an out-of-tree project, we have a custom top-level op.

The change makes both cases work conservatively and doesn't break anything upstream.

mlir/docs/Dialects/Affine.md

 ```
 
 ### Restrictions on Dimensions and Symbols
 
 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),
 5. the result of an [`affine.apply`
 operation](#affineapply-operation) that recursively takes as arguments any valid
 symbolic identifiers, or
 6. the result of a [`dim` operation](Standard.md#dim-operation) on either a
 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
 symbol is bound to, but also to induction variables of enclosing
 [`affine.for`](#affinefor-operation) and
 [`affine.parallel`](#affineparallel-operation) operations, and the result of an
 [`affine.apply` operation](#affineapply-operation) (which recursively may use
 other dimensions and symbols).

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

 /// uses in that region.
 static bool isTopLevelValue(Value value, Region *region) {
   if (auto arg = value.dyn_cast<BlockArgument>())
     return arg.getParentRegion() == region;
   return value.getDefiningOp()->getParentRegion() == 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;
   while (auto *parentOp = curOp->getParentOp()) {
     if (parentOp->hasTrait<OpTrait::AffineScope>())
       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
 // conditions:
 // *) It is valid as a symbol.
 // *) It is an induction variable.
 // *) It is the result of affine apply operation with dimension id arguments.
 bool mlir::isValidDim(Value value) {
if (isTopLevelValue(value))
     return true;
 
   if (auto *defOp = value.getDefiningOp())
     return isValidSymbol(value, getAffineScope(defOp));
 
   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;
   }
 
   // Constant operation is ok.
   Attribute operandCst;
   if (matchPattern(defOp, m_Constant(&operandCst)))
     return true;
 
   // Affine apply operation is ok if all of its operands are ok.
   if (auto applyOp = dyn_cast<AffineApplyOp>(defOp))
     return applyOp.isValidSymbol(region);
 
   // Dim op results could be valid symbols at any level.
   if (auto dimOp = dyn_cast<DimOp>(defOp))
     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);
 
   return false;
 }
 
 // Returns true if 'value' is a valid index to an affine operation (e.g.
 // affine.load, affine.store, affine.dma_start, affine.dma_wait) where