[PATCH] D72129: [mlir] Add in-dialect lowering of gpu.all_reduce.

[Alex Zinenko via Phabricator via llvm-commits]

ftynse requested changes to this revision.
ftynse marked an inline comment as done.
ftynse added a comment.
This revision now requires changes to proceed.

Almost there, only minor things necessary to improve understanding.

The approach you took for building code is very similar to our motivation for declarative builders (aka EDSC), maybe it's time to reconsider those again.



================
Comment at: mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp:150
+    auto bufferType =
+        MemRefType::get({kWarpSize}, valueType, ArrayRef<AffineMap>{}, 3);
+    return funcOp.addWorkgroupAttribution(bufferType);
----------------
csigg wrote:
> ftynse wrote:
> > Nit: let's have a named constant for the address space 
> Punted by just adding a local variable.
I'll add that later anyway :)


================
Comment at: mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp:268
+        [&] {
+          predicatedOpsFactory();
+          return ArrayRef<Value >();
----------------
I'd add an assertion that `predicatedOpsFactory()` does not return any values that it might expect to be passed to the continuation.


================
Comment at: mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp:354
+    };
+    while (funcOp.walk(callback).wasInterrupted()) {
+    }
----------------
Could you please add a comment on why this approach is necessary?  I understand that you need to operate on the GPUFuncOp level because you modify the GPUFuncOp itself, which would be incorrect from a nested operation (AllReduceOp). I don't understand why do you need to interrupt the walk after each rewrite. Is it because of some state invalidation?


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D72129/new/

https://reviews.llvm.org/D72129