Add gpu::LaunchOp::addKernelArgument.

[herhut]

Summary

Code motion into a gpu::LaunchOp region requires operands of moved
instructions to be threaded through operands of the gpu.launch to
ensure that the gpu.launch remaines closed from above.

To enable this, gpu.launch operations are now created with extensible
operand storage. The overhead is expected to be low given that
gpu.launch is a relatively rare operation.

[merge_guards_bot]

Unit tests: fail. 62155 tests passed, 5 failed and 811 were skipped.

failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[ftynse]

mlir/include/mlir/Dialect/GPU/GPUOps.td

488

This says "gpu.launch" but the line above says "gpu.func". Let's use one everywhere and say the other is equivalent.

[herhut]

mlir/include/mlir/Dialect/GPU/GPUOps.td

488

But they are not. When used inside of a launch, it cannot have operands. I could maybe state that gpu.launch is considered equivalent to a void function? I found it surprising that launch now has a return (as opposed to the terminator). That moves it closer to a function where it should feel more like a loop. WDYT about adding the terminator op back or is that too many operations?

[ftynse]

mlir/include/mlir/Dialect/GPU/GPUOps.td

488

Then it's even more confusing than I thought. I'm fine with having a separate terminator for launch.

[rriddle]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

285

nit: You can also use append to add arguments:

emitError().append(..., ..., ...).attachNote(...).append(..., ..., ...);

[mehdi_amini]

Code motion into a gpu::LaunchOp region requires operands of moved instructions to be threaded through operands of the gpu.launch to ensure that the gpu.launch remaines closed from above.

Can we revisit lifting this requirement? This was something that was raised from the beginning of the development of this op.

[herhut]

In D73465#1843770, @mehdi_amini wrote:

Code motion into a gpu::LaunchOp region requires operands of moved instructions to be threaded through operands of the gpu.launch to ensure that the gpu.launch remaines closed from above.

Can we revisit lifting this requirement? This was something that was raised from the beginning of the development of this op.

Should we just lift the requirement or completely forego the concept of passing arguments into the launch? Just lifting the requirement would be less of a breaking change but there is no good use of having the operands.

[herhut]

Split out gpu dialect cleanup parts.

[herhut]

mlir/include/mlir/Dialect/GPU/GPUOps.td

488

I will add it back in a new diff.

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

285

Thanks! I will do this in a new diff.

[merge_guards_bot]

Unit tests: fail. 62155 tests passed, 5 failed and 811 were skipped.

failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[mehdi_amini]

In D73465#1844126, @herhut wrote:

In D73465#1843770, @mehdi_amini wrote:

Code motion into a gpu::LaunchOp region requires operands of moved instructions to be threaded through operands of the gpu.launch to ensure that the gpu.launch remaines closed from above.

Can we revisit lifting this requirement? This was something that was raised from the beginning of the development of this op.

Should we just lift the requirement or completely forego the concept of passing arguments into the launch? Just lifting the requirement would be less of a breaking change but there is no good use of having the operands.

Right: I'm not sure what is the use for having the operands at all?

[herhut]

I went for lifting the requirement that gpu.launch needs to be closed from above. So we no longer have any arguments.

See https://reviews.llvm.org/D73769

mlir/include/mlir/Dialect/GPU/GPUOps.td

let extraClassDeclaration = [{
 
     /// Get the SSA values of the kernel arguments.
     iterator_range<Block::args_iterator> getKernelArguments();
 
     /// Erase the `index`-th kernel argument.  Both the entry block argument and
     /// the operand will be dropped.  The block argument must not have any uses.
     void eraseKernelArgument(unsigned index);
 
+    /// Add the given value as a kernel argument. Returns the corresponding newly
+    /// added BlockArgument.
+    BlockArgument addKernelArgument(Value argument);
+
     static StringRef getBlocksKeyword() { return "blocks"; }
     static StringRef getThreadsKeyword() { return "threads"; }
     static StringRef getArgsKeyword() { return "args"; }
 
     /// The number of launch configuration operands, placed at the leading
     /// positions of the operand list.
     static constexpr unsigned kNumConfigOperands = 6;
 
def GPU_ReturnOp : GPU_Op<"return", [Terminator]>, Arguments<(ins)>,
     Results<(outs)> {
   let summary = "Terminator for GPU launch regions.";
   let description = [{
     A terminator operation for regions that appear in the body of `gpu.launch`
     operation.  These regions are not expected to return any value so the
     terminator takes no operands.
   }];
 
   let parser = [{ return success(); }];
   let printer = [{ p << getOperationName(); }];
 }
 
 def GPU_YieldOp : GPU_Op<"yield", [Terminator]>,
     Arguments<(ins Variadic<AnyType>:$values)> {
   let summary = "GPU yield operation";
   let description = [{
     "gpu.yield" is a special terminator operation for blocks inside regions

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

 
 void LaunchOp::build(Builder *builder, OperationState &result, Value gridSizeX,
                      Value gridSizeY, Value gridSizeZ, Value blockSizeX,
                      Value blockSizeY, Value blockSizeZ, ValueRange operands) {
   // Add grid and block sizes as op operands, followed by the data operands.
   result.addOperands(
       {gridSizeX, gridSizeY, gridSizeZ, blockSizeX, blockSizeY, blockSizeZ});
   result.addOperands(operands);
+  // We want to be able to add operands later, for instance due to code motion.
+  result.setOperandListToResizable();
 
   // Create a kernel body region with kNumConfigRegionAttributes + N arguments,
   // where the first kNumConfigRegionAttributes arguments have `index` type and
   // the rest have the same types as the data operands.
   Region *kernelRegion = result.addRegion();
   Block *body = new Block();
   body->addArguments(
       std::vector<Type>(kNumConfigRegionAttributes, builder->getIndexType()));
for (Block &block : op.body()) {
     if (block.empty())
       continue;
     if (block.back().getNumSuccessors() != 0)
       continue;
     if (!isa<gpu::ReturnOp>(&block.back())) {
       return block.back()
                  .emitError("expected 'gpu.terminator' or a terminator with "
                             "successors")
                  .attachNote(op.getLoc())
              << "in '" << LaunchOp::getOperationName() << "' body region";
     }
   }
 
   return success();
 }
 
 // Pretty-print the kernel grid/block size assignment as

 void LaunchOp::eraseKernelArgument(unsigned index) {
   Block &entryBlock = body().front();
   assert(index < entryBlock.getNumArguments() - kNumConfigRegionAttributes &&
          "kernel argument index overflow");
   entryBlock.eraseArgument(kNumConfigRegionAttributes + index);
   getOperation()->eraseOperand(kNumConfigOperands + index);
 }
 
+BlockArgument LaunchOp::addKernelArgument(Value value) {
+  Block &entryBlock = body().front();
+  Operation *op = getOperation();
+  llvm::SmallVector<Value, 8> operands(op->getOperands());
+  operands.push_back(value);
+  op->setOperands(operands);
+  return entryBlock.addArgument(value.getType());
+}
+
 namespace {
 // Clone any known constants passed as operands to the kernel into its body.
 class PropagateConstantBounds : public OpRewritePattern<LaunchOp> {
   using OpRewritePattern<LaunchOp>::OpRewritePattern;
 
   PatternMatchResult matchAndRewrite(LaunchOp launchOp,
                                      PatternRewriter &rewriter) const override {
     rewriter.startRootUpdate(launchOp);