[MLIR] Introduce op trait PolyhedralScope (revised)

[bondhugula]

Summary

(A previous version of this, dd2c639c3cd397dfef941186fb85c82e4e918425, was
reverted.)

Introduce op trait PolyhedralScope for ops to define a new scope for
polyhedral optimization / affine dialect purposes, thus generalizing
such scopes beyond FuncOp. Ops to which this trait is attached will
define a new scope for the consideration of SSA values as valid symbols
for the purposes of polyhedral analysis and optimization. Update methods
that check for dim/symbol validity to work based on this trait.

[bondhugula]

ModuleOp is now marked with this trait as well. The trait lives in IR - as discussed here (https://reviews.llvm.org/D78863?id=260322#inline-722790), the trait makes sense even for those that don't depend on the affine dialect, and it's clean/logical to mark FuncOp/ModuleOp with this trait. The cyclic dependence b/w Affine and IR/ is now eliminated.

[bondhugula]

add test case where there is no surrounding func op

[ftynse]

In D79060#2009214, @bondhugula wrote:

ModuleOp is now marked with this trait as well. The trait lives in IR - as discussed here (https://reviews.llvm.org/D78863?id=260322#inline-722790), the trait makes sense even for those that don't depend on the affine dialect, and it's clean/logical to mark FuncOp/ModuleOp with this trait. The cyclic dependence b/w Affine and IR/ is now eliminated.

I don't think we reached final consensus on this point, but I vote for landing this regardless of the current implementation and discussing the layering separately on discourse. (If @rriddle is about to split OpTraits into several parts, all traits will then move in a consistent way).

Please wait for @rriddle to approve.

[bondhugula]

In D79060#2009621, @ftynse wrote:

In D79060#2009214, @bondhugula wrote:

ModuleOp is now marked with this trait as well. The trait lives in IR - as discussed here (https://reviews.llvm.org/D78863?id=260322#inline-722790), the trait makes sense even for those that don't depend on the affine dialect, and it's clean/logical to mark FuncOp/ModuleOp with this trait. The cyclic dependence b/w Affine and IR/ is now eliminated.

I don't think we reached final consensus on this point, but I vote for landing this regardless of the current implementation and discussing the layering separately on discourse. (If @rriddle is about to split OpTraits into several parts, all traits will then move in a consistent way).

Please wait for @rriddle to approve.

I agree that there may be a better place for this trait to reside. But since a split may be in the works for other traits as well, we can have that happen separately first.

[rriddle]

I'm fine with landing for now, but it should really be moved out when possible.

[lattner]

Would it make sense to call this "AffineScope". I don't believe we use the word "Polyhedral" anywhere. This was intentional because of the many preexisting notions of what "Polyhedral" means, but in any case we should be consistent.

[bondhugula]

In D79060#2018246, @lattner wrote:

Would it make sense to call this "AffineScope". I don't believe we use the word "Polyhedral" anywhere. This was intentional because of the many preexisting notions of what "Polyhedral" means, but in any case we should be consistent.

Yes, AffineScope would appear to be more consistent. Sounds good to rename.

mlir/docs/Dialects/Affine.md

 #affine_map2to3 = affine_map<(d0, d1)[s0] -> (d0, d1 + s0, d1 - s0)>
 // Binds %N to the s0 symbol in affine_map2to3.
 %x = alloc()[%N] : memref<40x50xf32, #affine_map2to3>
 ```
 
 ### Restrictions on Dimensions and Symbols
 
 The affine dialect imposes certain restrictions on dimension and symbolic
-identifiers to enable powerful analysis and transformation. A symbolic
-identifier can be bound to an SSA value that is either an argument to the
-function, a value defined at the top level of that function (outside of all
-loops and if operations), the result of a
-[`constant` operation](Standard.md#constant-operation), or the result of an
-[`affine.apply` operation](#affineapply-operation) that recursively takes as
-arguments any symbolic identifiers, or the result of a [`dim`
-operation](Standard.md#dim-operation) on either a memref that is a function
-argument or a memref where the corresponding dimension is either static or a
-dynamic one in turn bound to a symbolic identifier.  Dimensions may be bound not
-only to anything that a symbol is bound to, but also to induction variables of
-enclosing [`affine.for`](#affinefor-affineforop) and
-[`affine.parallel`](#affineparallel-affineparallelop) operations, and the
-result of an
+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 `PolyhedralScope` (eg. `FuncOp`),
+2. a value defined at the top level of a `PolyhedralScope` op (i.e., immediately
+enclosed by the latter),
+3. a value that dominates the `PolyhedralScope` 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 `PolyhedralScope` op or a memref where the
+corresponding dimension is either static or a dynamic one in turn bound to a
+valid symbol.
+
+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).
 
 ### Affine Expressions
 
 Syntax:
 
 ```

mlir/docs/Traits.md

 foo.region_op {
   foo.yield %result : i32
 }
 ```
 
 This trait is an important structural property of the IR, and enables operations
 to have [passes](PassManagement.md) scheduled under them.
 
+
+### PolyhedralScope
+
+*   `OpTrait::PolyhedralScope` -- `PolyhedralScope`
+
+This trait is carried by region holding operations that define a new scope for
+the purposes of polyhedral optimization and the affine dialect in particular.
+Any SSA values of 'index' type that either dominate such operations, or are
+defined at the top-level of such operations, or appear as region arguments for
+such operations automatically become valid symbols for the polyhedral scope
+defined by that operation. As a result, such SSA values could be used as the
+operands or index operands of various affine dialect operations like affine.for,
+affine.load, and affine.store.  The polyhedral scope defined by an operation
+with this trait includes all operations in its region excluding operations that
+are nested inside of other operations that themselves have this trait.
+
 ### Single Block with Implicit Terminator
 
 *   `OpTrait::SingleBlockImplicitTerminator<typename TerminatorOpType>` :
     `SingleBlockImplicitTerminator<string op>`
 
 This trait provides APIs and verifiers for operations with regions that have a
 single block that must terminate with `TerminatorOpType`.
 

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

 class AffineApplyOp;
 class AffineBound;
 class AffineDimExpr;
 class AffineValueMap;
 class AffineTerminatorOp;
 class FlatAffineConstraints;
 class OpBuilder;
 
-/// A utility function to check if a value is defined at the top level of a
-/// function. A value of index type defined at the top level is always a valid
-/// symbol.
+/// A utility function to check if a value is defined at the top level of an
+/// op with trait `PolyhedralScope` or is a region argument for such an op. A
+/// value of index type defined at the top level is always a valid symbol for
+/// all its uses.
 bool isTopLevelValue(Value value);
 
 /// AffineDmaStartOp starts a non-blocking DMA operation that transfers data
 /// from a source memref to a destination memref. The source and destination
 /// memref need not be of the same dimensionality, but need to have the same
 /// elemental type. The operands include the source and destination memref's
 /// each followed by its indices, size of the data transfer in terms of the
 /// number of elements (of the elemental type of the memref), a tag memref with
public:
   static StringRef getTagMapAttrName() { return "tag_map"; }
   static ParseResult parse(OpAsmParser &parser, OperationState &result);
   void print(OpAsmPrinter &p);
   LogicalResult verify();
   LogicalResult fold(ArrayRef<Attribute> cstOperands,
                      SmallVectorImpl<OpFoldResult> &results);
 };
 
-/// Returns true if the given Value can be used as a dimension id.
+/// Returns true if the given Value can be used as a dimension id in the region
+/// of the closest surrounding op that has the trait `PolyhedralScope`.
 bool isValidDim(Value value);
 
-/// Returns true if the given Value can be used as a symbol.
+/// Returns true if the given Value can be used as a dimension id in `region`,
+/// i.e., for all its uses in `region`.
+bool isValidDim(Value value, Region *region);
+
+/// Returns true if the given value can be used as a symbol in the region of the
+/// closest surrounding op that has the trait `PolyhedralScope`.
 bool isValidSymbol(Value value);
 
+/// Returns true if the given Value can be used as a symbol for `region`, i.e.,
+/// for all its uses in `region`.
+bool isValidSymbol(Value value, Region *region);
+
 /// Modifies both `map` and `operands` in-place so as to:
 /// 1. drop duplicate operands
 /// 2. drop unused dims and symbols from map
 /// 3. promote valid symbols to symbolic operands in case they appeared as
 ///    dimensional operands
 /// 4. propagate constant operands and drop them
 void canonicalizeMapAndOperands(AffineMap *map,
                                 SmallVectorImpl<Value> *operands);

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

def AffineApplyOp : Affine_Op<"apply", [NoSideEffect]> {
 
   let extraClassDeclaration = [{
     /// Returns the affine map to be applied by this operation.
     AffineMap getAffineMap() { return map(); }
 
     /// Returns the affine value map computed from this operation.
     AffineValueMap getAffineValueMap();
 
-    /// Returns true if the result of this operation can be used as dimension id.
+    /// Returns true if the result of this operation can be used as dimension id
+    /// in the region of the closest surrounding op with trait PolyhedralScope.
     bool isValidDim();
 
-    /// Returns true if the result of this operation is a symbol.
+    /// Returns true if the result of this operation can be used as dimension id
+    /// within 'region', i.e., for all its uses with `region`.
+    bool isValidDim(Region *region);
+
+    /// Returns true if the result of this operation is a symbol in the region
+    /// of the closest surrounding op that has the trait PolyhedralScope.
     bool isValidSymbol();
 
+    /// Returns true if the result of this operation is a symbol for all its
+    /// uses in `region`.
+    bool isValidSymbol(Region *region);
+
     operand_range getMapOperands() { return getOperands(); }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 def AffineForOp : Affine_Op<"for",

mlir/include/mlir/IR/Function.h

 // Function Operation.
 //===--------------------------------------------------------------------===//
 
 /// FuncOp represents a function, or an operation containing one region that
 /// forms a CFG(Control Flow Graph). The region of a function is not allowed to
 /// implicitly capture global values, and all external references must use
 /// Function arguments or attributes that establish a symbolic connection(e.g.
 /// symbols referenced by name via a string attribute).
-class FuncOp : public Op<FuncOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
-                         OpTrait::IsIsolatedFromAbove, OpTrait::FunctionLike,
-                         OpTrait::AutomaticAllocationScope,
-                         CallableOpInterface::Trait, SymbolOpInterface::Trait> {
+class FuncOp
+    : public Op<FuncOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
+                OpTrait::AutomaticAllocationScope, OpTrait::PolyhedralScope,
 public:
   using Op::Op;
   using Op::print;
 
   static StringRef getOperationName() { return "func"; }
 
   static FuncOp create(Location location, StringRef name, FunctionType type,
                        ArrayRef<NamedAttribute> attrs = {});

mlir/include/mlir/IR/Module.h

 
 /// ModuleOp represents a module, or an operation containing one region with a
 /// single block containing opaque operations. The region of a module is not
 /// allowed to implicitly capture global values, and all external references
 /// must use symbolic references via attributes(e.g. via a string name).
 class ModuleOp
     : public Op<
           ModuleOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
-          OpTrait::IsIsolatedFromAbove, OpTrait::SymbolTable,
+          OpTrait::IsIsolatedFromAbove, OpTrait::PolyhedralScope,
+          OpTrait::SymbolTable,
           OpTrait::SingleBlockImplicitTerminator<ModuleTerminatorOp>::Impl,
           SymbolOpInterface::Trait> {
 public:
   using Op::Op;
   using Op::print;
 
   static StringRef getOperationName() { return "module"; }
 

mlir/include/mlir/IR/OpBase.td

 // X op Y == Y op X
 def Commutative  : NativeOpTrait<"IsCommutative">;
 // Op behaves like a constant.
 def ConstantLike : NativeOpTrait<"ConstantLike">;
 // Op behaves like a function.
 def FunctionLike : NativeOpTrait<"FunctionLike">;
 // Op is isolated from above.
 def IsolatedFromAbove : NativeOpTrait<"IsIsolatedFromAbove">;
+// Op defines a polyhedral scope.
+def PolyhedralScope : NativeOpTrait<"PolyhedralScope">;
 // Op results are float or vectors/tensors thereof.
 def ResultsAreFloatLike : NativeOpTrait<"ResultsAreFloatLike">;
 // Op has the same operand type.
 def SameTypeOperands : NativeOpTrait<"SameTypeOperands">;
 // Op has same shape for all operands.
 def SameOperandsShape : NativeOpTrait<"SameOperandsShape">;
 // Op has same operand and result shape.
 def SameOperandsAndResultShape : NativeOpTrait<"SameOperandsAndResultShape">;

mlir/include/mlir/IR/OpDefinition.h

public:
   static LogicalResult verifyTrait(Operation *op) {
     for (auto &region : op->getRegions())
       if (!region.isIsolatedFromAbove(op->getLoc()))
         return failure();
     return success();
   }
 };
 
+/// A trait of region holding operations that defines a new scope for polyhedral
+/// optimization purposes. Any SSA values of 'index' type that either dominate
+/// such an operation or are used at the top-level of such an operation
+/// automatically become valid symbols for the polyhedral scope defined by that
+/// operation. For more details, see `Traits.md#PolyhedralScope`.
+template <typename ConcreteType>
+class PolyhedralScope : public TraitBase<ConcreteType, PolyhedralScope> {
+public:
+  static LogicalResult verifyTrait(Operation *op) {
+    static_assert(!ConcreteType::template hasTrait<ZeroRegion>(),
+                  "expected operation to have one or more regions");
+    return success();
+  }
+};
+
 /// A trait of region holding operations that define a new scope for automatic
 /// allocations, i.e., allocations that are freed when control is transferred
 /// back from the operation's region. Any operations performing such allocations
 /// (for eg. std.alloca) will have their allocations automatically freed at
 /// their closest enclosing operation with this trait.
 template <typename ConcreteType>
 class AutomaticAllocationScope
     : public TraitBase<ConcreteType, AutomaticAllocationScope> {

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

 /// Materialize a single constant operation from a given attribute value with
 /// the desired resultant type.
 Operation *AffineDialect::materializeConstant(OpBuilder &builder,
                                               Attribute value, Type type,
                                               Location loc) {
   return builder.create<ConstantOp>(loc, type, value);
 }
 
-/// A utility function to check if a given region is attached to a function.
-static bool isFunctionRegion(Region *region) {
-  return llvm::isa<FuncOp>(region->getParentOp());
+/// A utility function to check if a value is defined at the top level of an
+/// op with trait `PolyhedralScope`. 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::PolyhedralScope>();
+  return value.getDefiningOp()
+      ->getParentOp()
+      ->hasTrait<OpTrait::PolyhedralScope>();
 }
 
-/// A utility function to check if a value is defined at the top level of a
-/// function. A value of index type defined at the top level is always a valid
-/// symbol.
-bool mlir::isTopLevelValue(Value value) {
+/// A utility function to check if a value is defined at the top level of
+/// `region` or is an argument of `region`. A value of index type defined at the
+/// top level of a `PolyhedralScope` region is always a valid symbol for all
+/// uses in that region.
+static bool isTopLevelValue(Value value, Region *region) {
   if (auto arg = value.dyn_cast<BlockArgument>())
-    return isFunctionRegion(arg.getOwner()->getParent());
-  return isFunctionRegion(value.getDefiningOp()->getParentRegion());
+    return arg.getParentRegion() == region;
+  return value.getDefiningOp()->getParentOp() == region->getParentOp();
 }
 
-// Value can be used as a dimension id if it is valid as a symbol, or
-// it is an induction variable, or it is a result of affine apply operation
-// with dimension id arguments.
+/// Returns the closest region enclosing `op` that is held by an operation with
+/// trait `PolyhedralScope`.
+//  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::PolyhedralScope>())
+      return curOp->getParentRegion();
+    curOp = parentOp;
+  }
+  llvm_unreachable("op doesn't have an enclosing polyhedral scope");
+}
+
+// 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) {
   // The value must be an index type.
   if (!value.getType().isIndex())
     return false;
 
-  if (auto *op = value.getDefiningOp()) {
-    // Top level operation or constant operation is ok.
-    if (isFunctionRegion(op->getParentRegion()) || isa<ConstantOp>(op))
-      return true;
-    // Affine apply operation is ok if all of its operands are ok.
-    if (auto applyOp = dyn_cast<AffineApplyOp>(op))
-      return applyOp.isValidDim();
-    // The dim op is okay if its operand memref/tensor is defined at the top
-    // level.
-    if (auto dimOp = dyn_cast<DimOp>(op))
-      return isTopLevelValue(dimOp.getOperand());
-    return false;
-  }
-  // This value has to be a block argument of a FuncOp, an 'affine.for', or an
-  // 'affine.parallel'.
-  auto *parentOp = value.cast<BlockArgument>().getOwner()->getParentOp();
-  return isa<FuncOp>(parentOp) || isa<AffineForOp>(parentOp) ||
-         isa<AffineParallelOp>(parentOp);
-}
+  if (auto *defOp = value.getDefiningOp())
+    return isValidDim(value, getAffineScope(defOp));
+
+  // This value has to be a block argument for an op that has the
+  // `PolyhedralScope` trait or for an affine.for or affine.parallel.
+  auto *parentOp = value.cast<BlockArgument>().getOwner()->getParentOp();
+  return parentOp->hasTrait<OpTrait::PolyhedralScope>() ||
+         isa<AffineForOp>(parentOp) || isa<AffineParallelOp>(parentOp);
+}
+
+// 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 an affine apply operation with dimension id operands.
+bool mlir::isValidDim(Value value, Region *region) {
+  // The value must be an index type.
+  if (!value.getType().isIndex())
+    return false;
+
+  // All valid symbols are okay.
+  if (isValidSymbol(value, region))
+
+  auto *op = value.getDefiningOp();
+  if (!op) {
+    // This value has to be a block argument for an affine.for or an
+    // affine.parallel.
+    auto *parentOp = value.cast<BlockArgument>().getOwner()->getParentOp();
+    return isa<AffineForOp>(parentOp) || isa<AffineParallelOp>(parentOp);
+  }
+
+    return applyOp.isValidDim(region);
 
 /// Returns true if the 'index' dimension of the `memref` defined by
-/// `memrefDefOp` is a statically  shaped one or defined using a valid symbol.
+/// `memrefDefOp` is a statically  shaped one or defined using a valid symbol
+/// for `region`.
 template <typename AnyMemRefDefOp>
-static bool isMemRefSizeValidSymbol(AnyMemRefDefOp memrefDefOp,
-                                    unsigned index) {
+bool isMemRefSizeValidSymbol(AnyMemRefDefOp memrefDefOp, unsigned index,
+                             Region *region) {
   auto memRefType = memrefDefOp.getType();
   // Statically shaped.
   if (!memRefType.isDynamicDim(index))
     return true;
   // Get the position of the dimension among dynamic dimensions;
   unsigned dynamicDimPos = memRefType.getDynamicDimIndex(index);
-  return isValidSymbol(
-      *(memrefDefOp.getDynamicSizes().begin() + dynamicDimPos));
+  return isValidSymbol(*(memrefDefOp.getDynamicSizes().begin() + dynamicDimPos),
+                       region);
 }
 
-/// Returns true if the result of the dim op is a valid symbol.
-static bool isDimOpValidSymbol(DimOp dimOp) {
+/// Returns true if the result of the dim op is a valid symbol for `region`.
+static bool isDimOpValidSymbol(DimOp dimOp, Region *region) {
   // The dim op is okay if its operand memref/tensor is defined at the top
   // level.
   if (isTopLevelValue(dimOp.getOperand()))
     return true;
 
   // The dim op is also okay if its operand memref/tensor is a view/subview
   // whose corresponding size is a valid symbol.
   unsigned index = dimOp.getIndex();
   if (auto viewOp = dyn_cast<ViewOp>(dimOp.getOperand().getDefiningOp()))
-    return isMemRefSizeValidSymbol<ViewOp>(viewOp, index);
+    return isMemRefSizeValidSymbol<ViewOp>(viewOp, index, region);
   if (auto subViewOp = dyn_cast<SubViewOp>(dimOp.getOperand().getDefiningOp()))
-    return isMemRefSizeValidSymbol<SubViewOp>(subViewOp, index);
+    return isMemRefSizeValidSymbol<SubViewOp>(subViewOp, index, region);
   if (auto allocOp = dyn_cast<AllocOp>(dimOp.getOperand().getDefiningOp()))
-    return isMemRefSizeValidSymbol<AllocOp>(allocOp, index);
+    return isMemRefSizeValidSymbol<AllocOp>(allocOp, index, region);
   return false;
 }
 
-// Value can be used as a symbol if it is a constant, or it is defined at
-// the top level, or it is a result of affine apply operation with symbol
-// arguments, or a result of the dim op on a memref satisfying certain
-// constraints.
+// A value can be used as a symbol (at all its use sites) iff it meets one of
+// the following conditions:
+// *) It is a constant.
+// *) Its defining op or block arg appearance is immediately enclosed by an op
+//    with `PolyhedralScope` trait.
+// *) It is the result of an affine.apply operation with symbol operands.
+// *) It is a result of the dim op on a memref whose corresponding size is a
+//    valid symbol.
 bool mlir::isValidSymbol(Value value) {
   // The value must be an index type.
   if (!value.getType().isIndex())
     return false;
 
-  if (auto *op = value.getDefiningOp()) {
-    // Top level operation or constant operation is ok.
-    if (isFunctionRegion(op->getParentRegion()) || isa<ConstantOp>(op))
-      return true;
-    // Affine apply operation is ok if all of its operands are ok.
-    if (auto applyOp = dyn_cast<AffineApplyOp>(op))
-      return applyOp.isValidSymbol();
-    if (auto dimOp = dyn_cast<DimOp>(op)) {
-      return isDimOpValidSymbol(dimOp);
-    }
+  // Check that the value is a top level value.
+  if (isTopLevelValue(value))
+
+  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.
+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))
+    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 (auto *parentOpRegion = region->getParentOp()->getParentRegion())
+        return isValidSymbol(value, parentOpRegion);
+    return false;
   }
-  // Otherwise, check that the value is a top level value.
-  return isTopLevelValue(value);
+
+  // 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 (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).
-// Returns false otherwise.
-static bool isValidAffineIndexOperand(Value value) {
-  return isValidDim(value) || isValidSymbol(value);
+// affine.load, affine.store, affine.dma_start, affine.dma_wait) where
+// `region` provides the polyhedral symbol scope. Returns false otherwise.
+static bool isValidAffineIndexOperand(Value value, Region *region) {
+  return isValidDim(value, region) || isValidSymbol(value, region);
 }
 
 /// Utility function to verify that a set of operands are valid dimension and
 /// symbol identifiers. The operands should be laid out such that the dimension
 /// operands are before the symbol operands. This function returns failure if
 /// there was an invalid operand. An operation is provided to emit any necessary
 /// errors.
 template <typename OpTy>
 static LogicalResult
 verifyDimAndSymbolIdentifiers(OpTy &op, Operation::operand_range operands,
                               unsigned numDims) {
   unsigned opIt = 0;
   for (auto operand : operands) {
     if (opIt++ < numDims) {
-      if (!isValidDim(operand))
+      if (!isValidDim(operand, getAffineScope(op)))
         return op.emitOpError("operand cannot be used as a dimension id");
-    } else if (!isValidSymbol(operand)) {
+    } else if (!isValidSymbol(operand, getAffineScope(op))) {
       return op.emitOpError("operand cannot be used as a symbol");
     }
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // AffineApplyOp

 
 // The result of the affine apply operation can be used as a dimension id if all
 // its operands are valid dimension ids.
 bool AffineApplyOp::isValidDim() {
   return llvm::all_of(getOperands(),
                       [](Value op) { return mlir::isValidDim(op); });
 }
 
+// The result of the affine apply operation can be used as a dimension id if all
+// its operands are valid dimension ids with the parent operation of `region`
+// defining the polyhedral scope for symbols.
+bool AffineApplyOp::isValidDim(Region *region) {
+  return llvm::all_of(getOperands(),
+                      [&](Value op) { return ::isValidDim(op, region); });
+}
+
 // The result of the affine apply operation can be used as a symbol if all its
 // operands are symbols.
 bool AffineApplyOp::isValidSymbol() {
   return llvm::all_of(getOperands(),
                       [](Value op) { return mlir::isValidSymbol(op); });
 }
 
+// The result of the affine apply operation can be used as a symbol in `region`
+// if all its operands are symbols in `region`.
+bool AffineApplyOp::isValidSymbol(Region *region) {
+  return llvm::all_of(getOperands(), [&](Value operand) {
+    return mlir::isValidSymbol(operand, region);
+  });
+}
+
 OpFoldResult AffineApplyOp::fold(ArrayRef<Attribute> operands) {
   auto map = getAffineMap();
 
   // Fold dims and symbols to existing values.
   auto expr = map.getResult(0);
   if (auto dim = expr.dyn_cast<AffineDimExpr>())
     return getOperand(dim.getPosition());
   if (auto sym = expr.dyn_cast<AffineSymbolExpr>())
LogicalResult AffineDmaStartOp::verify() {
   unsigned numInputsAllMaps = getSrcMap().getNumInputs() +
                               getDstMap().getNumInputs() +
                               getTagMap().getNumInputs();
   if (getNumOperands() != numInputsAllMaps + 3 + 1 &&
       getNumOperands() != numInputsAllMaps + 3 + 1 + 2) {
     return emitOpError("incorrect number of operands");
   }
 
+  Region *scope = getAffineScope(*this);
   for (auto idx : getSrcIndices()) {
     if (!idx.getType().isIndex())
       return emitOpError("src index to dma_start must have 'index' type");
-    if (!isValidAffineIndexOperand(idx))
+    if (!isValidAffineIndexOperand(idx, scope))
       return emitOpError("src index must be a dimension or symbol identifier");
   }
   for (auto idx : getDstIndices()) {
     if (!idx.getType().isIndex())
       return emitOpError("dst index to dma_start must have 'index' type");
-    if (!isValidAffineIndexOperand(idx))
+    if (!isValidAffineIndexOperand(idx, scope))
       return emitOpError("dst index must be a dimension or symbol identifier");
   }
   for (auto idx : getTagIndices()) {
     if (!idx.getType().isIndex())
       return emitOpError("tag index to dma_start must have 'index' type");
-    if (!isValidAffineIndexOperand(idx))
+    if (!isValidAffineIndexOperand(idx, scope))
       return emitOpError("tag index must be a dimension or symbol identifier");
   }
   return success();
 }
 
 LogicalResult AffineDmaStartOp::fold(ArrayRef<Attribute> cstOperands,
                                      SmallVectorImpl<OpFoldResult> &results) {
   /// dma_start(memrefcast) -> dma_start
if (tagMapOperands.size() != tagMapAttr.getValue().getNumInputs())
     return parser.emitError(parser.getNameLoc(),
                             "tag memref operand count != to map.numInputs");
   return success();
 }
 
 LogicalResult AffineDmaWaitOp::verify() {
   if (!getOperand(0).getType().isa<MemRefType>())
     return emitOpError("expected DMA tag to be of memref type");
+  Region *scope = getAffineScope(*this);
   for (auto idx : getTagIndices()) {
     if (!idx.getType().isIndex())
       return emitOpError("index to dma_wait must have 'index' type");
-    if (!isValidAffineIndexOperand(idx))
+    if (!isValidAffineIndexOperand(idx, scope))
       return emitOpError("index must be a dimension or symbol identifier");
   }
   return success();
 }
 
 LogicalResult AffineDmaWaitOp::fold(ArrayRef<Attribute> cstOperands,
                                     SmallVectorImpl<OpFoldResult> &results) {
   /// dma_wait(memrefcast) -> dma_wait
if (mapAttr) {
     if (map.getNumInputs() != op.getNumOperands() - 1)
       return op.emitOpError("expects as many subscripts as affine map inputs");
   } else {
     if (op.getMemRefType().getRank() != op.getNumOperands() - 1)
       return op.emitOpError(
           "expects the number of subscripts to be equal to memref rank");
   }
 
+  Region *scope = getAffineScope(op);
   for (auto idx : op.getMapOperands()) {
     if (!idx.getType().isIndex())
       return op.emitOpError("index to load must have 'index' type");
-    if (!isValidAffineIndexOperand(idx))
+    if (!isValidAffineIndexOperand(idx, scope))
       return op.emitOpError("index must be a dimension or symbol identifier");
   }
   return success();
 }
 
 void AffineLoadOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<SimplifyAffineOp<AffineLoadOp>>(context);
if (mapAttr) {
     if (map.getNumInputs() != op.getNumOperands() - 2)
       return op.emitOpError("expects as many subscripts as affine map inputs");
   } else {
     if (op.getMemRefType().getRank() != op.getNumOperands() - 2)
       return op.emitOpError(
           "expects the number of subscripts to be equal to memref rank");
   }
 
+  Region *scope = getAffineScope(op);
   for (auto idx : op.getMapOperands()) {
     if (!idx.getType().isIndex())
       return op.emitOpError("index to store must have 'index' type");
-    if (!isValidAffineIndexOperand(idx))
+    if (!isValidAffineIndexOperand(idx, scope))
       return op.emitOpError("index must be a dimension or symbol identifier");
   }
   return success();
 }
 
 void AffineStoreOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<SimplifyAffineOp<AffineStoreOp>>(context);
if (map.getNumResults() != op.getMemRefType().getRank())
                             " memref rank");
     if (map.getNumInputs() + 1 != op.getNumOperands())
       return op.emitOpError("too few operands");
   } else {
     if (op.getNumOperands() != 1)
       return op.emitOpError("too few operands");
   }
 
+  Region *scope = getAffineScope(op);
   for (auto idx : op.getMapOperands()) {
-    if (!isValidAffineIndexOperand(idx))
+    if (!isValidAffineIndexOperand(idx, scope))
       return op.emitOpError("index must be a dimension or symbol identifier");
   }
   return success();
 }
 
 void AffinePrefetchOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   // prefetch(memrefcast) -> prefetch

mlir/test/Dialect/Affine/invalid.mlir

 
 #set0 = affine_set<(i)[N] : (i >= 0, N - i >= 0)>
 
 func @affine_if_invalid_dimop_dim(%arg0: index, %arg1: index, %arg2: index, %arg3: index) {
   affine.for %n0 = 0 to 7 {
     %0 = alloc(%arg0, %arg1, %arg2, %arg3) : memref<?x?x?x?xf32>
     %dim = dim %0, 0 : memref<?x?x?x?xf32>
 
-    // expected-error@+1 {{operand cannot be used as a dimension id}}
+    // expected-error@+1 {{operand cannot be used as a symbol}}
     affine.if #set0(%dim)[%n0] {}
   }
   return
 }
 
 // -----
 
 func @affine_store_missing_l_square(%C: memref<4096x4096xf32>) {

mlir/test/Dialect/Affine/ops.mlir

affine.for %arg4 = 0 to %13 step 264 {
       }
     }
   }
   return
 }
 
 // -----
 
+// Test symbol constraints for ops with PolyhedralScope trait.
+
+// CHECK-LABEL: func @valid_symbol_polyhedral_scope
+func @valid_symbol_polyhedral_scope(%n : index, %A : memref<?xf32>) {
+  test.polyhedral_scope {
+    %c1 = constant 1 : index
+    %l = subi %n, %c1 : index
+    // %l, %n are valid symbols since test.polyhedral_scope defines a new
+    // polyhedral scope.
+    affine.for %i = %l to %n {
+      %m = subi %l, %i : index
+      test.polyhedral_scope {
+        // %m and %n are valid symbols.
+        affine.for %j = %m to %n {
+          %v = affine.load %A[%n - 1] : memref<?xf32>
+          affine.store %v, %A[%n - 1] : memref<?xf32>
+        }
+        "terminate"() : () -> ()
+      }
+    }
+    "terminate"() : () -> ()
+  }
+  return
+}
+
+// -----
+
+// Test the fact that module op always provides a polyhedral scope.
+
+%idx = "test.foo"() : () -> (index)
+"test.func"() ({
+^bb0(%A : memref<?xf32>):
+  affine.load %A[%idx] : memref<?xf32>
+  "terminate"() : () -> ()
+}) : () -> ()
+
+// -----
+
 // CHECK-LABEL: @parallel
 // CHECK-SAME: (%[[N:.*]]: index)
 func @parallel(%N : index) {
   // CHECK: affine.parallel (%[[I0:.*]], %[[J0:.*]]) = (0, 0) to (symbol(%[[N]]), 100) step (10, 10)
   affine.parallel (%i0, %j0) = (0, 0) to (symbol(%N), 100) step (10, 10) {
     // CHECK-NEXT: affine.parallel (%{{.*}}, %{{.*}}) = (%[[I0]], %[[J0]]) to (%[[I0]] + 10, %[[J0]] + 10)
     affine.parallel (%i1, %j1) = (%i0, %j0) to (%i0 + 10, %j0 + 10) {
     }
   }
   return
 }

mlir/test/lib/Dialect/Test/TestDialect.cpp

 static void print(OpAsmPrinter &p, IsolatedRegionOp op) {
   p << "test.isolated_region ";
   p.printOperand(op.getOperand());
   p.shadowRegionArgs(op.region(), op.getOperand());
   p.printRegion(op.region(), /*printEntryBlockArgs=*/false);
 }
 
 //===----------------------------------------------------------------------===//
+// Test PolyhedralScopeOp
+//===----------------------------------------------------------------------===//
+
+static ParseResult parsePolyhedralScopeOp(OpAsmParser &parser,
+                                          OperationState &result) {
+  // Parse the body region, and reuse the operand info as the argument info.
+  Region *body = result.addRegion();
+  return parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{});
+}
+
+static void print(OpAsmPrinter &p, PolyhedralScopeOp op) {
+  p << "test.polyhedral_scope ";
+  p.printRegion(op.region(), /*printEntryBlockArgs=*/false);
+}
+
+//===----------------------------------------------------------------------===//
 // Test parser.
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseWrappedKeywordOp(OpAsmParser &parser,
                                          OperationState &result) {
   StringRef keyword;
   if (parser.parseKeyword(&keyword))
     return failure();

mlir/test/lib/Dialect/Test/TestOps.td

 //===-- TestOps.td - Test dialect operation definitions ----*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef TEST_OPS
 #define TEST_OPS
 
+include "mlir/Dialect/Affine/IR/AffineOpsBase.td"
 include "mlir/IR/OpBase.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/SideEffects.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffects.td"
def IsolatedRegionOp : TEST_Op<"isolated_region", [IsolatedFromAbove]> {
   }];
 
   let arguments = (ins Index);
   let regions = (region SizedRegion<1>:$region);
   let parser = [{ return ::parse$cppClass(parser, result); }];
   let printer = [{ return ::print(p, *this); }];
 }
 
+def PolyhedralScopeOp : TEST_Op<"polyhedral_scope", [PolyhedralScope]> {
+  let summary =  "polyhedral scope operation";
+  let description = [{
+    Test op that defines a new polyhedral scope.
+  }];
+
+  let regions = (region SizedRegion<1>:$region);
+  let parser = [{ return ::parse$cppClass(parser, result); }];
+  let printer = [{ return ::print(p, *this); }];
+}
+
 def WrappingRegionOp : TEST_Op<"wrapping_region",
     [SingleBlockImplicitTerminator<"TestReturnOp">]> {
   let summary =  "wrapping region operation";
   let description = [{
     Test op wrapping another op in a region, to test calling
     parseGenericOperation from the custom parser.
   }];