Page MenuHomePhabricator

[mlir][gpu] Add pass to make GPU ops within a region execute asynchronously.
ClosedPublic

Authored by csigg on Oct 22 2020, 1:05 AM.

Details

Summary

Do not use the pass yet, except in a test.

Diff Detail

Unit TestsFailed

TimeTest
40 mslinux > MLIR.Dialect/GPU::async-region.mlir
Script: -- : 'RUN: at line 1'; /mnt/disks/ssd0/agent/llvm-project/build/bin/mlir-opt -gpu-async-region /mnt/disks/ssd0/agent/llvm-project/mlir/test/Dialect/GPU/async-region.mlir | /mnt/disks/ssd0/agent/llvm-project/build/bin/FileCheck /mnt/disks/ssd0/agent/llvm-project/mlir/test/Dialect/GPU/async-region.mlir
130 mswindows > MLIR.Dialect/GPU::async-region.mlir
Script: -- : 'RUN: at line 1'; c:\ws\w16-1\llvm-project\premerge-checks\build\bin\mlir-opt.exe -gpu-async-region C:\ws\w16-1\llvm-project\premerge-checks\mlir\test\Dialect\GPU\async-region.mlir | c:\ws\w16-1\llvm-project\premerge-checks\build\bin\filecheck.exe C:\ws\w16-1\llvm-project\premerge-checks\mlir\test\Dialect\GPU\async-region.mlir

Event Timeline

csigg created this revision.Oct 22 2020, 1:05 AM
csigg requested review of this revision.Oct 22 2020, 1:05 AM
herhut requested changes to this revision.Oct 22 2020, 6:44 AM
herhut added inline comments.
mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
111

I find it somewhat hard to read that this is a carried state. A better name or comment would already help a lot.

133

Is it useful to formulate this as a pattern? It could equally well be just a walk over all operations, as there is no matching going on, at all.

This revision now requires changes to proceed.Oct 22 2020, 6:44 AM
csigg marked an inline comment as done.Oct 22 2020, 10:50 AM
csigg added inline comments.
mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
133

I'm using notifyMatchFailure, and maybe/theoretically there could be different patterns in the same pass. But to be honest I mainly used RewritePattern because I vaguely remember that River didn't like the raw walk in the all-reduce lowering. If you think a raw walk is better here, I'm happy to change it.

rriddle added inline comments.Oct 22 2020, 11:03 AM
mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
133

I don't exactly remember the details there. A normal walk is fine, but I tend to err more towards patterns because they compose better/have better logging/remove the need to write transformation application logic when things are more complicated than a simple walk. Not to say that you shouldn't just use a simple walk here, but I'd say it depends on how this pass is intended to evolve.

mehdi_amini added inline comments.Oct 22 2020, 4:14 PM
mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
133

Because there is state in the pattern keeping a value, this is to me a strong indication that we should not use a pattern here. This is just hiding the traversal logic, while it seems core to the algorithm.

csigg updated this revision to Diff 300170.Oct 22 2020, 11:28 PM

Use raw walk instead of a RewritePattern.

This does seem a little cleaner, at least at the moment:

  • It's easier to make the pass fail when some preconditions aren't met.
  • The rewriter.replaceOp() assumes that we replace the root, which is not what we did.
sanjoy.google added inline comments.Oct 23 2020, 12:02 AM
mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
39

The more common term for this is to "thread the token through", WDYT?

80

Is this correct if the region has cycles or diamond control flow?

csigg added inline comments.Oct 23 2020, 12:36 PM
mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
39

Thanks for the suggestion. Yes, that sound better. I will include that change in my next update (it's always a bit of a dance, so I'm waiting if there are more requested changes).

80

I thought control flow requires separate blocks, each with their own terminator (which triggers a host synchronization and clears the currentToken). Is that not correct?

133

I've removed the pattern, but it did not contain state. The Callback containing the state was instantiated inside the matchAndRewrite().

csigg updated this revision to Diff 300609.Oct 26 2020, 2:53 AM

Error out on encountering pre-existing gpu.wait.

sanjoy.google added inline comments.Oct 28 2020, 10:32 PM
mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
80

You're right, I didn't notice that bit.

herhut accepted this revision.Oct 29 2020, 9:06 AM

Thanks!

This revision is now accepted and ready to land.Oct 29 2020, 9:06 AM
csigg updated this revision to Diff 301739.Oct 29 2020, 1:53 PM

Rebase.

This revision was landed with ongoing or failed builds.Oct 29 2020, 2:18 PM
This revision was automatically updated to reflect the committed changes.