[PATCH] D73465: Add gpu::LaunchOp::addKernelArgument.

[Stephan Herhut via Phabricator via llvm-commits]

herhut created this revision.
herhut added a reviewer: ftynse.
Herald added subscribers: llvm-commits, liufengdb, lucyrfox, mgester, arpith-jacob, csigg, antiagainst, shauheen, burmako, jpienaar, rriddle, mehdi_amini.
Herald added a reviewer: nicolasvasilache.
Herald added a project: LLVM.

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.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D73465

Files:
  mlir/include/mlir/Dialect/GPU/GPUOps.td
  mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
  mlir/test/Dialect/GPU/invalid.mlir


Index: mlir/test/Dialect/GPU/invalid.mlir
===================================================================
--- mlir/test/Dialect/GPU/invalid.mlir
+++ mlir/test/Dialect/GPU/invalid.mlir
@@ -79,7 +79,7 @@
   // @expected-note at +1 {{in 'gpu.launch' body region}}
   gpu.launch blocks(%bx, %by, %bz) in (%sbx = %sz, %sby = %sz, %sbz = %sz)
              threads(%tx, %ty, %tz) in (%stx = %sz, %sty = %sz, %stz = %sz) {
-    // @expected-error at +1 {{expected 'gpu.terminator' or a terminator with successors}}
+    // @expected-error at +1 {{expected 'gpu.return' or a terminator with successors}}
     return
   }
   return
Index: mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
===================================================================
--- mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -201,6 +201,8 @@
   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
@@ -270,16 +272,16 @@
   }
 
   // Block terminators without successors are expected to exit the kernel region
-  // and must be `gpu.launch`.
+  // and must be `gpu.return`.
   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")
+      return (block.back().emitError()
+              << "expected '" << gpu::ReturnOp::getOperationName() << "' or a "
+              << "terminator with successors")
                  .attachNote(op.getLoc())
              << "in '" << LaunchOp::getOperationName() << "' body region";
     }
@@ -449,6 +451,15 @@
   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> {
Index: mlir/include/mlir/Dialect/GPU/GPUOps.td
===================================================================
--- mlir/include/mlir/Dialect/GPU/GPUOps.td
+++ mlir/include/mlir/Dialect/GPU/GPUOps.td
@@ -454,6 +454,10 @@
     /// 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"; }
@@ -474,11 +478,14 @@
 
 def GPU_ReturnOp : GPU_Op<"return", [Terminator]>, Arguments<(ins)>,
     Results<(outs)> {
-  let summary = "Terminator for GPU launch regions.";
+  let summary = "Terminator for GPU launch regions and GPU functions.";
   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.
+    operation and regions that appear in the body of `gpu.func` functions. The
+    operands to the `gpu.return` are the result values returned by an 
+    incovation of the `gpu.func`.
+    As `gpu.launch` does not return any values, it is expected that a
+    corresponding `gpu.return` also has no operands.
   }];
 
   let parser = [{ return success(); }];


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D73465.240539.patch
Type: text/x-patch
Size: 4148 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200127/54c7926f/attachment.bin>