[MLIR] Introduce op trait PolyhedralScope

[bondhugula]

Summary

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 affine ops and their analysis/transforms. Update methods that check
for dim/symbol validity to work based on this trait.

Note: This op trait could be moved to the affine dialect so that only
those who depend on it are able to use it. However, this would mean that
FuncOp will have to be treated specially in the symbol checks, i.e.,
(FuncOp || op with trait PolyhedralScope) would replace (op with trait
PolyhedralScope) for all purposes. Keeping this trait in the core IR
allows FuncOp to be marked PolyhedralScope.

[bondhugula]

Add trait documentation

[bondhugula]

Fix comment

[ftynse]

Looks fine to me with comments adressed

mlir/docs/Dialects/Affine.md

64

Nit: can we make this an actual markdown list?

mlir/include/mlir/IR/OpDefinition.h

1060

Shouldn't we rather change that op actually has regions? An op may not have ZeroRegion trait and still have no associated regions.

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

157–158

Did you mean parentOp->hasTrait<PolyhedralScope>() instead of isolated from above?

mlir/test/Dialect/Affine/ops.mlir

122

It is a bit unusual to have a dialect test depend on ops from the test dialect, but looks okay given the nature of the test.

[bondhugula]

Address review comments

[bondhugula]

Thanks for the review!

mlir/docs/Dialects/Affine.md

64

Sure - that would be better.

mlir/include/mlir/IR/OpDefinition.h

1060

One could imagine an op that has some instances of it with zero regions and some instances with a region. One would in such cases let the op have the PolyhedralScope trait: we don't have such an example in the code base or a scenario I can think of right away but no regions is actually no harm. So, I thought we'll only disallow if it's statically a zero region op. That said, this check should have actually been a static_assert:

static_assert(ConcreteType::template hasTrait<ZeroRegion>(),
                  "expected operation to have one or more regions");

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

157–158

Thanks - forgot to update this.

mlir/test/Dialect/Affine/ops.mlir

122

Hmm... there is one in parser.mlir also that depends on test.isolated_region; can't think of a way to test these without the test dialect. But once there is affine.execute_region, that could replace this.

[rriddle]

mlir/include/mlir/IR/OpDefinition.h

1056

Is this trait applicable to anything outside of the affine dialect? If not I don't really see why it should be here and not in the affine dialect instead.

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

104

value.getParentRegion() == region?

249

Can we use m_Constant instead?

[bondhugula]

mlir/include/mlir/IR/OpDefinition.h

1056

I was in two minds on this. Could you see the second para of the commit summary? We'd like to mark FuncOp with PolyhedralScope.

[rriddle]

mlir/include/mlir/IR/OpDefinition.h

1056

Ah, sorry I missed that. A lot of times I skip the commit message and jump right in.

Yes, marking FuncOp makes things a bit difficult. This has come up before w.r.t interfaces as well. I'd like to avoid adding special dialect traits here if they aren't general(especially because I'd like to trim this file a bit). I would expect that you could also use FunctionLike(instead of FuncOp) as an indicator of a polyhedral scope. I'm leaning more towards moving the trait to Affine unless you can foresee uses outside of that dialect.

[bondhugula]

mlir/include/mlir/IR/OpDefinition.h

1056

There could be things that don't depend on the Affine dialect but would want to use the 'PolyhedralScope' on their FuncLike or other region ops. That's because their lowering could potentially go through the affine dialect much later which their region holding op lasts through - is this meaningful?

But when I initially started with keeping this in the Affine dialect, all the clean and compact code that we have now becomes weird. To be precise:

op->hasTrait<OpTrait::PolyhedralScope> becomes
op->hasTrait<OpTrait::PolyhedralScope> || isa<FuncOp>(op)

We can't right away use FuncLikeOp (although we want to go in that direction) because not all func op's may want to opt into this; you'd want to whitelist the desired ones by attaching the PolyhedralScope trait - they may or may not depend on the affine dialect. Not having PolyhedralScope attached to FuncOp would mean a simple getParentWithTrait<PolyhedralScope> becomes an entire while loop that tries to find an op that either has the PolyhedralScope trait or is a FuncOp.

[rriddle]

mlir/include/mlir/IR/OpDefinition.h

1056

I agree that it is easier to manipulate when there is one interface for detection. Can we have an Affine/Traits.h that defines this and just include it in Function.h? It isn't ideal, but we already have to do this for interfaces anyways? I'm just trying to avoid the future where many non-core traits get pulled into IR/ because of the current weirdness with FuncOp.

[bondhugula]

Create Affine/Traits.h and move polyhedral scope trait there.

[bondhugula]

Resolve remaining review comments.

[rriddle]

LGTM after resolving the other comments.

mlir/include/mlir/IR/OpBase.td

1633

Can you move this to AffineOpsBase.td or somewhere similar?

[andydavis1]

Thanks Uday...

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

139

Same comment as for isSymbol. Make the conditions a list, so its easier for people to see the set of conditions.

224

There are a lot of "ors" in this description, and this is the place most people will be looking for what the definition of what a symbol is. Could you make this a bullet list:
Value can be used as a symbol iff it meets one of the following conditions:
*) It is a constant.
*) It is defined at the top level of 'region'...
*) ...

