[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

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

[bondhugula]

mlir/docs/Dialects/Affine.md

612

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

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

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

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

630–631 ↗

(On Diff #258551)

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

632–635 ↗

(On Diff #258551)

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...

656 ↗

(On Diff #258551)

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

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

108 ↗

(On Diff #258551)

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

121 ↗

(On Diff #258551)

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

128 ↗

(On Diff #258551)

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

166–167 ↗

(On Diff #258551)

I cannot relate this comment to the code below.

169 ↗

(On Diff #258551)

Shouldn't this also check for "KnownIsolatedFromAbove" ?

1068 ↗

(On Diff #258551)

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

2290 ↗

(On Diff #258551)

Prefer ValueRange to ArrayRef<Value>

2299 ↗

(On Diff #258551)

Nit: drop trivial braces

2302 ↗

(On Diff #258551)

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

2303 ↗

(On Diff #258551)

You already pushed it in line 2295

2339 ↗

(On Diff #258551)

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

2341 ↗

(On Diff #258551)

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

2404 ↗

(On Diff #258551)

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

630–631 ↗

(On Diff #258551)

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.

632–635 ↗

(On Diff #258551)

Sure, dropped these lines.

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

108 ↗

(On Diff #258551)

Actually, it is - thanks!

121 ↗

(On Diff #258551)

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 ↗

(On Diff #258551)

Thanks.

169 ↗

(On Diff #258551)

This was checked as part of isValidSymbol above.

1068 ↗

(On Diff #258551)

Sure.

2302 ↗

(On Diff #258551)

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

2303 ↗

(On Diff #258551)

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

630–631 ↗

(On Diff #258551)

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 ↗

(On Diff #259980)

You don't need an if here.

2309 ↗

(On Diff #259980)

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

630–631 ↗

(On Diff #258551)

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 ↗

(On Diff #259980)

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

2309 ↗

(On Diff #259980)

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

656 ↗

(On Diff #258551)

@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

120 ↗

(On Diff #353183)

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

2673 ↗

(On Diff #353183)

We can now do builder.createBlock

[bondhugula]

Address review comments.

[bondhugula]

Avoid createBlock

[bondhugula]

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

2673 ↗

(On Diff #353183)

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

120 ↗

(On Diff #353183)

(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/docs/Dialects/Affine.md

 #affine_map2to3 = (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`
+identifiers to enable powerful analysis and transformation. An SSA value is a
+valid symbol if it is either (1) a region argument for an op that is either
+"isolated from above" (like the FuncOp) or is an affine graybox op, (2) a value
+defined at the top level of (outside of all loops, if operations, or other
+operations with regions) of an affine graybox op or an op "isolated from above",
+(3) a value that dominates the closest enclosing affine graybox or an op
+"isolated from above", (4) the result of a [`constant`
+operation](Standard.md#constant-operation), (4) the result of an [`affine.apply`
+operation](#affineapply-operation) that recursively takes as arguments any
+symbolic identifiers, or (5) 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` operations](#affinefor-operation), and the result of an
-[`affine.apply` operation](#affineapply-operation) (which recursively may use
-other dimensions and symbols).
+dynamic one in turn bound to a symbolic identifier. Note that as a result of
+(3), symbol validity is sensitive to the location at which the value binds to
+the symbol. Dimensions may be bound not only to anything that a symbol is bound
+to, but also to induction variables of enclosing [`affine.for`
+operations](#affinefor-operation), and the result of an [`affine.apply`
+operation](#affineapply-operation) (which recursively may use other dimensions
+and symbols).
 
 ### Affine Expressions
 
 Syntax:
 
 ```
 affine-expr ::= `(` affine-expr `)`
               | affine-expr `+` affine-expr

 
 Affine terminator is a special terminator operation for blocks inside affine
 loops ([`affine.for`](#affinefor-operation)) and branches
 ([`affine.if`](#affineif-operation)). It unconditionally transmits the control
 flow to the successor of the operation enclosing the region.
 
 *Rationale*: bodies of affine operations are [blocks](../LangRef.md#blocks) that
 must have terminators. Loops and branches represent structured control flow and
 should not accept arbitrary branches as terminators.
 
 This operation does _not_ have a custom syntax. However, affine control
 operations omit the terminator in their custom syntax for brevity.

mlir/include/mlir/Dialect/AffineOps/AffineOps.h

 namespace mlir {
 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 graybox. A value of index type defined
+/// at the top level is always a valid symbol.
 bool isTopLevelValue(Value value);
 
 class AffineOpsDialect : public Dialect {
 public:
   AffineOpsDialect(MLIRContext *context);
   static StringRef getDialectNamespace() { return "affine"; }
 
   /// Materialize a single constant operation from a given attribute value with
public:
   static void build(Builder *builder, OperationState &result, AffineMap map,
                     ValueRange operands);
 
   /// Returns the affine map to be applied by this operation.
   AffineMap getAffineMap() {
     return getAttrOfType<AffineMapAttr>("map").getValue();
   }
 
-  /// 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 its immediately surrounding affine scope.
   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 the region of the op 'opWithRegion'.
+  bool isValidDim(Operation *opWithRegion);
+
+  /// Returns true if the result of this operation is a symbol in its
+  /// immediately surrounding affine scope.
   bool isValidSymbol();
 
+  /// Returns true if the result of this operation is a symbol in the region of
+  /// 'opWithRegion'.
+  bool isValidSymbol(Operation *opWithRegion);
+
   static StringRef getOperationName() { return "affine.apply"; }
 
   operand_range getMapOperands() { return getOperands(); }
 
   // Hooks to customize behavior of this op.
   static ParseResult parse(OpAsmParser &parser, OperationState &result);
   void print(OpAsmPrinter &p);
   LogicalResult verify();
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
+/// op that is isolated from above or an affine graybox 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 for an op with
+/// a region.
+bool isValidDim(Value value, Operation *opWithRegion);
+
+/// Returns true if the given value can be used as a symbol in the closest
+/// op that is isolated from above or an affine graybox 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 an
+/// op with a region.
+bool isValidSymbol(Value value, Operation *opWithRegion);
+
 /// 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/AffineOps/AffineOps.td

def AffineTerminatorOp :
   // No custom parsing/printing form.
   let parser = ?;
   let printer = ?;
 
   // Fully specified by traits.
   let verifier = ?;
 }
 
+def AffineGrayBoxOp : Affine_Op<"graybox">,
+    Arguments<(ins Variadic<AnyMemRef>:$operands)> {
+  let summary = "graybox operation";
+  let description =
+      "The affine graybox op introduces a new symbol context for affine "
+      "operations. It holds a single region,  region's blocks. which can be a "
+      "list of one or more blocks. 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. Its region's "
+      "blocks can have terminators the same way as current MLIR functions "
+      "(FuncOp) can. Control from any return ops from the top level of its "
+      "region returns to right after the affine.graybox op. Its control flow "
+      "thus conforms to the control flow semantics of regions, i.e., control "
+      "always returns to the immediate enclosing (parent) op. The results of a "
+      "graybox op match 1:1 with the return values from its region's blocks";
+
+  let regions = (region AnyRegion:$region);
+
+  // TODO: builders.
+
+  // TODO: canonicalizations related to memrefs.
+  let hasCanonicalizer = 0;
+}
+
 #endif // AFFINE_OPS

mlir/lib/Dialect/AffineOps/AffineOps.cpp

 /// Materialize a single constant operation from a given attribute value with
 /// the desired resultant type.
 Operation *AffineOpsDialect::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 graybox.
+static bool isIsolatedOrGrayBoxRegion(Region *region) {
+  return region->getParentOp()->isKnownIsolatedFromAbove() ||
+         isa<AffineGrayBoxOp>(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 graybox. 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 isIsolatedOrGrayBoxRegion(arg->getOwner()->getParent());
+  return isIsolatedOrGrayBoxRegion(value->getDefiningOp()->getParentRegion());
+}
+
+/// A utility function to check if a value is defined at the top level of
+/// 'opWithRegion'. A value of index type defined at the top level is always a
+/// valid symbol.
+static bool isTopLevelValue(Value value, Operation *opWithRegion) {
+  assert(opWithRegion->getNumRegions() > 0 &&
+         "only to be called on ops with regions");
+  if (auto arg = value.dyn_cast<BlockArgument>())
+    return arg->getOwner()->getParentOp() == opWithRegion;
+  return value->getDefiningOp()->getParentOp() == opWithRegion;
+}
+
+/// Returns the closest op surrounding 'op' that is either an AffineGrayBoxOp 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/utitlies.
+static Operation *getAffineScope(Operation *op) {
+  // TODO: make this compact by introducing a variadic pack on getParentOfType.
+  auto *curOp = op;
+  while ((curOp = curOp->getParentOp()))
+    if (llvm::isa<AffineGrayBoxOp>(curOp) || curOp->isKnownIsolatedFromAbove())
+      return curOp;
+
+  assert(false && "op doesn't have a parent op");
+  return nullptr;
 }
 
 // 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 for a FuncOp or an affine.for.
-  auto *parentOp = value.cast<BlockArgument>()->getOwner()->getParentOp();
-  return isa<FuncOp>(parentOp) || isa<AffineForOp>(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. (A graybox can't have index type arguments.)
+  auto *parentOp = value.cast<BlockArgument>()->getOwner()->getParentOp();
+  return parentOp->isKnownIsolatedFromAbove() || isa<AffineForOp>(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 affine apply operation with
+// dimension id arguments.
+bool mlir::isValidDim(Value value, Operation *opWithRegion) {
+  assert(opWithRegion->getNumRegions() > 0 &&
+         "only to be called on ops with regions");
+  // The value must be an index type.
+  if (!value->getType().isIndex())
+    return false;
+
+  auto *op = value->getDefiningOp();
+  if (!op) {
+    // This value has to be a block argument for a FuncOp or an affine.for.
+    auto *parentOp = value.cast<BlockArgument>()->getOwner()->getParentOp();
+    return parentOp->isKnownIsolatedFromAbove() || isa<AffineForOp>(parentOp);
+  }
+
+  if (::isTopLevelValue(value, opWithRegion) || isa<ConstantOp>(op))
+    return applyOp.isValidDim(opWithRegion);
 
 /// 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>
-bool isMemRefSizeValidSymbol(AnyMemRefDefOp memrefDefOp, unsigned index) {
+bool isMemRefSizeValidSymbol(AnyMemRefDefOp memrefDefOp, unsigned index,
+                             Operation *op) {
+  assert(op->getNumRegions() > 0 && "only to be called on ops with regions");
   auto memRefType = memrefDefOp.getType();
   // Statically shaped.
   if (!ShapedType::isDynamic(memRefType.getDimSize(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),
+                       op);
 }
 
 /// Returns true if the result of the dim op is a valid symbol.
-static bool isDimOpValidSymbol(DimOp dimOp) {
+static bool isDimOpValidSymbol(DimOp dimOp, Operation *op) {
+  assert(op->getNumRegions() > 0 && "only to be called on ops with regions");
   // 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, op);
   if (auto subViewOp = dyn_cast<SubViewOp>(dimOp.getOperand()->getDefiningOp()))
-    return isMemRefSizeValidSymbol<SubViewOp>(subViewOp, index);
+    return isMemRefSizeValidSymbol<SubViewOp>(subViewOp, index, op);
   if (auto allocOp = dyn_cast<AllocOp>(dimOp.getOperand()->getDefiningOp()))
-    return isMemRefSizeValidSymbol<AllocOp>(allocOp, index);
+    return isMemRefSizeValidSymbol<AllocOp>(allocOp, index, op);
   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 (graybox or func op) 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);
-    }
-  }
-  // Otherwise, check that the value is a top level value.
-  return isTopLevelValue(value);
+  // 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 the region of an op 'opWithRegion' if it is
+// a constant, or it is defined at the top level of 'opWithRegion' or dominates
+// 'opWithRegion', 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, Operation *opWithRegion) {
+  assert(opWithRegion->getNumRegions() > 0 &&
+         "only to be called on ops with regions");
+
+  // The value must be an index type.
+  if (!value->getType().isIndex())
+    return false;
+
+  // A top-level value is a valid symbol.
+  if (::isTopLevelValue(value, opWithRegion))
+    return true;
+
+  auto *defOp = value->getDefiningOp();
+  if (!defOp)
+    // A block argument that is not a top-level value isn't a valid symbol.
+    return false;
+
+  // 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(opWithRegion);
+
+  // Dim op results could be valid symbols at any level.
+  if (auto dimOp = dyn_cast<DimOp>(defOp))
+    return isDimOpValidSymbol(dimOp, opWithRegion);
+
+  // Check for values dominating 'opWithRegion'.
+  if (auto *parentOp = opWithRegion->getParentOp())
+    if (!parentOp->isKnownIsolatedFromAbove())
+      return isValidSymbol(value, parentOp);
+
+  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 the
+// region of 'op'. Returns false otherwise.
+static bool isValidAffineIndexOperand(Value value, Operation *op) {
+  return isValidDim(value, op) || isValidSymbol(value, op);
 }
 
 /// 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(Operation *opWithRegion) {
+  return llvm::all_of(getOperands(), [&](Value op) {
+    return mlir::isValidDim(op, opWithRegion);
+  });
+}
+
 // 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(Operation *opWithRegion) {
+  return llvm::all_of(getOperands(), [&](Value operand) {
+    return mlir::isValidSymbol(operand, opWithRegion);
+  });
+}
+
 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");
   }
 
+  auto *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))
+    auto *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))
+
+    auto *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))
+    auto *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))
+    auto *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);
 }
 
 //===----------------------------------------------------------------------===//
+// AffineGrayBoxOp
+//===----------------------------------------------------------------------===//
+//
+
+static LogicalResult verify(AffineGrayBoxOp op) {
+  // All memref uses in the graybox region should be explicitly captured.
+  // FIXME: change this walk to an affine walk that doesn't walk inner
+  // grayboxes.
+  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 a graybox argument or locally defined.
+  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 grayboxes.
+      if (defOp->getParentOfType<AffineGrayBoxOp>() == op)
+        continue;
+    return op.emitOpError("incoming memref not explicitly captured");
+  }
+  return success();
+}
+
+static ParseResult parseAffineGrayBoxOp(OpAsmParser &parser,
+                                        OperationState &result) {
+  // Sizes of the grid and block.
+  SmallVector<OpAsmParser::OperandType, 4> memrefs;
+
+  // Region arguments to be created.
+  SmallVector<OpAsmParser::OperandType, 4> regionMemRefs;
+
+  auto argLoc = parser.getCurrentLocation();
+
+  // Parse the memref assignments.
+  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");
+
+  SmallVector<Type, 4> memrefTypes;
+  if (parser.parseOptionalColonTypeList(memrefTypes))
+    return failure();
+
+  if (parser.resolveOperands(memrefs, memrefTypes, argLoc, result.operands))
+    return failure();
+
+  // Introduce the body region and parse it.  The region has
+  // kNumConfigRegionAttributes leading arguments that correspond to
+  // block/thread identifiers and grid/block sizes, all of the `index` type.
+  // Follow the actual kernel arguments.
+  Region *body = result.addRegion();
+  if (parser.parseRegion(*body, regionMemRefs, memrefTypes) ||
+      parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+
+  // Parse the optional attribute list.
+  if (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, AffineGrayBoxOp op) {
+  p << AffineGrayBoxOp::getOperationName() << " [";
+  p.printOperands(op.region().front().getArguments());
+  p << "] = (";
+  auto operands = op.getOperands();
+  p.printOperands(operands);
+  p << ") ";
+  if (!operands.empty())
+    p << ": " << operands.getTypes();
+
+  p.printRegion(op.region(),
+                /*printEntryBlockArgs=*/false,
+                /*printBlockTerminators=*/true);
+
+  p.printOptionalAttrDict(op.getAttrs());
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/AffineOps/AffineOps.cpp.inc"

mlir/test/AffineOps/graybox.mlir

+// RUN: mlir-opt %s | FileCheck %s
+
+// CHECK-LABEL: @arbitrary_bound
+func @arbitrary_bound(%n : index) {
+  affine.for %i = 0 to %n {
+    affine.graybox [] = () {
+      // %pow can now be used as a loop bound.
+      %pow = call @powi(%i) : (index) ->  index
+      affine.for %j = 0 to %pow {
+        "foo"() : () -> ()
+      }
+      return
+    }
+    // CHECK:      affine.graybox [] = () {
+    // CHECK-NEXT:   call @powi
+    // CHECK-NEXT:   affine.for
+    // CHECK-NEXT:     "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.graybox [{{.*}}] = ({{.*}}) : memref<128xi32>, memref<1024xf32>
+    affine.graybox [%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
+    }
+  }
+  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.graybox [%{{.*}}] = (%{{.*}}) : memref<100xf32> {
+    affine.graybox [%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 graybox).
+      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 @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.graybox
+    affine.graybox [%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/AffineOps/ops.mlir

affine.for %arg4 = 0 to %13 step 264 {
       %24 = dim %20, 0 : memref<?x?xf32, (d0, d1)[s0, s1, s2] -> (d0 * s1 + d1 * s2 + s0)>
       affine.for %arg5 = 0 to %24 step 768 {
         "foo"() : () -> ()
       }
     }
   }
   return
 }
+
+// -----
+
+// Test symbol restrictions with ops isolated from above.
+
+// CHECK-LABEL: func @valid_symbol_isolated_region
+func @valid_symbol_isolated_region(%n : index) {
+  test.isolated_region %n {
+    %c1 = constant 1 : index
+    %l = subi %n, %c1 : index
+    // %l, %n are valid symbols since test.isolated_region is known to be
+    // isolated from above.
+    affine.for %i = %l to %n {
+    }
+    return
+  }
+  return
+}