[MLIR] Introduce affine.execute_region op

[bondhugula]

Summary

The affine.execute_region op executes its region exactly once while
defining a new polyhedral scope for its region for analysis and
transformation purposes, i.e., a new symbol context is defined for
operations appearing in its region. This allows the polyhedral form to
be used in a wider context without the need for function outlining.
The op explicitly captures only memrefs and is lowered readily to an
std.execute_region.

[bondhugula]

• rename getParentAffineScope -> getAffineScope

[bondhugula]

Update doc comments.

[bondhugula]

• invalid test cases
• update affine dialect doc

[bondhugula]

• fix comments

[bondhugula]

• add missed result parsing
• fix verifier

[mehdi_amini]

mlir/docs/Dialects/Affine.md

612 ↗

(On Diff #236236)

I don’t think you addressed my concerns on this topic?

[bondhugula]

mlir/docs/Dialects/Affine.md

612 ↗

(On Diff #236236)

I think I responded to everything and included all of the arguments in the RFC:
https://github.com/bondhugula/llvm-project/blob/graybox/mlir/rfc/rfc-graybox.md
Could you just provide a summary list of the concerns you still have - either here or on that thread as you prefer?

[mehdi_amini]

mlir/docs/Dialects/Affine.md

612 ↗

(On Diff #236236)

I explained my concerns in the original thread https://groups.google.com/a/tensorflow.org/d/msg/mlir/O5PXVbtlSng/3SXmxDiLAAAJ

Here is what I wrote:

I am trying to not consider affine at all here. I wrote these example to try to illustrate how MLIR region/op interaction are structured opaquely to be able to derive cross-dialect invariants in general.
The invariant I am presenting above is independent from any dialect, let me abstract the type further:

func @foo(%value : !dialect.type) {
op.with_region {
any.op(%value) : (!dialect.type) -> ()
}
}

If I look at it generically, here is my take on it:
a) The op.with_region defines the semantic of its immediate region, so it can either accept or reject any.op.
b) Let's assume that op.with_region does not know anything about any.op (no traits, no prior knowledge, the any.op could be unregistered at this point).
c) For this IR to be valid, op.with_region must be accepting unknown op (like any.op).
d) From the perspective of op.with_region, the any.op is like an opaque call to some code it cannot see.

But what if any.op has a region?

func @foo(%value : !dialect.type) {
op.with_region {
any.op(%value) ({
^bb(%value_inside):
// do something with %value_inside (explicitly captured)
}) : (!dialect.type) -> ()
}}

Here what changes is that:
e) any.op has a region now. Unless op.with_region forbids unknown operation from having a region in its verifier, and it does not have specific handling for any.op, then the IR should be valid.
f) Since %value_inside is explicitly captured, without knowing specifically any.op, then the uses of %value_inside cannot be restricted by op.with_region but only by any_op.
g) For any practical purpose here, there should not be any difference between this form and the first one above.

Finally, what if any.op has implicit capture?

func @foo(%value : !dialect.type) {
op.with_region {
any.op() ({
// do something with %value (implicit capture instead of explicit)
}) : (!dialect.type) -> ()
}}

Now:
h) any.op() is implicitly capturing %value.
g) Without more information about any.op (traits, etc.), this should be equivalent to the explicit capture case: if the IR was valid the first and second case, then it should be valid here.

If we don't have these properties, and if op.with_region can constrain the validity of the region attached to any.op, then any.op is not longer in control of the semantics of the enclosed region. No transformation can operate on any.op without knowing all of the enclosing operations, since these can add arbitrary restrictions.
For example, this is a valid IR (you can pipe this in mlir-opt right now):

module {
"d1.op1" () ({
"d2.op2" () ({
module {
func @bar() {
return
}
func @foo() {
call @bar() : () -> ()
return
}
}
"d2.done" () : () -> ()
}) : () -> ()
}) : () -> ()
}

If I get the inner @foo function, and would like to inline the call to @bar, what do I need to check to ensure I can?
If the FuncOp defines the semantic of the region, then the FuncOp should control itself whether it allows to inline or not, and I should query FuncOp for @foo, CallOp for the call-site, FuncOp for @bar, and likely the op inside @bar that I am about to inline.

If you allow to put restriction on what can happen inside @foo(), based on the enclosing operation, then you can't inline unless you ensure that all the enclosing operation will be happy with it (so you need to check the enclosing modules, but also "d1.op1" and "d2.op2").

Basically, this would be breaking the composability of the IR: you couldn't assemble independent pieces and reason about them independently. I don't know why we would want that, here we really want to reason about the functions in the inner module independently if they are surrounded by "d1.op1"() and "d2.op2"() like here (otherwise none of the current passes in MLIR are correct).

I don't think you answered these points that explain why I am not convinced it is OK to have explicit capture just for memref.
I don't think it is necessary to have explicit capture of memref either by the way, dropping this may help getting forward right now.

You answered the email above in the thread, but you didn't address it I believe, you wrote "I'll respond to the connection to the affine.graybox proposal in another post" but I didn't see another post after that.