[bondhugula]

Resolve remaining review comments.

mlir/include/mlir/IR/OpDefinition.h

1056

Sure - that sounds good. Done.

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

104

Thanks.

224

Sure.

[bondhugula]

Rebase.

[gribozavr2]

Sorry, I reverted this change in https://github.com/llvm/llvm-project/commit/ef06016d73390d5b380018cc0d16003b4ed4a35a. @ftynse understands the problem and will respond here a bit later.

[bondhugula]

In D78863#2007528, @gribozavr2 wrote:

Sorry, I reverted this change in https://github.com/llvm/llvm-project/commit/ef06016d73390d5b380018cc0d16003b4ed4a35a. @ftynse understands the problem and will respond here a bit later.

I had tested this locally with Clang and GCC. The harbormaster tests were already failing due to unrelated clang changes. I'm curious to know what the issue was. Thanks.

[ftynse]

In D78863#2007528, @gribozavr2 wrote:

Sorry, I reverted this change in https://github.com/llvm/llvm-project/commit/ef06016d73390d5b380018cc0d16003b4ed4a35a. @ftynse understands the problem and will respond here a bit later.

This broke IREE and also was problematic with Bazel builds. There are two issues, and I posit that one of them caused another.

1. IREE pass pipeline triggered llvm_unreachable("op doesn't have an enclosing polyhedral scope"); (AffineOps.cpp:119). We investigated this and found that this pipeline _first_ transforms std.func into gpu.func, and then runs transformations that depend transitively on Affine within gpu.func. The latter not having the PolyhedralScope trait, and with no other function like op around, the loop in getAffineScope completed and the assertion was triggered. An obvious solution to this is to give gpu.func op the trait. However, I think this failure rather indicates a deeper problem with the affine scope modeling: nothing in the semantics of the Affine Ops requires (and verifies) that they _only_ appear within some op with a PolyhedralScope. So the control flow path with "llvm_unrechable" actually _is_ reachable by valid IR. We can either add such requirement, or consider top-level ModuleOp (which should be there by default) to also be a polyhedral scope. This needs further discussion.
2. The introduction of Affine/Traits.h created a circular dependency between "IR" and "AffineDialect" libraries (AffineDialect depends on IR because it uses core IR constructs, and IR depends on AffineDialect because Function has a trait defined in Affine). It is actually possible to circumvent the issue in Bazel build files we use, e.g., for TensorFlow, but it is considered to be a hack. Since this specific issue was discussed in the review, let's iterate one more time on it. It party stems from the original lack of separation between core IR and Affine+Standard dialects. Factoring Function out of the "IR" library can be one way to break the cycle. However, the issue of dependency on the affine dialect only for the purpose of having a trait remains. Generally, any dialect with region-carrying ops that wishes to contain ops from the Affine dialect has to depend on it, e.g. GPU. This sounds wrong layering-wise. Since we don't want dialect-specific traits to live in IR, another possible separation is to have dialect-specific traits live in a library (and filepath) separate from both IR and Dialect, with both IR and Dialect depending on this library.

[bondhugula]

In D78863#2007580, @ftynse wrote:

In D78863#2007528, @gribozavr2 wrote:

Sorry, I reverted this change in https://github.com/llvm/llvm-project/commit/ef06016d73390d5b380018cc0d16003b4ed4a35a. @ftynse understands the problem and will respond here a bit later.

This broke IREE and also was problematic with Bazel builds. There are two issues, and I posit that one of them caused another.

1. IREE pass pipeline triggered llvm_unreachable("op doesn't have an enclosing polyhedral scope"); (AffineOps.cpp:119). We investigated this and found that this pipeline _first_ transforms std.func into gpu.func, and then runs transformations that depend transitively on Affine within gpu.func. The latter not having the PolyhedralScope trait, and with no other function like op around, the loop in getAffineScope completed and the assertion was triggered. An obvious solution to this is to give gpu.func op the trait. However, I think this failure rather indicates a deeper problem with the affine scope modeling: nothing in the semantics of the Affine Ops requires (and verifies) that they _only_ appear within some op with a PolyhedralScope. So the control flow path with "llvm_unrechable" actually _is_ reachable by valid IR. We can either add such requirement, or consider top-level ModuleOp (which should be there by default) to also be a polyhedral scope. This needs further discussion.

