Page MenuHomePhabricator

[MLIR] Introduce affine graybox op
Needs ReviewPublic

Authored by bondhugula on Jan 4 2020, 9:41 PM.

Details

Summary
- to allow the polyhedral form to be used in a wider context without the
  need for function outlining

- expand affine symbol rules to account for affine graybox ops as well as
  any op that is isolated from above; the rules for valid symbols for
  polyhedral purposes were based on FuncOp alone; generalize this to any
  op isolated from above; allows affine.for/if to be used in other
  function like ops.

Signed-off-by: Uday Bondhugula <uday@polymagelabs.com>

Diff Detail

Unit TestsFailed

TimeTest
20 msMLIR.Dialect/Affine::graybox.mlir
Script: -- : 'RUN: at line 1'; /mnt/disks/ssd0/agent/workspace/amd64_debian_testing_clang8/build/bin/mlir-opt -allow-unregistered-dialect /mnt/disks/ssd0/agent/workspace/amd64_debian_testing_clang8/mlir/test/Dialect/Affine/graybox.mlir | /mnt/disks/ssd0/agent/workspace/amd64_debian_testing_clang8/build/bin/FileCheck /mnt/disks/ssd0/agent/workspace/amd64_debian_testing_clang8/mlir/test/Dialect/Affine/graybox.mlir
40 msMLIR.Dialect/Affine::ops.mlir
Script: -- : 'RUN: at line 1'; /mnt/disks/ssd0/agent/workspace/amd64_debian_testing_clang8/build/bin/mlir-opt -allow-unregistered-dialect -split-input-file /mnt/disks/ssd0/agent/workspace/amd64_debian_testing_clang8/mlir/test/Dialect/Affine/ops.mlir | /mnt/disks/ssd0/agent/workspace/amd64_debian_testing_clang8/build/bin/FileCheck /mnt/disks/ssd0/agent/workspace/amd64_debian_testing_clang8/mlir/test/Dialect/Affine/ops.mlir

Event Timeline

bondhugula created this revision.Jan 4 2020, 9:41 PM
bondhugula updated this revision to Diff 236224.Jan 4 2020, 9:42 PM
bondhugula removed a reviewer: nicolasvasilache.
  • rename getParentAffineScope -> getAffineScope

Update doc comments.

  • invalid test cases
  • update affine dialect doc
bondhugula updated this revision to Diff 236231.Jan 5 2020, 3:01 AM
  • fix comments
bondhugula updated this revision to Diff 236236.Jan 5 2020, 4:59 AM
  • add missed result parsing
  • fix verifier
mehdi_amini requested changes to this revision.Jan 5 2020, 5:49 AM
mehdi_amini added inline comments.
mlir/docs/Dialects/Affine.md
454

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

This revision now requires changes to proceed.Jan 5 2020, 5:49 AM
bondhugula marked an inline comment as done.Jan 5 2020, 7:06 AM
bondhugula added inline comments.
mlir/docs/Dialects/Affine.md
454

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 added inline comments.Jan 5 2020, 2:35 PM
mlir/docs/Dialects/Affine.md
454

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 marked an inline comment as done.Jan 5 2020, 8:31 PM
bondhugula added inline comments.
mlir/docs/Dialects/Affine.md
454

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 marked an inline comment as done.Jan 6 2020, 12:59 AM
bondhugula added inline comments.
mlir/docs/Dialects/Affine.md
454

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 retitled this revision from [LLVM] [MLIR] Introduce affine graybox op to [MLIR] Introduce affine graybox op.Tue, Mar 24, 6:35 AM
bondhugula edited the summary of this revision. (Show Details)

Rebase.

changes after rebase - unregistered ops