(I'm on vacation till 1/10, expect some delays in answers)

[bondhugula]

mlir/docs/Dialects/Affine.md

612 ↗

(On Diff #236236)

The downsides of not explicitly capturing memrefs on the graybox are discussed in the RFC here.
https://github.com/bondhugula/llvm-project/blob/graybox/mlir/rfc/rfc-graybox.md#rationale-and-design-alternatives-what-to-explicitly-capture

Another way to think about this is that: because grayboxes introduce a new symbol context, most polyhedral walks would like to conceptually see the graybox just as a "call" with those memrefs passed to start with. Other arbitrary/unknown ops with regions don't start a symbol context and so the affine walks will just walk through such ops (just like they walk through affine.fors/ifs that are encountered). OTOH, walkAffine will not walk through grayboxes from the top. If you don't explicitly capture, the key is that most polyhedral/affine passes will have to stop/check every op for a graybox and then scan the interior of that op for memrefs if it turns out to be a graybox. (For the future, this would even go against multithreading polyhedral passes to run concurrently on different func's and grayboxes in an isolated way, but I'm not bringing this up now in the RFC). With an explicit capture, things would just be a regular operand scan of ops scanned with walkAffine. For the future cases where you really need more precise information on what's happening to the memrefs inside the graybox (just like you may want to in the case of call's), that can be done as needed. Non-memref values on the other hand just move transparently across the boundaries of grayboxes in the regular SSA passes/canonicalizations or hybrid polyhedral/SSA ones.

My concerns with explicit capture are actually very different from yours: that they make it harder to move IR across without actually updating the memrefs being used (either hoist from inside to outside or sink). You'd have to check if you are moving past a graybox and then remap memrefs (consider scalar replacement on affine load/stores, LICM as examples).

But I still strongly feel that explicit capture of the memrefs is the right tradeoff to start with (even if perhaps not the right one eventually) - we can reevaluate its impact and drop if necessary.

Best,
Uday

[bondhugula]

mlir/docs/Dialects/Affine.md

612 ↗

(On Diff #236236)

And finally reg. your points on unknown ops with regions that explicitly capture in addition to perhaps implicitly: that's an issue separate from affine.graybox and that exists in the codebase as is now. If you have an op whose operands bind to its region's arguments in ways that are unknown, the current affine passes/utilities don't check for example if the memrefs inside are shadows, how memory footprints / dependences should be accurate computed, etc. (in fact, they are already incorrect because they'd walk through that op and treat those memrefs as distinct). The solution to this would depend on the pass/utility: some should just do nothing in the presence of unknown ops that have one or more regions (we shouldn't even say 'explicit capture' here because we don't even know if it's really capturing in spite of having operands). Since that same op could also implicitly capture (as you mention in one case), choosing to not look inside may not be an option depending on the pass (unless of course the op is known to be isolated from above). But how is all this related to explicit capture in affine.grayboxes? The latter is a known op where the binding between its operands and region's arguments is well-defined.

[bondhugula]

Rebase.

[bondhugula]

changes after rebase - unregistered ops

[bondhugula]

Rename affine.graybox -> affine.execute_region. Rebase + updates to
context-sensitive valid dim/symbol checking.

[ftynse]

I am very sorry, for some reason this diff was not appearing on my phabricator todo list until recently.... Please do not hesitate to ping me by email if I take more than a couple of workdays to iterate.

In general, this change makes sense to me and looks like a proper extension of the affine modeling. My only design concern is that one can circumvent the explicit-capture mechanism for memrefs. A direct solution would be to disallow any values of memref type to be defined in the region, but I have not considered the implications of this on the expressiveness.

There are several places where some action is performed twice (e.g., parsing attribute dict), requesting change for these.

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

629–630

I understand the idea of the restriction, but it looks like it can be circumvented by:

%0 = ... : memref<...>
%1 = cheat.wrap %0 : memref<...> -> !cheat.opaque
affine.execute_region {
  // This defines the memref inside the region, so seemingly complies with the semantics.
  %2 = cheat.unwrap %1 : !cheat.opaque -> memref<...>
  affine.load %2[...]
}

631–634

I would avoid referring FuncOp, it is not special in any sense, and it actually allows any terminator. Neither would I remind that blocks must have terminators, it's a core IR requirement.

"returns to right after the affine.execute_region op" sounds unnecessarily complex to me. "std.return" terminator placed inside blocks of the "affine.execute_region" returns the control flow to the "affine.execute_region". Since we already mentioned that "execution_region" executes the region exactly ones, it is naturally implied that the control flow will be further transferred to the control-flow successor of "execution_region...

655

Syntax nit: I'd consider omitting empty [] = (), looks distracting

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

108

Nit: wouldn't it be easier to accept the region as argument instead of the operation that contains it?

121

Would variadic template help you? isKnownIsolatedFromAbove isn't a class...

128

The assert message is misleading. The op may have parent ops, just none that satisfy the "affine scope" conditions.

Also, use llvm_unreachable instead of assert(false) and drop the return value

163–164

I cannot relate this comment to the code below.

166

Shouldn't this also check for "KnownIsolatedFromAbove" ?

1054

Nit: we tend to use auto when it increases readability, Operation * would look just fine here

2276

Prefer ValueRange to ArrayRef<Value>

2285

Nit: drop trivial braces

2288

If you have ValueRange, this entire vector manipulation gets replaced by memrefs.getTypes()

2289

You already pushed it in line 2295

2325

Avoid capturing llvm::enumerate by-reference. An enumerator only keeps iterators, so taking copying it is cheap, and we avoid running into potential problems with implicit lifetime extension

2327

Since you already have an enumerator, you might as well mention the position of the mismatching argument.

2390

You already parser the attr dict 4 lines above.

[bondhugula]

In D72223#2001800, @ftynse wrote:

I am very sorry, for some reason this diff was not appearing on my phabricator todo list until recently.... Please do not hesitate to ping me by email if I take more than a couple of workdays to iterate.

In general, this change makes sense to me and looks like a proper extension of the affine modeling. My only design concern is that one can circumvent the explicit-capture mechanism for memrefs. A direct solution would be to disallow any values of memref type to be defined in the region, but I have not considered the implications of this on the expressiveness.

There are several places where some action is performed twice (e.g., parsing attribute dict), requesting change for these.

Thanks for the detailed review! Yes, we need to still converge on the memref explicit capture. But I'll address these other straightforward changes first.

[bondhugula]

Address review comments.

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

629–630

Independently of affine.execute_region, such a problem exists with memref non-dereferencing ops and it can / has to be tackled. For example, consider this variation of your snippet (no affine.execute_region).

%0 = ... : memref<...>
%1 = cheat.wrap %0 : memref<...> -> !cheat.opaque
%2 = cheat.load %1[0, 0] : !cheat.opaque
cheat.store %v, %1[0, 0] : !cheat.opaque
call @foo(%1) : (!cheat.opaque) -> ()
affine.load %0[0, 0] : memref<...>

Note that all the current passes/utilities including affine store to load fwd'ing, invariant load hoisting/scalar rep, dependence analysis itself, the affine loop fusion would do the wrong thing here because there is an escape side channel like you show. So, the thing you are pointing to is interesting and has to be tackled, but if we step back and take another look, this is the same as the larger issue of dealing with escaping or aliasing and not specific to execute_region. A solution to deal with these is to actually detect/treat unknown memref non-dereferencing ops that define SSA values (like your cheat.unwrap) and bail out in their presence (depending on what we are doing). For eg. the dependence information isn't going to be accurate in their presence. The point of explicit captures is that you know which memrefs are going in, but it isn't free of the problem of side-channel escapes that manifest in straightline code themselves. But whenever we don't have such escapes (and we can always detect that being conservative), which I believe is the common scenario, having the captures does serve its intended purpose -- you still won't have to look inside and do something special with these ops in all passes/utilities.

631–634

Sure, dropped these lines.

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

108

Actually, it is - thanks!

121

Right - it won't; the comment is probably stale (it was FuncOp and AffineExecuteRegionOp earlier). Anyway this code will be updated to make use of a new op trait 'PolyhedralScope' in another revision, which I'll submit as this one's parent. AffineExecuteRegion will be marked with this trait and other ops that want to define new scopes can have that trait.

128

Thanks.

166

This was checked as part of isValidSymbol above.

1054

Sure.

2288

Sure, this was really old code - that didn't get updated post ValueRange/TypeRange migrations!

2289

Thanks!

[ftynse]

Once again, this disappeared from my attention list... Is it due to "unresolved" grand-parent diff where I'm not listed as a reviewer?

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

629–630

In the general case, I suppose it can be even worse than that.

func @foo(%arg0: !cheat.opaque, %arg1: memref<..>) { // opaque and memref alias
  cheat.do_cheat %arg0 // this affects the memref
  affine.load %arg1[...]
}

which looks like we either need a powerful and abstract enough way to describe the aliasing between objects of different types, or to treat any side-effecting operation conservatively.

Anyway, I agree with the argument that the proposed approach is no worse than what we already have in Affine and I would like to make progress on this.

[mehdi_amini]

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

2305

You don't need an if here.

2309

Seems like you can fix the FIXME and just write: if(op->isProperAncestor(memref.getDefiningOp()) continue;

[bondhugula]

In D72223#2005186, @ftynse wrote:

Once again, this disappeared from my attention list... Is it due to "unresolved" grand-parent diff where I'm not listed as a reviewer?

That shouldn't happen logically - if it helps, I can add you as a reviewer there as well!

[bondhugula]

Rebase on polyhedral scope trait (simplifies this revision) + address remaining concerns

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

629–630

which looks like we either need a powerful and abstract enough way to describe the
aliasing between objects of different types, or to treat any side-effecting operation

That's right, thanks.

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

2305

You actually do :-) (Notice the continue above is guarded again inside. ) - although I could better use an else here.

2309

Thanks, but the block arg owner check above also has to be fixed similarly to allow it to be any descendent of 'op'. Done - used a lambda with an all_of.

[bondhugula]

Rebase on upstream tip. Bring code to date.

[bondhugula]

Minor update to test cases.

[bondhugula]

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

655

@ftynse Unfortunately, the parser API won't support this unless we want to disallow [] = () and even there we won't be able to emit the right error message. (Basically we won't know whether to parse the =.)

[ftynse]

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

111

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

2280

We can now do builder.createBlock

[bondhugula]

Address review comments.

[bondhugula]

Avoid createBlock

[bondhugula]

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

2280

Unfortunately, using this messes up the insertion point subsequently - outweighs the idiomacy it provides and I'd like to avoid using a scope guard.

[bondhugula]

@mehdi_amini Could you take a look and let me know if you have any concerns with bringing this in? One thing this unblocks is an inliner into an affine.for/if.

[mehdi_amini]

What does this op bring on top of scf.execute_region? Can scf.execute_region carry the AffineScope trait?

Somehow related, in another thread of discussion, we discussed reversing the AffineScope traits: replacing it by a trait that would mean the opposite (let's say NotAnAffineScope as a straw man) and consider every operation that don't define NotAnAffineScope. This would enable to use any op unknown to affine that define a region inside an affine scope, it would just create implicitly a new scope. We could avoid the need to tag operations like FuncOp with AffineScope since that would be the default. Instead only affine operation would need tagging.

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

111

(formatting still to be fixed here)

[mehdi_amini]

Re-read the RFC on "graybox", I am still entirely unconvinced by the current design and the explicit capture of memref used here. I think this is less flexible and not justified all-in-all. believe the more extensible alternative is to support opaque nested regions in general and treat them conservatively as new scope (basically what I wrote above about reversing the trait).
But I'm not willing to die on this hill today, if @ftynse agrees with you that affine should stay a bit more isolated, and that this explicit capture is really worth it right now, then feel free to move forward.

[bondhugula]

In D72223#2845669, @mehdi_amini wrote:

What does this op bring on top of scf.execute_region? Can scf.execute_region carry the AffineScope trait?

scf.execute_region can carry the affinescope trait (or equivalently can be free of the "ExpandAffineScope" trait which is the complement of AffineScope and which is what we want to replace AffineScope with). The real difference between the two is just going to be memref captures.

Somehow related, in another thread of discussion, we discussed reversing the AffineScope traits: replacing it by a trait that would mean the opposite (let's say NotAnAffineScope as a straw man) and consider every operation that don't define NotAnAffineScope. This would enable to use any op unknown to affine that define a region inside an affine scope, it would just create implicitly a new scope. We could avoid the need to tag operations like FuncOp with AffineScope since that would be the default. Instead only affine operation would need tagging.

Yes, we already discussed all of this. Just for the name, "ExpandsAffineScope" is more meaningful since it's adding to affine scope started by an op above that didn't have that trait. Any region holding op that doesn't have ExpandsAffineScope starts such a scope. Only affine.for, affine.if, and affine.parallel will have this trait. It's the complement of the current AffineScope.

[bondhugula]

In D72223#2845698, @mehdi_amini wrote:

Re-read the RFC on "graybox", I am still entirely unconvinced by the current design and the explicit capture of memref used here. I think this is less flexible and not justified all-in-all. believe the more extensible alternative is to support opaque nested regions in general and treat them conservatively as new scope (basically what I wrote above about reversing the trait).
But I'm not willing to die on this hill today, if @ftynse agrees with you that affine should stay a bit more isolated, and that this explicit capture is really worth it right now, then feel free to move forward.

Actually, if you don't want explicit memref captures, you can just use scf.execute_region as is! There would be no difference besides the fact that the latter would be from the scf dialect mixed with affine dialect ops and being looked at by affine passes/utilities etc. for transformation purposes, which is fine. I have felt that having a list of memrefs on the operand list of an affine.execute_region simplifies everything (passes/utilities) doing a walk from the top - they deal with that op transparently just like any other non-region holding op that has memref operands. But I'm willing to hold on to this revision until the use case beyond what can already be served by scf.execute_region arrives. We need to perhaps first think about extending affine passes properly in the presence of scf.execute_region ops and reevaluate this.

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 isolated from above or an affine execute_region. A value of index type
+/// defined at the top level is always a valid symbol.
 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 closest
+/// surrounding op that is isolated from above or an affine execute_region
+/// enclosing this value's definition or block argument appearance.
 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`.
+bool isValidDim(Value value, Region *region);
+
+/// Returns true if the given value can be used as a symbol in the closest
+/// op that is isolated from above or an affine execute_region op enclosing this
+/// value's definition or block argument appearance.
 bool isValidSymbol(Value value);
 
+/// Returns true if the given Value can be used as a symbol for `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.
     bool isValidDim();
 
+    /// Returns true if the result of this operation can be used as dimension id
+    /// within the region 'region'.
+    bool isValidDim(Region *region);
+
     /// Returns true if the result of this operation is a symbol.
     bool isValidSymbol();
 
+    /// Returns true if the result of this operation is a symbol in 'region'.
+    bool isValidSymbol(Region *region);
+
     operand_range getMapOperands() { return getOperands(); }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 def AffineForOp : Affine_Op<"for",
     [ImplicitAffineTerminator, RecursiveSideEffects,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>]> {
   let summary = "for operation";
   let description = [{
     Syntax:
 
     ```
     operation   ::= `affine.for` ssa-id `=` lower-bound `to` upper-bound
                     (`step` integer-literal)? `{` op* `}`
 
     lower-bound ::= `max`? affine-map-attribute dim-and-symbol-use-list | shorthand-bound
     upper-bound ::= `min`? affine-map-attribute dim-and-symbol-use-list | shorthand-bound
def AffineTerminatorOp :
   // No custom parsing/printing form.
   let parser = ?;
   let printer = ?;
 
   // Fully specified by traits.
   let verifier = ?;
 }
 
+def AffineExecuteRegionOp : Affine_Op<"execute_region"> {
+  let summary = "execute_region operation";
+  let description = [{
+      The `affine.execute_region` op introduces a new symbol context for affine
+      operations. It holds a single region,  which can be a list of one or more
+      blocks, and its semantics are to execute its region exactly once. The op's
+      region can have zero or more arguments, each of which can only be a
+      memref. The operands bind 1:1 to its region's arguments. The op can't use
+      any memrefs defined outside of it, but can use any other SSA values that
+      dominate it. The results of a execute_region op match 1:1 with the return
+      values from its region's blocks;
+
+      Examples:
+
+      ```mlir
+      affine.for %i = 0 to 128 {
+        affine.execute_region [%rI, %rM] = (%I, %M)
+              : (memref<128xi32>, memref<24xf32>) -> () {
+          %idx = affine.load %rI[%i] : memref<128xi32>
+          %index = index_cast %idx : i32 to index
+          affine.load %rM[%index]: memref<24xf32>
+          return
+        }
+      }
+      ```
+
+      ```mlir
+      affine.for %i = 0 to %n {
+        affine.execute_region [] = () : () -> () {
+          // %pow can now be used as a loop bound.
+          %pow = call @powi(%i) : (index) ->  index
+          affine.for %j = 0 to %pow {
+            "foo"() : () -> ()
+          }
+          return
+        }
+      }
+      ```
+
+      ```mlir
+      affine.for %i = 0 to %n {
+        affine.execute_region : () -> () {
+          // %pow can now be used as a loop bound.
+          %pow = call @powi(%i) : (index) ->  index
+          affine.for %j = 0 to %pow {
+            "foo"() : () -> ()
+          }
+          return
+        }
+      }
+      ```
+  }];
+
+  let arguments =  (ins Variadic<AnyMemRef>:$operands);
+  let results = (outs Variadic<AnyType>);
+
+  let regions = (region AnyRegion:$region);
+
+  let skipDefaultBuilders = 1;
+
+  let builders = [
+    OpBuilder<"Builder *builder, OperationState &result, "
+              "ValueRange memrefs">
+  ];
+
+  // TODO: canonicalizations related to memrefs.
+  let hasCanonicalizer = 0;
+}
 #endif // AFFINE_OPS

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 given region is attached to an op isolated
+/// from above or an affine scope.
+static bool isIsolatedOrAffineExecuteRegion(Region *region) {
+  return region->getParentOp()->isKnownIsolatedFromAbove() ||
+         isa<AffineExecuteRegionOp>(region->getParentOp());
 }
 
-/// 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 isolated from above or an affine.execute_region op. 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 isFunctionRegion(arg.getOwner()->getParent());
-  return isFunctionRegion(value.getDefiningOp()->getParentRegion());
+    return isIsolatedOrAffineExecuteRegion(arg.getOwner()->getParent());
+  return isIsolatedOrAffineExecuteRegion(
+      value.getDefiningOp()->getParentRegion());
+}
+
+/// A utility function to check if a value is defined at the top level of
+/// `region`. A value of index type defined at the top level is always a
+/// valid symbol.
+static bool isTopLevelValue(Value value, Region *region) {
+  if (auto arg = value.dyn_cast<BlockArgument>())
+    return arg.getOwner()->getParentOp() == region->getParentOp();
+  return value.getDefiningOp()->getParentOp() == region->getParentOp();
+}
+
+/// Returns the closest region surrounding 'op' that is held either by an
+/// AffineExecuteRegionOp or an op isolated from above (eg. FuncOp). Asserts if
+/// called on a top-level op.
+//  TODO: getAffineScope should be publicly exposed for affine passes/utilities.
+static Region *getAffineScope(Operation *op) {
+  auto *curOp = op;
+  while (auto *parentOp = curOp->getParentOp()) {
+    if (llvm::isa<AffineExecuteRegionOp>(parentOp) ||
+        parentOp->isKnownIsolatedFromAbove())
+      return curOp->getParentRegion();
+    curOp = parentOp;
+  }
+
+  llvm_unreachable("op doesn't have an enclosing affine scope");
 }
 
 // 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 *op = value.getDefiningOp())
+    return isValidDim(value, getAffineScope(op));
+
+  // This value has to be a block argument for an op isolated from above or an
+  // affine.for. (An affine.execute_region can't have index type arguments.)
+  auto *parentOp = value.cast<BlockArgument>().getOwner()->getParentOp();
+  return parentOp->isKnownIsolatedFromAbove() || 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 arguments.
+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 a affine.for or
+    // 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 'op'.
 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) {
+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.
+// the top level of the enclosing affine scope (affine.execute_region or an op
+// isolated from above) or dominates such a scope, or it is a result of 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) {
   // 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 *op = value.getDefiningOp())
+    return isValidSymbol(value, getAffineScope(op));
+
+  return false;
+
+// Value can be used as a symbol in `region` if it is a constant, or it is
+// defined at the top level of 'region' or dominates 'region's parent op, or it
+// 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) inside
+// 'region'. 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>

 
 // 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.
+bool AffineApplyOp::isValidDim(Region *region) {
+  return llvm::all_of(getOperands(),
+                      [&](Value op) { return mlir::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 if all its
+// operands are symbols.
+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

 }
 
 LogicalResult AffineDmaWaitOp::verify() {
   if (!getOperand(0).getType().isa<MemRefType>())
     return emitOpError("expected DMA tag to be of memref type");
   for (auto idx : getTagIndices()) {
     if (!idx.getType().isIndex())
       return emitOpError("index to dma_wait must have 'index' type");
-    if (!isValidAffineIndexOperand(idx))
+    Region *scope = getAffineScope(*this);
+    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 (getMemRefType().getRank() != getNumOperands() - 1)
       return emitOpError(
           "expects the number of subscripts to be equal to memref rank");
   }
 
   for (auto idx : getMapOperands()) {
     if (!idx.getType().isIndex())
       return emitOpError("index to load must have 'index' type");
-    if (!isValidAffineIndexOperand(idx))
+    Region *scope = getAffineScope(*this);
+    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 (getMemRefType().getRank() != getNumOperands() - 2)
       return emitOpError(
           "expects the number of subscripts to be equal to memref rank");
   }
 
   for (auto idx : getMapOperands()) {
     if (!idx.getType().isIndex())
       return emitOpError("index to store must have 'index' type");
-    if (!isValidAffineIndexOperand(idx))
+    Region *scope = getAffineScope(*this);
+    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 (mapAttr) {
     if (map.getNumInputs() + 1 != op.getNumOperands())
       return op.emitOpError("too few operands");
   } else {
     if (op.getNumOperands() != 1)
       return op.emitOpError("too few operands");
   }
 
   for (auto idx : op.getMapOperands()) {
-    if (!isValidAffineIndexOperand(idx))
+    Region *scope = getAffineScope(op);
+    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
   results.insert<SimplifyAffineOp<AffinePrefetchOp>>(context);
 }
 
 LogicalResult AffinePrefetchOp::fold(ArrayRef<Attribute> cstOperands,
                                      SmallVectorImpl<OpFoldResult> &results) {
   /// prefetch(memrefcast) -> prefetch
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
+// AffineExecuteRegionOp
+//===----------------------------------------------------------------------===//
+//
+
+// TODO: missing region body.
+void AffineExecuteRegionOp::build(Builder *builder, OperationState &result,
+                                  ValueRange memrefs) {
+  // Create a region and an empty entry block. The arguments of the region are
+  // the supplied memrefs.
+  Region *region = result.addRegion();
+  Block *body = new Block();
+  region->push_back(body);
+  body->addArguments(memrefs.getTypes());
+
+  // Set the operands list as resizable so that we can add memrefs.
+  result.setOperandListToResizable();
+}
+
+static LogicalResult verify(AffineExecuteRegionOp op) {
+  // All memref uses in the execute_region region should be explicitly captured.
+  // FIXME: change this walk to an affine walk that doesn't walk inner
+  // execute_regions.
+  DenseSet<Value> memrefsUsed;
+  op.region().walk([&](Operation *innerOp) {
+    for (auto v : innerOp->getOperands())
+      if (v.getType().isa<MemRefType>())
+        memrefsUsed.insert(v);
+  });
+
+  // For each memref use, ensure either an execute_region argument or a local
+  // def.
+  for (auto memref : memrefsUsed) {
+    if (auto arg = memref.dyn_cast<BlockArgument>())
+      if (arg.getOwner()->getParent()->getParentOp() == op)
+        continue;
+    if (auto *defOp = memref.getDefiningOp())
+      // FIXME: this will only work if the memrefs collected above didn't
+      // include any from inner execute_regions.
+      if (defOp->getParentOfType<AffineExecuteRegionOp>() == op)
+        continue;
+    return op.emitOpError("incoming memref not explicitly captured");
+  }
+
+  // Verify that the region arguments match operands.
+  auto &entryBlock = op.region().front();
+  if (entryBlock.getNumArguments() != op.getNumOperands())
+    return op.emitOpError("region argument count does not match operand count");
+
+  for (auto argEn : llvm::enumerate(entryBlock.getArguments())) {
+    if (op.getOperand(argEn.index()).getType() != argEn.value().getType())
+      return op.emitOpError("region argument ")
+             << argEn.index() << " does not match corresponding operand";
+  }
+
+  return success();
+}
+
+// Custom form syntax.
+//
+// (ssa-id `=`)? `affine.execute_region` (`[` memref-region-arg-list `]`
+//                                   `=` `(` memref-use-list `)`)?
+//                  `:` memref-type-list-parens `->` function-result-type `{`
+//    block+
+// `}`
+//
+// Ex:
+//
+//  affine.execute_region [%rI, %rM] = (%I, %M)
+//        : (memref<128xi32>, memref<1024xf32>) -> () {
+//      %idx = affine.load %rI[%i] : memref<128xi32>
+//      %index = index_cast %idx : i32 to index
+//      affine.load %rM[%index]: memref<1024xf32>
+//      return
+//    }
+//
+static ParseResult parseAffineExecuteRegionOp(OpAsmParser &parser,
+                                              OperationState &result) {
+  // Memref operands.
+  SmallVector<OpAsmParser::OperandType, 4> memrefs;
+
+  // Region arguments to be created.
+  SmallVector<OpAsmParser::OperandType, 4> regionMemRefs;
+
+  // The execute_region op has the same type signature as a function.
+  FunctionType opType;
+
+  // Parse the memref assignments.
+  auto argLoc = parser.getCurrentLocation();
+  if (parser.parseRegionArgumentList(regionMemRefs,
+                                     OpAsmParser::Delimiter::Square) ||
+      parser.parseEqual() ||
+      parser.parseOperandList(memrefs, OpAsmParser::Delimiter::Paren))
+    return failure();
+
+  if (memrefs.size() != regionMemRefs.size())
+    return parser.emitError(parser.getNameLoc(),
+                            "incorrect number of memref captures");
+
+  if (parser.parseColonType(opType) ||
+      parser.addTypesToList(opType.getResults(), result.types))
+    return failure();
+
+  auto memrefTypes = opType.getInputs();
+  if (parser.resolveOperands(memrefs, memrefTypes, argLoc, result.operands))
+    return failure();
+
+  // Introduce and parse body region, and the optional attribute list.
+  Region *body = result.addRegion();
+  if (parser.parseRegion(*body, regionMemRefs, memrefTypes) ||
+      parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+
+  // Set the operands list as resizable so that we can modify operands.
+  result.setOperandListToResizable();
+  return success();
+}
+
+static void print(OpAsmPrinter &p, AffineExecuteRegionOp op) {
+  p << AffineExecuteRegionOp::getOperationName() << " [";
+  // TODO: consider shadowing region arguments.
+  p.printOperands(op.region().front().getArguments());
+  p << "] = (";
+  auto operands = op.getOperands();
+  p.printOperands(operands);
+  p << ") ";
+
+  SmallVector<Type, 4> argTypes(op.getOperandTypes());
+  p << " : "
+    << FunctionType::get(argTypes, op.getResultTypes(),
+                         op.getOperation()->getContext());
+
+  p.printRegion(op.region(),
+                /*printEntryBlockArgs=*/false,
+                /*printBlockTerminators=*/true);
+
+  p.printOptionalAttrDict(op.getAttrs());
+}
+
+//===----------------------------------------------------------------------===//
 // AffineParallelOp
 //===----------------------------------------------------------------------===//
 
 void AffineParallelOp::build(Builder *builder, OperationState &result,
                              ArrayRef<int64_t> ranges) {
   SmallVector<AffineExpr, 8> lbExprs(ranges.size(),
                                      builder->getAffineConstantExpr(0));
   auto lbMap = AffineMap::get(0, 0, lbExprs, builder->getContext());

mlir/test/Dialect/Affine/execute-region.mlir

+// RUN: mlir-opt %s | FileCheck %s
+
+// CHECK-LABEL: @arbitrary_bound
+func @arbitrary_bound(%n : index) {
+  affine.for %i = 0 to %n {
+    affine.execute_region [] = () : () -> () {
+      // %pow can now be used as a loop bound.
+      %pow = call @powi(%i) : (index) ->  index
+      affine.for %j = 0 to %pow {
+        "test.foo"() : () -> ()
+      }
+      return
+    }
+    // CHECK:      affine.execute_region [] = () : () -> () {
+    // CHECK-NEXT:   call @powi
+    // CHECK-NEXT:   affine.for
+    // CHECK-NEXT:     "test.foo"()
+    // CHECK-NEXT:   }
+    // CHECK-NEXT:   return
+    // CHECK-NEXT: }
+  }
+  return
+}
+
+func @powi(index) -> index
+
+// CHECK-LABEL: func @arbitrary_mem_access
+func @arbitrary_mem_access(%I: memref<128xi32>, %M: memref<1024xf32>) {
+  affine.for %i = 0 to 128 {
+    // CHECK: %{{.*}} = affine.execute_region [{{.*}}] = ({{.*}}) : (memref<128xi32>, memref<1024xf32>) -> f32
+    %ret = affine.execute_region [%rI, %rM] = (%I, %M) : (memref<128xi32>, memref<1024xf32>) -> f32 {
+      %idx = affine.load %rI[%i] : memref<128xi32>
+      %index = index_cast %idx : i32 to index
+      %v = affine.load %rM[%index]: memref<1024xf32>
+      return %v : f32
+    }
+  }
+  return
+}
+
+// CHECK-LABEL: @symbol_check
+func @symbol_check(%B: memref<100xi32>, %A: memref<100xf32>) {
+  %cf1 = constant 1.0 : f32
+  affine.for %i = 0 to 100 {
+    %v = affine.load %B[%i] : memref<100xi32>
+    %vo = index_cast %v : i32 to index
+    // CHECK: affine.execute_region [%{{.*}}] = (%{{.*}}) : (memref<100xf32>) -> () {
+    affine.execute_region [%rA] = (%A) : (memref<100xf32>) -> () {
+      // %vi is now a symbol here.
+      %vi = index_cast %v : i32 to index
+      affine.load %rA[%vi] : memref<100xf32>
+      // %vo is also a symbol (dominates the execute_region).
+      affine.load %rA[%vo] : memref<100xf32>
+      return
+    }
+    // CHECK:        index_cast
+    // CHECK-NEXT:   affine.load
+    // CHECK-NEXT:   affine.load
+    // CHECK-NEXT:   return
+    // CHECK-NEXT: }
+  }
+  return
+}
+
+// CHECK-LABEL: func @test_more_symbol_validity
+func @test_more_symbol_validity(%A: memref<100xf32>, %pos : index) {
+  %c5 = constant 5 : index
+  affine.for %i = 0 to 100 {
+    %sym = call @external() : () -> (index)
+    affine.execute_region [%rA] = (%A) : (memref<100xf32>) -> () {
+      affine.load %rA[symbol(%pos) + symbol(%sym) + %c5] : memref<100xf32>
+      return
+    }
+  }
+  affine.execute_region [%rA] = (%A) : (memref<100xf32>) -> () {
+    affine.load %rA[symbol(%pos) + %c5] : memref<100xf32>
+    return
+  }
+  return
+}
+
+func @external() -> (index)
+
+// CHECK-LABEL: func @search
+func @search(%A : memref<?x?xi32>, %S : memref<?xi32>, %key : i32) {
+  %ni = dim %A, 0 : memref<?x?xi32>
+  %c1 = constant 1 : index
+  // This loop can be parallelized.
+  affine.for %i = 0 to %ni {
+    // CHECK: affine.execute_region
+    affine.execute_region [%rA, %rS] = (%A, %S) : (memref<?x?xi32>, memref<?xi32>) -> () {
+      %c0 = constant 0 : index
+      %nj = dim %rA, 1 : memref<?x?xi32>
+      br ^bb1(%c0 : index)
+
+    ^bb1(%j: index):
+      %p1 = cmpi "slt", %j, %nj : index
+      cond_br %p1, ^bb2(%j : index), ^bb5
+
+    ^bb2(%j_arg : index):
+      %v = affine.load %rA[%i, %j_arg] : memref<?x?xi32>
+      %p2 = cmpi "eq", %v, %key : i32
+    cond_br %p2, ^bb3(%j_arg : index), ^bb4(%j_arg : index)
+
+    ^bb3(%j_arg2: index):
+      %j_int = index_cast %j_arg2 : index to i32
+      affine.store %j_int, %rS[%i] : memref<?xi32>
+      br ^bb5
+
+    ^bb4(%j_arg3 : index):
+      %jinc = addi %j_arg3, %c1 : index
+      br ^bb1(%jinc : index)
+
+    ^bb5:
+      return
+    }
+  }
+  return
+}

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>) {
func @affine_min(%arg0 : index, %arg1 : index, %arg2 : index) {
   // expected-error@+1 {{operand count and affine map dimension and symbol count must match}}
   %0 = affine.min affine_map<(d0) -> (d0)> ()
 
   return
 }
 
 // -----
 
+// CHECK-LABEL: @affine.execute_region_missing_capture
+func @affine.execute_region_missing_capture(%M : memref<2xi32>) {
+  affine.for %i = 0 to 10 {
+    affine.execute_region [] = () : () -> () {
+    // expected-error@-1 {{incoming memref not explicitly captured}}
+      affine.load %M[%i] : memref<2xi32>
+    }
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @affine.execute_region_wrong_capture
+func @affine.execute_region_wrong_capture(%s : index) {
+  affine.execute_region [%rS] = (%s) : (index) -> () {
+    // expected-error@-1 {{operand #0 must be memref}}
+    "use"(%s) : (index) -> ()
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @affine.execute_region_wrong_capture
+func @affine.execute_region_wrong_capture(%A : memref<2xi32>) {
+  affine.execute_region [] = (%A) : (memref<2xi32>) -> () {
+    // expected-error@-1 {{incorrect number of memref captures}}
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @affine.execute_region_region_type_mismatch
+func @affine.execute_region_region_type_mismatch(%A : memref<2xi32>) {
+  "affine.execute_region"(%A) ({
+    // expected-error@-1 {{region argument 0 does not match corresponding operand}}
+    ^bb0(%rA : memref<4xi32>):
+      return
+  }) : (memref<2xi32>) -> ()
+}
+
+// -----
+
+// CHECK-LABEL: @affine.execute_region_region_arg_count_mismatch
+func @affine.execute_region_region_arg_count_mismatch(%A : memref<2xi32>) {
+  "affine.execute_region"(%A) ({
+    // expected-error@-1 {{region argument count does not match operand count}}
+    ^bb0:
+      return
+  }) : (memref<2xi32>) -> ()
+  return
+}
+
+// -----
+
 // CHECK-LABEL: @affine_max
 func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
   // expected-error@+1 {{operand count and affine map dimension and symbol count must match}}
   %0 = affine.max affine_map<(d0) -> (d0)> (%arg0, %arg1)
 
   return
 }