As you hint, marking the gpu.func with PolyhedralScope isn't really a solution, although we may want to do that in the future. Marking the ModuleOp PolyhedralScope appears to be the right fix to me - the root op is always a valid polyhedral scope.

2. The introduction of Affine/Traits.h created a circular dependency between "IR" and "AffineDialect" libraries (AffineDialect depends on IR because it uses core IR constructs, and IR depends on AffineDialect because Function has a trait defined in Affine). It is actually possible to circumvent the issue in Bazel build files we use, e.g., for TensorFlow, but it is considered to be a hack. Since this specific issue was discussed in the review, let's iterate one more time on it. It party stems from the original lack of separation between core IR and Affine+Standard dialects. Factoring Function out of the "IR" library can be one way to break the cycle. However, the issue of dependency on the affine dialect only for the purpose of having a trait remains. Generally, any dialect with region-carrying ops that wishes to contain ops from the Affine dialect has to depend on it, e.g. GPU. This sounds wrong layering-wise. Since we don't want dialect-specific traits to live in IR, another possible separation is to have dialect-specific traits live in a library (and filepath) separate from both IR and Dialect, with both IR and Dialect depending on this library.

This was due to oversight. It was River's suggestion to include Affine/Traits.h from Function, but he perhaps intended for it to be in IR/ - I earlier had it in OpDefinition.h where it didn't create a cyclic dependence.

[ftynse]

In D78863#2007673, @bondhugula wrote:

In D78863#2007580, @ftynse wrote:

In D78863#2007528, @gribozavr2 wrote:

Sorry, I reverted this change in https://github.com/llvm/llvm-project/commit/ef06016d73390d5b380018cc0d16003b4ed4a35a. @ftynse understands the problem and will respond here a bit later.

This broke IREE and also was problematic with Bazel builds. There are two issues, and I posit that one of them caused another.

1. IREE pass pipeline triggered llvm_unreachable("op doesn't have an enclosing polyhedral scope"); (AffineOps.cpp:119). We investigated this and found that this pipeline _first_ transforms std.func into gpu.func, and then runs transformations that depend transitively on Affine within gpu.func. The latter not having the PolyhedralScope trait, and with no other function like op around, the loop in getAffineScope completed and the assertion was triggered. An obvious solution to this is to give gpu.func op the trait. However, I think this failure rather indicates a deeper problem with the affine scope modeling: nothing in the semantics of the Affine Ops requires (and verifies) that they _only_ appear within some op with a PolyhedralScope. So the control flow path with "llvm_unrechable" actually _is_ reachable by valid IR. We can either add such requirement, or consider top-level ModuleOp (which should be there by default) to also be a polyhedral scope. This needs further discussion.

As you hint, marking the gpu.func with PolyhedralScope isn't really a solution, although we may want to do that in the future. Marking the ModuleOp PolyhedralScope appears to be the right fix to me - the root op is always a valid polyhedral scope.

I could not convince myself it would work in that IREE case, but it probably would. isValidDim was not checking for GPUFuncOp as a symbol definition previously, so it shouldn't affect the behavior.

2. The introduction of Affine/Traits.h created a circular dependency between "IR" and "AffineDialect" libraries (AffineDialect depends on IR because it uses core IR constructs, and IR depends on AffineDialect because Function has a trait defined in Affine). It is actually possible to circumvent the issue in Bazel build files we use, e.g., for TensorFlow, but it is considered to be a hack. Since this specific issue was discussed in the review, let's iterate one more time on it. It party stems from the original lack of separation between core IR and Affine+Standard dialects. Factoring Function out of the "IR" library can be one way to break the cycle. However, the issue of dependency on the affine dialect only for the purpose of having a trait remains. Generally, any dialect with region-carrying ops that wishes to contain ops from the Affine dialect has to depend on it, e.g. GPU. This sounds wrong layering-wise. Since we don't want dialect-specific traits to live in IR, another possible separation is to have dialect-specific traits live in a library (and filepath) separate from both IR and Dialect, with both IR and Dialect depending on this library.

This was due to oversight. It was River's suggestion to include Affine/Traits.h from Function, but he perhaps intended for it to be in IR/ - I earlier had it in OpDefinition.h where it didn't create a cyclic dependence.

He mentioned there is already a similar circular dependency between Interfaces and IR, which is indeed the case, so it looks like he meant what you did...

[bondhugula]

