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

[Christian Sigg via Phabricator via llvm-commits]

csigg added a comment.

Thanks a lot for the review, Alex!



================
Comment at: mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp:29
+
+  explicit GpuAllReduceRewriter(gpu::GPUFuncOp funcOp_,
+                                gpu::AllReduceOp reduceOp_,
----------------
ftynse wrote:
> Nit: `explicit` is unnecessary for multi-argument constructors.
Removed. 

Nit the nit: `explicit` does prevent copy-initialization from initializer list.


================
Comment at: mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp:88
+    // The number of active threads in the warp, not clamped to 32.
+    Value *activeWidth =
+        create<SubIOp>(loc, blockSize, numThreadsWithSmallerWarpId);
----------------
ftynse wrote:
> Not sure I follow here. `numThreadsWithSmallerWarpId` is the equivalent of `warp_id * warp_size`, or `floor(linear_thread_id, warp_size) * warp_size`. Subtracting that from `blockSize` will give you the number of threads in all warps starting from the current warp. From skimming through the consumer (createWarpReduce), it looks like what it expects is the number of threads in the _current_ warp.
Correct. The activeWidth does not need to be clamped to subgroup width though. I added two comments.


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


================
Comment at: mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp:243
+
+    auto addBranch = [&](ValueRange operands) {
+      create<BranchOp>(loc, continueBlock, operands);
----------------
ftynse wrote:
> Is this lambda worse it? 
Haha, I like your comment. Very subtle ;-) 


================
Comment at: mlir/test/Dialect/GPU/all-reduce.mlir:7
+  gpu.func @kernel(%arg0 : f32) attributes { gpu.kernel } {
+    // CHECK:   %[[vc0_i32:[a-z_0-9]+]] = constant 0 : i32
+    // CHECK:   %[[vc31_i32:[a-z_0-9]+]] = constant 31 : i32
----------------
ftynse wrote:
> ftynse wrote:
> > csigg wrote:
> > > These CHECKs were generated from the output with:
> > > 
> > > ```
> > > sed -r \
> > > -e 's|\t+//.*||' \
> > > -e 's|%([a-z_0-9]+) = |%[[v\1:[a-z_0-9]+]] = |' \
> > > -e 's|\(%([a-z_0-9]+): ([a-z_0-9]+)\):|(%[[v\1:[a-z_0-9]+]]: \2):|' \
> > > -e 's|%([a-z_0-9]+)|%[[v\1]]|g' \
> > > -e 's|bb([0-9]+)|bb[[#b\1]]|g' \
> > > -e 's|^ |    // CHECK:|'
> > > ```
> > > 
> > > and manaul edits from there.
> > > 
> > We have https://github.com/llvm/llvm-project/blob/master/mlir/utils/generate-test-checks.py, have you tried it?
> Is there a reason why you can't use `.*` as a regexp for names? Seeing the full with ranges makes it hard to read the test. I think we also support capital letters in the names.
Replaced with generate-test-checks.py result.


================
Comment at: mlir/test/Dialect/GPU/all-reduce.mlir:7-176
+    // CHECK:   %[[vc0_i32:[a-z_0-9]+]] = constant 0 : i32
+    // CHECK:   %[[vc31_i32:[a-z_0-9]+]] = constant 31 : i32
+    // CHECK:   %[[vc0:[a-z_0-9]+]] = constant 0 : index
+    // CHECK:   %[[vc32_i32:[a-z_0-9]+]] = constant 32 : i32
+    // CHECK:   %[[vc1_i32:[a-z_0-9]+]] = constant 1 : i32
+    // CHECK:   %[[vc2_i32:[a-z_0-9]+]] = constant 2 : i32
+    // CHECK:   %[[vc4_i32:[a-z_0-9]+]] = constant 4 : i32
----------------
csigg wrote:
> ftynse wrote:
> > ftynse wrote:
> > > csigg wrote:
> > > > These CHECKs were generated from the output with:
> > > > 
> > > > ```
> > > > sed -r \
> > > > -e 's|\t+//.*||' \
> > > > -e 's|%([a-z_0-9]+) = |%[[v\1:[a-z_0-9]+]] = |' \
> > > > -e 's|\(%([a-z_0-9]+): ([a-z_0-9]+)\):|(%[[v\1:[a-z_0-9]+]]: \2):|' \
> > > > -e 's|%([a-z_0-9]+)|%[[v\1]]|g' \
> > > > -e 's|bb([0-9]+)|bb[[#b\1]]|g' \
> > > > -e 's|^ |    // CHECK:|'
> > > > ```
> > > > 
> > > > and manaul edits from there.
> > > > 
> > > We have https://github.com/llvm/llvm-project/blob/master/mlir/utils/generate-test-checks.py, have you tried it?
> > Is there a reason why you can't use `.*` as a regexp for names? Seeing the full with ranges makes it hard to read the test. I think we also support capital letters in the names.
> Replaced with generate-test-checks.py result.
Works very well, I wish I knew about this (well, Mehdi mentioned it but I couldn't find it).


Repository:
  rG LLVM Github Monorepo

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

https://reviews.llvm.org/D72129