@ftynse, @rriddle As I fix this, should the commit title / summary be "Revert "Revert ...."" or should it just be the original title and summary "[MLIR] Introduce ..."? Thanks.

[ftynse]

In D78863#2009104, @bondhugula wrote:

@ftynse, @rriddle As I fix this, should the commit title / summary be "Revert "Revert ...."" or should it just be the original title and summary "[MLIR] Introduce ..."? Thanks.

I think any mention of the previous commit series is fine.

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
-[`afffine.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:
   void print(OpAsmPrinter &p);
   LogicalResult verify();
   static void getCanonicalizationPatterns(OwningRewritePatternList &results,
                                           MLIRContext *context);
   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

 /// 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::Symbol,
                 OpTrait::FunctionLike, OpTrait::AutomaticAllocationScope,
-                CallableOpInterface::Trait> {
+                OpTrait::PolyhedralScope, CallableOpInterface::Trait> {
 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/OpBase.td

 // PredOpTrait is an op trait implemented by way of a predicate on the op.
 class PredOpTrait<string descr, Pred pred> : OpTrait {
   string description = descr;
   Pred predicate = pred;
 }
 
 // Op defines an automatic allocation scope.
 def AutomaticAllocationScope : NativeOpTrait<"AutomaticAllocationScope">;
+// Op defines a polyhedral scope.
+def PolyhedralScope : NativeOpTrait<"PolyhedralScope">;
 // Op supports operand broadcast behavior.
 def ResultsBroadcastableShape :
   NativeOpTrait<"ResultsBroadcastableShape">;
 // TODO: Alias of the above, remove post integrate.
 def Broadcastable : NativeOpTrait<"ResultsBroadcastableShape">;
 // X op Y == Y op X
 def Commutative  : NativeOpTrait<"IsCommutative">;
 // Op behaves like a constant.

mlir/include/mlir/IR/OpDefinition.h

 public:
   static LogicalResult verifyTrait(Operation *op) {
     if (op->hasTrait<ZeroRegion>())
       return op->emitOpError("is expected to have regions");
     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();
+  }
+};
+
 /// This class provides APIs and verifiers for ops with regions having a single
 /// block that must terminate with `TerminatorOpType`.
 template <typename TerminatorOpType> struct SingleBlockImplicitTerminator {
   template <typename ConcreteType>
   class Impl : public TraitBase<ConcreteType, Impl> {
   public:
     static LogicalResult verifyTrait(Operation *op) {
       for (unsigned i = 0, e = op->getNumRegions(); i < e; ++i) {

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.getOwner()->getParentOp() == region->getParentOp();
+  return value.getDefiningOp()->getParentOp() == region->getParentOp();
+}
+
+/// 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");
 }
 
-// Value can be used as a dimension id if it is valid as a symbol, or
+// A 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.
 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 if it is valid as a symbol, or it is an
+// induction variable, or it is a 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.
+// Value can be used as a symbol (at all its use sites) either if it is a
+// constant, or its defining op or block arg appearance is immediately enclosed
+// by an op with `PolyhedralScope` trait, or is the result of an affine.apply
+// operation with symbol operands, or 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;
+
+// Value can be used as a symbol for `region` if it is a constant, or is defined
+// at the top level of 'region' or is its argument, or dominates `region`'s
+// parent op, or is the result of an affine apply operation with symbol
+// arguments, or 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.
+  if (isa<ConstantOp>(defOp))
+    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 ::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() != getNumOperands() - 1)
       return emitOpError("expects as many subscripts as affine map inputs");
   } else {
     if (getMemRefType().getRank() != getNumOperands() - 1)
       return emitOpError(
           "expects the number of subscripts to be equal to memref rank");
   }
 
+  Region *scope = getAffineScope(*this);
   for (auto idx : getMapOperands()) {
     if (!idx.getType().isIndex())
       return emitOpError("index to load must have 'index' type");
-    if (!isValidAffineIndexOperand(idx))
+    if (!isValidAffineIndexOperand(idx, scope))
       return 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() != getNumOperands() - 2)
       return emitOpError("expects as many subscripts as affine map inputs");
   } else {
     if (getMemRefType().getRank() != getNumOperands() - 2)
       return emitOpError(
           "expects the number of subscripts to be equal to memref rank");
   }
 
+  Region *scope = getAffineScope(*this);
   for (auto idx : getMapOperands()) {
     if (!idx.getType().isIndex())
       return emitOpError("index to store must have 'index' type");
-    if (!isValidAffineIndexOperand(idx))
+    if (!isValidAffineIndexOperand(idx, scope))
       return 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
+}
+
+// -----
+
 // 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

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.
   }];