Make gpu.launch implicitly capture uses of values defined above.

[herhut]

Summary

In the original design, gpu.launch required explicit capture of uses
and passing them as operands to the gpu.launch operation. This was
motivated by infrastructure reestrictions rather than design. This
change lifts the requirement and removes the concept of kernel
arguments from gpu.launch. Instead, the kernel outlining
transformation now does the explicit capturing.

This is a breaking change for users of gpu.launch.

[merge_guards_bot]

Unit tests: pass. 62330 tests passed, 0 failed and 838 were skipped.

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/lib/Dialect/GPU/Transforms/KernelOutlining.cpp

148

ValueRange operands

(side nit: we re-export ArrayRef in the mlir namespace so drop llvm:: for it)

mlir/test/Dialect/GPU/invalid.mlir

4

6 operands

[mehdi_amini]

Nice!

(Nit: typo in the description "reestrictions").

[nicolasvasilache]

nit commit message "reestrictions"

[nicolasvasilache]

very cool, thanks Stephan!

[herhut]

Review comments.

[merge_guards_bot]

Unit tests: pass. 62330 tests passed, 0 failed and 838 were skipped.

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.

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

let extraClassDeclaration = [{
     /// The name of the symbolRef attribute specifying the name of the module
     /// containing the kernel to launch.
     static StringRef getKernelModuleAttrName() { return "kernel_module"; }
   }];
 
   let verifier = [{ return ::verify(*this); }];
 }
 
-def GPU_LaunchOp : GPU_Op<"launch", [IsolatedFromAbove]>,
+def GPU_LaunchOp : GPU_Op<"launch">,
     Arguments<(ins Index:$gridSizeX, Index:$gridSizeY, Index:$gridSizeZ,
-               Index:$blockSizeX, Index:$blockSizeY, Index:$blockSizeZ,
-               Variadic<AnyType>:$operands)>,
+               Index:$blockSizeX, Index:$blockSizeY, Index:$blockSizeZ)>,
     Results<(outs)> {
   let summary = "GPU kernel launch operation";
 
   let description = [{
     Launch a kernel on the specified grid of thread blocks. The body of the
     kernel is defined by the single region that this operation contains. The
-    operation takes at least six operands, with first three operands being grid
-    sizes along x,y,z dimensions, the following three arguments being block
-    sizes along x,y,z dimension, and the remaining operands are arguments of the
-    kernel. When a lower-dimensional kernel is required, unused sizes must be
-    explicitly set to `1`.
+    operation takes six operands, with first three operands being grid sizes
+    along x,y,z dimensions and the following three arguments being block sizes
+    along x,y,z dimension. When a lower-dimensional kernel is required,
+    unused sizes must be explicitly set to `1`.
 
-    The body region has at least _twelve_ arguments, grouped as follows:
+    The body region has _twelve_ arguments, grouped as follows:
 
     -   three arguments that contain block identifiers along x,y,z dimensions;
     -   three arguments that contain thread identifiers along x,y,z dimensions;
-    -   operands of the `gpu.launch` operation as is, including six leading
-        operands for grid and block sizes.
-
-    Operations inside the body region, and any operations in the nested regions,
-    are _not_ allowed to use values defined outside the _body_ region, as if
-    this region was a function. If necessary, values must be passed as kernel
-    arguments into the body region. Nested regions inside the kernel body are
-    allowed to use values defined in their ancestor regions as long as they
-    don't cross the kernel body region boundary.
+    -   operands of the `gpu.launch` operation as is (i.e. the operands for
+        grid and block sizes).
 
     Syntax:
 
     ```
     operation ::= `gpu.launch` `block` `(` ssa-id-list `)` `in` ssa-reassignment
                              `threads` `(` ssa-id-list `)` `in` ssa-reassignment
-                               (`args` ssa-reassignment `:` type-list)?
                                region attr-dict?
     ssa-reassignment ::= `(` ssa-id `=` ssa-use (`,` ssa-id `=` ssa-use)* `)`
     ```
 
     Example:
 
     ```mlir
     gpu.launch blocks(%bx, %by, %bz) in (%sz_bx = %0, %sz_by = %1, %sz_bz = %2)
-               threads(%tx, %ty, %tz) in (%sz_tx = %3, %sz_ty = %4, %sz_tz = %5)
-               args(%arg0 = %6, %arg1 = 7) : f32, memref<?xf32, 1> {
+               threads(%tx, %ty, %tz) in (%sz_tx = %3, %sz_ty = %4, %sz_tz = %5) {
       // Block and thread identifiers, as well as block/grid sizes are
       // immediately usable inside body region.
       "some_op"(%bx, %tx) : (index, index) -> ()
-      %42 = load %arg1[%bx] : memref<?xf32, 1>
+      // Assuming %val1 is defined outside the gpu.launch region.
+      %42 = load %val1[%bx] : memref<?xf32, 1>
     }
 
     // Generic syntax explains how the pretty syntax maps to the IR structure.
     "gpu.launch"(%cst, %cst, %c1,  // Grid sizes.
-                        %cst, %c1, %c1,   // Block sizes.
-                        %arg0, %arg1)     // Actual arguments.
+                 %cst, %c1, %c1)   // Block sizes.
+
         {/*attributes*/}
         // All sizes and identifiers have "index" size.
-        : (index, index, index, index, index, index, f32, memref<?xf32, 1>)
-            -> () {
+        : (index, index, index, index, index, index) -> () {
     // The operation passes block and thread identifiers, followed by grid and
-    // block sizes, followed by actual arguments to the entry block of the
-    // region.
+    // block sizes.
     ^bb0(%bx : index, %by : index, %bz : index,
          %tx : index, %ty : index, %tz : index,
          %num_bx : index, %num_by : index, %num_bz : index,
-         %num_tx : index, %num_ty : index, %num_tz : index,
-         %arg0 : f32, %arg1 : memref<?xf32, 1>):
+         %num_tx : index, %num_ty : index, %num_tz : index)
       "some_op"(%bx, %tx) : (index, index) -> ()
-      %3 = "std.load"(%arg1, %bx) : (memref<?xf32, 1>, index) -> f32
+      %3 = "std.load"(%val1, %bx) : (memref<?xf32, 1>, index) -> f32
     }
     ```
 
     Rationale: using operation/block arguments gives analyses a clear way of
     understanding that a value has additional semantics (e.g., we will need to
     know what value corresponds to threadIdx.x for coalescing). We can recover
     these properties by analyzing the operations producing values, but it is
     easier just to have that information by construction.
   }];
 
   let regions = (region AnyRegion:$body);
 
   let skipDefaultBuilders = 1;
 
   let builders = [
     OpBuilder<"Builder *builder, OperationState &result, Value gridSizeX,"
               "Value gridSizeY, Value gridSizeZ, Value blockSizeX,"
-              "Value blockSizeY, Value blockSizeZ,"
-              "ValueRange operands">
+              "Value blockSizeY, Value blockSizeZ">
   ];
 
-  let hasCanonicalizer = 1;
-
   let extraClassDeclaration = [{
     /// Get the SSA values corresponding to kernel block identifiers.
     KernelDim3 getBlockIds();
     /// Get the SSA values corresponding to kernel thread identifiers.
     KernelDim3 getThreadIds();
     /// Get the SSA values corresponding to kernel grid size.
     KernelDim3 getGridSize();
     /// Get the SSA values corresponding to kernel block size.
     KernelDim3 getBlockSize();
-    /// Get the operand values passed as kernel arguments.
-    operand_range getKernelOperandValues();
-    /// Get the operand types passed as kernel arguments.
-    operand_type_range getKernelOperandTypes();
 
     /// Get the SSA values passed as operands to specify the grid size.
     KernelDim3 getGridSizeOperandValues();
     /// Get the SSA values passed as operands to specify the block size.
     KernelDim3 getBlockSizeOperandValues();
 
-    /// 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);
-
     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;
 
     /// The number of region attributes containing the launch configuration,
     /// placed in the leading positions of the argument list.
     static constexpr unsigned kNumConfigRegionAttributes = 12;

mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp

static LogicalResult createLaunchFromOp(OpTy rootForOp,
   SmallVector<Value, 3> numWorkGroups3D(3, one), workGroupSize3D(3, one);
   for (auto numWorkGroup : enumerate(numWorkGroups)) {
     numWorkGroups3D[numWorkGroup.index()] = numWorkGroup.value();
   }
   for (auto workGroupSize : enumerate(workGroupSizes)) {
     workGroupSize3D[workGroupSize.index()] = workGroupSize.value();
   }
 
-  // Get the values used within the region of the rootForOp but defined above
-  // it.
-  llvm::SetVector<Value> valuesToForwardSet;
-  getUsedValuesDefinedAbove(rootForOp.region(), rootForOp.region(),
-                            valuesToForwardSet);
-  // Also add the values used for the lb, ub, and step of the rootForOp.
-  valuesToForwardSet.insert(rootForOp.getOperands().begin(),
-                            rootForOp.getOperands().end());
-  auto valuesToForward = valuesToForwardSet.takeVector();
   auto launchOp = builder.create<gpu::LaunchOp>(
       rootForOp.getLoc(), numWorkGroups3D[0], numWorkGroups3D[1],
       numWorkGroups3D[2], workGroupSize3D[0], workGroupSize3D[1],
-      workGroupSize3D[2], valuesToForward);
+      workGroupSize3D[2]);
   if (failed(createLaunchBody(builder, rootForOp, launchOp,
                               numWorkGroups.size(), workGroupSizes.size()))) {
     return failure();
   }
 
-  // Replace values that are used within the region of the launchOp but are
-  // defined outside. They all are replaced with kernel arguments.
-  for (auto pair :
-       llvm::zip_first(valuesToForward, launchOp.getKernelArguments())) {
-    Value from = std::get<0>(pair);
-    Value to = std::get<1>(pair);
-    replaceAllUsesInRegionWith(from, to, launchOp.body());
-  }
   return success();
 }
 
 // Replace the rooted at "rootForOp" with a GPU launch operation.  This expects
 // "innermostForOp" to point to the last loop to be transformed to the kernel,
 // and to have (numBlockDims + numThreadDims) perfectly nested loops between
 // "rootForOp" and "innermostForOp".
 // TODO(ravishankarm) : This method can be modified to use the
void LoopToGpuConverter::createLaunch(OpTy rootForOp, OpTy innermostForOp,
   Value gridSizeX = numBlockDims > 0 ? dims[0] : constOne;
   Value gridSizeY = numBlockDims > 1 ? dims[1] : constOne;
   Value gridSizeZ = numBlockDims > 2 ? dims[2] : constOne;
   Value blockSizeX = numThreadDims > 0 ? dims[numBlockDims] : constOne;
   Value blockSizeY = numThreadDims > 1 ? dims[numBlockDims + 1] : constOne;
   Value blockSizeZ = numThreadDims > 2 ? dims[numBlockDims + 2] : constOne;
 
   // Create a launch op and move the body region of the innermost loop to the
-  // launch op.  Pass the values defined outside the outermost loop and used
-  // inside the innermost loop and loop lower bounds as kernel data arguments.
-  // Still assuming perfect nesting so there are no values other than induction
-  // variables that are defined in one loop and used in deeper loops.
-  llvm::SetVector<Value> valuesToForwardSet;
-  getUsedValuesDefinedAbove(innermostForOp.region(), rootForOp.region(),
-                            valuesToForwardSet);
-  auto valuesToForward = valuesToForwardSet.takeVector();
-  auto originallyForwardedValues = valuesToForward.size();
-  valuesToForward.insert(valuesToForward.end(), lbs.begin(), lbs.end());
-  valuesToForward.insert(valuesToForward.end(), steps.begin(), steps.end());
+  // launch op.
   auto launchOp = builder.create<gpu::LaunchOp>(
       rootForOp.getLoc(), gridSizeX, gridSizeY, gridSizeZ, blockSizeX,
-      blockSizeY, blockSizeZ, valuesToForward);
-  valuesToForward.resize(originallyForwardedValues);
+      blockSizeY, blockSizeZ);
 
   // Replace the loop terminator (loops contain only a single block) with the
-  // gpu return and move the operations from the loop body block to the gpu
+  // gpu terminator and move the operations from the loop body block to the gpu
   // launch body block.  Do not move the entire block because of the difference
   // in block arguments.
   Operation &terminator = innermostForOp.getBody()->back();
   Location terminatorLoc = terminator.getLoc();
   terminator.erase();
   builder.setInsertionPointToEnd(innermostForOp.getBody());
   builder.create<gpu::TerminatorOp>(terminatorLoc, llvm::None);
   launchOp.body().front().getOperations().splice(
       launchOp.body().front().begin(),
       innermostForOp.getBody()->getOperations());
 
   // Remap the loop iterators to use block/thread identifiers instead.  Loops
   // may iterate from LB with step S whereas GPU thread/block ids always iterate
   // from 0 to N with step 1.  Therefore, loop induction variables are replaced
   // with (gpu-thread/block-id * S) + LB.
   builder.setInsertionPointToStart(&launchOp.body().front());
-  auto lbArgumentIt = std::next(launchOp.getKernelArguments().begin(),
-                                originallyForwardedValues);
-  auto stepArgumentIt = std::next(lbArgumentIt, lbs.size());
+  auto lbArgumentIt = lbs.begin();
+  auto stepArgumentIt = steps.begin();
   for (auto en : llvm::enumerate(ivs)) {
     Value id =
         en.index() < numBlockDims
             ? getDim3Value(launchOp.getBlockIds(), en.index())
             : getDim3Value(launchOp.getThreadIds(), en.index() - numBlockDims);
     Value step = steps[en.index()];
     if (!isConstantOne(step))
       id = builder.create<MulIOp>(rootForOp.getLoc(), step, id);
 
     Value ivReplacement =
         builder.create<AddIOp>(rootForOp.getLoc(), *lbArgumentIt, id);
     en.value().replaceAllUsesWith(ivReplacement);
-    replaceAllUsesInRegionWith(steps[en.index()], *stepArgumentIt,
-                               launchOp.body());
     std::advance(lbArgumentIt, 1);
     std::advance(stepArgumentIt, 1);
   }
 
-  // Remap the values defined outside the body to use kernel arguments instead.
-  // The list of kernel arguments also contains the lower bounds for loops at
-  // trailing positions, make sure we don't touch those.
-  for (auto pair :
-       llvm::zip_first(valuesToForward, launchOp.getKernelArguments())) {
-    Value from = std::get<0>(pair);
-    Value to = std::get<1>(pair);
-    replaceAllUsesInRegionWith(from, to, launchOp.body());
-  }
-
   // We are done and can erase the original outermost loop.
   rootForOp.erase();
 }
 
 // Generic loop to GPU kernel conversion function.
 template <typename OpTy>
 static LogicalResult convertLoopNestToGPULaunch(OpTy forOp,
                                                 unsigned numBlockDims,

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

 }
 
 //===----------------------------------------------------------------------===//
 // LaunchOp
 //===----------------------------------------------------------------------===//
 
 void LaunchOp::build(Builder *builder, OperationState &result, Value gridSizeX,
                      Value gridSizeY, Value gridSizeZ, Value blockSizeX,
-                     Value blockSizeY, Value blockSizeZ, ValueRange operands) {
+                     Value blockSizeY, Value blockSizeZ) {
   // Add grid and block sizes as op operands, followed by the data operands.
   result.addOperands(
       {gridSizeX, gridSizeY, gridSizeZ, blockSizeX, blockSizeY, blockSizeZ});
-  result.addOperands(operands);
 
   // 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()));
-  body->addArguments(llvm::to_vector<4>(operands.getTypes()));
   kernelRegion->push_back(body);
 }
 
 KernelDim3 LaunchOp::getBlockIds() {
   assert(!body().getBlocks().empty() && "FuncOp body must not be empty.");
   auto args = body().getBlocks().front().getArguments();
   return KernelDim3{args[0], args[1], args[2]};
 }

 }
 
 KernelDim3 LaunchOp::getBlockSize() {
   assert(!body().getBlocks().empty() && "FuncOp body must not be empty.");
   auto args = body().getBlocks().front().getArguments();
   return KernelDim3{args[9], args[10], args[11]};
 }
 
-LaunchOp::operand_range LaunchOp::getKernelOperandValues() {
-  return llvm::drop_begin(getOperands(), kNumConfigOperands);
-}
-
-LaunchOp::operand_type_range LaunchOp::getKernelOperandTypes() {
-  return llvm::drop_begin(getOperandTypes(), kNumConfigOperands);
-}
-
 KernelDim3 LaunchOp::getGridSizeOperandValues() {
   return KernelDim3{getOperand(0), getOperand(1), getOperand(2)};
 }
 
 KernelDim3 LaunchOp::getBlockSizeOperandValues() {
   return KernelDim3{getOperand(3), getOperand(4), getOperand(5)};
 }
 
-iterator_range<Block::args_iterator> LaunchOp::getKernelArguments() {
-  auto args = body().getBlocks().front().getArguments();
-  return llvm::drop_begin(args, LaunchOp::kNumConfigRegionAttributes);
-}
-
 static LogicalResult verify(LaunchOp op) {
   // Kernel launch takes kNumConfigOperands leading operands for grid/block
   // sizes and transforms them into kNumConfigRegionAttributes region arguments
   // for block/thread identifiers and grid/block sizes.
   if (!op.body().empty()) {
     Block &entryBlock = op.body().front();
     if (entryBlock.getNumArguments() !=
         LaunchOp::kNumConfigOperands + op.getNumOperands())
static void printLaunchOp(OpAsmPrinter &p, LaunchOp op) {
   // Print the launch configuration.
   p << LaunchOp::getOperationName() << ' ' << op.getBlocksKeyword();
   printSizeAssignment(p, op.getGridSize(), operands.take_front(3),
                       op.getBlockIds());
   p << ' ' << op.getThreadsKeyword();
   printSizeAssignment(p, op.getBlockSize(), operands.slice(3, 3),
                       op.getThreadIds());
 
-  // From now on, the first kNumConfigOperands operands corresponding to grid
-  // and block sizes are irrelevant, so we can drop them.
-  operands = operands.drop_front(LaunchOp::kNumConfigOperands);
-
-  // Print the data argument remapping.
-  if (!op.body().empty() && !operands.empty()) {
-    p << ' ' << op.getArgsKeyword() << '(';
-    Block *entryBlock = &op.body().front();
-    interleaveComma(llvm::seq<int>(0, operands.size()), p, [&](int i) {
-      p << entryBlock->getArgument(LaunchOp::kNumConfigRegionAttributes + i)
-        << " = " << operands[i];
-    });
-    p << ") ";
-  }
-
-  // Print the types of data arguments.
-  if (!operands.empty())
-    p << ": " << operands.getTypes();
-
   p.printRegion(op.body(), /*printEntryBlockArgs=*/false);
   p.printOptionalAttrDict(op.getAttrs());
 }
 
 // Parse the size assignment blocks for blocks and threads.  These have the form
 //   (%region_arg, %region_arg, %region_arg) in
 //   (%region_arg = %operand, %region_arg = %operand, %region_arg = %operand)
 // where %region_arg are percent-identifiers for the region arguments to be
parseSizeAssignment(OpAsmParser &parser,
   }
 
   return parser.parseRParen();
 }
 
 // Parses a Launch operation.
 // operation ::= `gpu.launch` `blocks` `(` ssa-id-list `)` `in` ssa-reassignment
 //                           `threads` `(` ssa-id-list `)` `in` ssa-reassignment
-//                             (`args` ssa-reassignment `:` type-list)?
-//                             region attr-dict?
+//                            region attr-dict?
 // ssa-reassignment ::= `(` ssa-id `=` ssa-use (`,` ssa-id `=` ssa-use)* `)`
 static ParseResult parseLaunchOp(OpAsmParser &parser, OperationState &result) {
   // Sizes of the grid and block.
   SmallVector<OpAsmParser::OperandType, LaunchOp::kNumConfigOperands> sizes(
       LaunchOp::kNumConfigOperands);
   MutableArrayRef<OpAsmParser::OperandType> sizesRef(sizes);
 
   // Actual (data) operands passed to the kernel.
if (parser.parseKeyword(LaunchOp::getBlocksKeyword().data()) ||
       parser.parseKeyword(LaunchOp::getThreadsKeyword().data()) ||
       parseSizeAssignment(parser, sizesRef.drop_front(3),
                           regionArgsRef.slice(9, 3),
                           regionArgsRef.slice(3, 3)) ||
       parser.resolveOperands(sizes, parser.getBuilder().getIndexType(),
                              result.operands))
     return failure();
 
-  // If kernel argument renaming segment is present, parse it.  When present,
-  // the segment should have at least one element.  If this segment is present,
-  // so is the trailing type list.  Parse it as well and use the parsed types
-  // to resolve the operands passed to the kernel arguments.
-  SmallVector<Type, 4> dataTypes;
-  if (!parser.parseOptionalKeyword(LaunchOp::getArgsKeyword())) {
-    llvm::SMLoc argsLoc = parser.getCurrentLocation();
-
-    regionArgs.push_back({});
-    dataOperands.push_back({});
-    if (parser.parseLParen() || parser.parseRegionArgument(regionArgs.back()) ||
-        parser.parseEqual() || parser.parseOperand(dataOperands.back()))
-      return failure();
-
-    while (!parser.parseOptionalComma()) {
-      regionArgs.push_back({});
-      dataOperands.push_back({});
-      if (parser.parseRegionArgument(regionArgs.back()) ||
-          parser.parseEqual() || parser.parseOperand(dataOperands.back()))
-        return failure();
-    }
-
-    if (parser.parseRParen() || parser.parseColonTypeList(dataTypes) ||
-        parser.resolveOperands(dataOperands, dataTypes, argsLoc,
-                               result.operands))
-      return failure();
-  }
-
-  // Introduce the body region and parse it.  The region has
-  // kNumConfigRegionAttributes leading arguments that correspond to
+  // Introduce the body region and parse it. The region has
+  // kNumConfigRegionAttributes arguments that correspond to
   // block/thread identifiers and grid/block sizes, all of the `index` type.
-  // Follow the actual kernel arguments.
   Type index = parser.getBuilder().getIndexType();
-  dataTypes.insert(dataTypes.begin(), LaunchOp::kNumConfigRegionAttributes,
-                   index);
+  SmallVector<Type, LaunchOp::kNumConfigRegionAttributes> dataTypes(
+      LaunchOp::kNumConfigRegionAttributes, index);
   Region *body = result.addRegion();
   return failure(parser.parseRegion(*body, regionArgs, dataTypes) ||
                  parser.parseOptionalAttrDict(result.attributes));
 }
 
-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);
-}
-
-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);
-    PatternRewriter::InsertionGuard guard(rewriter);
-    rewriter.setInsertionPointToStart(&launchOp.body().front());
-
-    // Traverse operands passed to kernel and check if some of them are known
-    // constants.  If so, clone the constant operation inside the kernel region
-    // and use it instead of passing the value from the parent region.  Perform
-    // the traversal in the inverse order to simplify index arithmetics when
-    // dropping arguments.
-    auto operands = launchOp.getKernelOperandValues();
-    auto kernelArgs = launchOp.getKernelArguments();
-    bool found = false;
-    for (unsigned i = operands.size(); i > 0; --i) {
-      unsigned index = i - 1;
-      Value operand = operands[index];
-      if (!isa_and_nonnull<ConstantOp>(operand.getDefiningOp()))
-        continue;
-
-      found = true;
-      Value internalConstant =
-          rewriter.clone(*operand.getDefiningOp())->getResult(0);
-      Value kernelArg = *std::next(kernelArgs.begin(), index);
-      kernelArg.replaceAllUsesWith(internalConstant);
-      launchOp.eraseKernelArgument(index);
-    }
-
-    if (!found) {
-      rewriter.cancelRootUpdate(launchOp);
-      return matchFailure();
-    }
-
-    rewriter.finalizeRootUpdate(launchOp);
-    return matchSuccess();
-  }
-};
-} // end namespace
-
-void LaunchOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
-                                           MLIRContext *context) {
-  results.insert<PropagateConstantBounds>(context);
-}
-
 //===----------------------------------------------------------------------===//
 // LaunchFuncOp
 //===----------------------------------------------------------------------===//
 
 void LaunchFuncOp::build(Builder *builder, OperationState &result,
                          GPUFuncOp kernelFunc, Value gridSizeX, Value gridSizeY,
                          Value gridSizeZ, Value blockSizeX, Value blockSizeY,
                          Value blockSizeZ, ValueRange kernelOperands) {

mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp

 
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/RegionUtils.h"
 
 using namespace mlir;
 
 template <typename OpTy>
 static void createForAllDimensions(OpBuilder &builder, Location loc,
                                    SmallVectorImpl<Value> &values) {
   for (StringRef dim : {"x", "y", "z"}) {
     Value v = builder.create<OpTy>(loc, builder.getIndexType(),
auto newLaunch = LaunchBuilder.create<gpu::LaunchFuncOp>(
       launch.getLoc(), kernelFunc, launch.getGridSizeOperandValues(),
       launch.getBlockSizeOperandValues(), newLaunchArgs);
   launch.erase();
   return newLaunch;
 }
 
 // Outline the `gpu.launch` operation body into a kernel function. Replace
 // `gpu.terminator` operations by `gpu.return` in the generated function.
-static gpu::GPUFuncOp outlineKernelFunc(gpu::LaunchOp launchOp) {
+static gpu::GPUFuncOp outlineKernelFunc(gpu::LaunchOp launchOp,
+                                        llvm::SetVector<Value> &operands) {
   Location loc = launchOp.getLoc();
   // Create a builder with no insertion point, insertion will happen separately
   // due to symbol table manipulation.
   OpBuilder builder(launchOp.getContext());
 
-  SmallVector<Type, 4> kernelOperandTypes(launchOp.getKernelOperandTypes());
+  // Identify uses from values defined outside of the scope of the launch
+  // operation.
+  getUsedValuesDefinedAbove(launchOp.body(), operands);
+
+  SmallVector<Type, 4> kernelOperandTypes;
+  kernelOperandTypes.reserve(operands.size());
+  for (Value operand : operands) {
+    kernelOperandTypes.push_back(operand.getType());
+  }
   FunctionType type =
       FunctionType::get(kernelOperandTypes, {}, launchOp.getContext());
   std::string kernelFuncName =
       Twine(launchOp.getParentOfType<FuncOp>().getName(), "_kernel").str();
   auto outlinedFunc = builder.create<gpu::GPUFuncOp>(loc, kernelFuncName, type);
   outlinedFunc.setAttr(gpu::GPUDialect::getKernelFuncAttrName(),
                        builder.getUnitAttr());
   outlinedFunc.body().takeBody(launchOp.body());
   injectGpuIndexOperations(loc, outlinedFunc.body());
+  Block &entryBlock = outlinedFunc.body().front();
+  for (Value operand : operands) {
+    BlockArgument newArg = entryBlock.addArgument(operand.getType());
+    replaceAllUsesInRegionWith(operand, newArg, outlinedFunc.body());
+  }
   outlinedFunc.walk([](gpu::TerminatorOp op) {
     OpBuilder replacer(op);
     replacer.create<gpu::ReturnOp>(op.getLoc());
     op.erase();
   });
 
   return outlinedFunc;
 }
 
 // Replace `gpu.launch` operations with an `gpu.launch_func` operation launching
 // `kernelFunc`. The kernel func contains the body of the `gpu.launch` with
 // constant region arguments inlined.
 static void convertToLaunchFuncOp(gpu::LaunchOp &launchOp,
-                                  gpu::GPUFuncOp kernelFunc) {
+                                  gpu::GPUFuncOp kernelFunc,
+                                  ValueRange operands) {
   OpBuilder builder(launchOp);
   auto launchFuncOp = builder.create<gpu::LaunchFuncOp>(
       launchOp.getLoc(), kernelFunc, launchOp.getGridSizeOperandValues(),
-      launchOp.getBlockSizeOperandValues(), launchOp.getKernelOperandValues());
+      launchOp.getBlockSizeOperandValues(), operands);
   inlineBeneficiaryOps(kernelFunc, launchFuncOp);
   launchOp.erase();
 }
 
 namespace {
 
 /// Pass that moves the kernel of each LaunchOp into its separate nested module.
 ///
 /// This pass moves the kernel code of each LaunchOp into a function created
 /// inside a nested module. It also creates an external function of the same
 /// name in the parent module.
 ///
 /// The gpu.modules are intended to be compiled to a cubin blob independently in
 /// a separate pass. The external functions can then be annotated with the
 /// symbol of the cubin accessor function.
 class GpuKernelOutliningPass : public ModulePass<GpuKernelOutliningPass> {
 public:
   void runOnModule() override {
     SymbolTable symbolTable(getModule());
     bool modified = false;
     for (auto func : getModule().getOps<FuncOp>()) {
       // Insert just after the function.
       Block::iterator insertPt(func.getOperation()->getNextNode());
       func.walk([&](gpu::LaunchOp op) {
-        gpu::GPUFuncOp outlinedFunc = outlineKernelFunc(op);
+        llvm::SetVector<Value> operands;
+        gpu::GPUFuncOp outlinedFunc = outlineKernelFunc(op, operands);
 
         // Create nested module and insert outlinedFunc. The module will
         // originally get the same name as the function, but may be renamed on
         // insertion into the parent module.
         auto kernelModule = createKernelModule(outlinedFunc, symbolTable);
         symbolTable.insert(kernelModule, insertPt);
 
         // Potentially changes signature, pulling in constants.
-        convertToLaunchFuncOp(op, outlinedFunc);
+        convertToLaunchFuncOp(op, outlinedFunc, operands.getArrayRef());
         modified = true;
       });
     }
 
     // If any new module was inserted in this module, annotate this module as
     // a container module.
     if (modified)
       getModule().setAttr(gpu::GPUDialect::getContainerModuleAttrName(),

mlir/test/Conversion/LoopsToGPU/imperfect_2D.mlir

 // RUN: mlir-opt -convert-loop-op-to-gpu -gpu-num-workgroups=2,2 -gpu-workgroup-size=32,4 %s | FileCheck %s
 
 module {
   // arg2 = arg0 * transpose(arg1) ; with intermediate buffer and tile size passed as argument
   // CHECK: func {{@.*}}([[ARG0:%.*]]: memref<?x?xf32>, [[ARG1:%.*]]: memref<?x?xf32>, [[ARG2:%.*]]: memref<?x?xf32>, [[ARG3:%.*]]: index, [[ARG4:%.*]]: index)
   func @foo(%arg0: memref<?x?xf32>, %arg1 : memref<?x?xf32>, %arg2 : memref<?x?xf32>, %arg3 : index, %arg4 : index) {
     %0 = dim %arg0, 0 : memref<?x?xf32>
     %1 = dim %arg0, 1 : memref<?x?xf32>
     %c0 = constant 0 : index
     %c1 = constant 1 : index
-    // CHECK: gpu.launch blocks([[ARG5:%.*]], [[ARG6:%.*]], [[ARG7:%.*]]) in ([[ARG11:%.*]] = {{%.*}}, [[ARG12:%.*]] = {{%.*}}, [[ARG13:%.*]] = {{%.*}}) threads([[ARG8:%.*]], [[ARG9:%.*]], [[ARG10:%.*]]) in ([[ARG14:%.*]] = {{%.*}}, [[ARG15:%.*]] = {{%.*}}, [[ARG16:%.*]] = {{%.*}}) args([[ARG17:%.*]] = [[ARG3]], [[ARG18:%.*]] = [[ARG4]], [[ARG19:%.*]] = [[ARG1]], [[ARG20:%.*]] = {{%.*}}, {{%.*}} = {{%.*}}, [[ARG22:%.*]] = [[ARG0]], [[ARG23:%.*]] = [[ARG2]]
-    // CHECK: [[TEMP1:%.*]] = muli [[ARG17]], [[ARG6]] : index
+    // CHECK: gpu.launch blocks([[ARG5:%.*]], [[ARG6:%.*]], [[ARG7:%.*]]) in ([[ARG11:%.*]] = {{%.*}}, [[ARG12:%.*]] = {{%.*}}, [[ARG13:%.*]] = {{%.*}}) threads([[ARG8:%.*]], [[ARG9:%.*]], [[ARG10:%.*]]) in ([[ARG14:%.*]] = {{%.*}}, [[ARG15:%.*]] = {{%.*}}, [[ARG16:%.*]] = {{%.*}})
+    // CHECK: [[TEMP1:%.*]] = muli [[ARG3]], [[ARG6]] : index
     // CHECK: [[BLOCKLOOPYLB:%.*]] = addi {{%.*}}, [[TEMP1]] : index
-    // CHECK: [[BLOCKLOOPYSTEP:%.*]] = muli [[ARG17]], [[ARG12]] : index
+    // CHECK: [[BLOCKLOOPYSTEP:%.*]] = muli [[ARG3]], [[ARG12]] : index
     // CHECK: loop.for [[BLOCKLOOPYIV:%.*]] = [[BLOCKLOOPYLB]] to {{%.*}} step [[BLOCKLOOPYSTEP]]
     loop.for %iv1 = %c0 to %0 step %arg3 {
 
-      // CHECK: [[TEMP2:%.*]] = muli [[ARG18]], [[ARG5]] : index
+      // CHECK: [[TEMP2:%.*]] = muli [[ARG4]], [[ARG5]] : index
       // CHECK: [[BLOCKLOOPXLB:%.*]] = addi  {{%.*}}, [[TEMP2]] : index
-      // CHECK: [[BLOCKLOOPXSTEP:%.*]] = muli [[ARG18]], [[ARG11]] : index
+      // CHECK: [[BLOCKLOOPXSTEP:%.*]] = muli [[ARG4]], [[ARG11]] : index
       // CHECK: loop.for [[BLOCKLOOPXIV:%.*]] = [[BLOCKLOOPXLB]] to {{%.*}} step [[BLOCKLOOPXSTEP]]
 
       loop.for %iv2 = %c0 to %1 step %arg4 {
 
         // TODO: This is effectively shared memory. Lower it to a
         // shared memory.
         %2 = alloc(%arg3, %arg4) : memref<?x?xf32>
 
         // Load transpose tile
-        // CHECK: [[TEMP3:%.*]] = muli [[ARG20]], [[ARG9:%.*]] : index
+        // CHECK: [[TEMP3:%.*]] = muli [[ARG20:%.*]], [[ARG9:%.*]] : index
         // CHECK: [[THREADLOOP1YLB:%.*]] = addi {{%.*}}, [[TEMP3]] : index
         // CHECK: [[THREADLOOP1YSTEP:%.*]] = muli [[ARG20]], [[ARG15]] : index
         // CHECK: loop.for [[THREADLOOP1YIV:%.*]] = [[THREADLOOP1YLB]] to {{%.*}} step [[THREADLOOP1YSTEP]]
         loop.for %iv3 = %c0 to %arg3 step %c1 {
           // CHECK: [[TEMP4:%.*]] = muli [[ARG20]], [[ARG8]] : index
           // CHECK: [[THREADLOOP1XLB:%.*]] = addi {{%.*}}, [[TEMP4]] : index
           // CHECK: [[THREADLOOP1XSTEP:%.*]] = muli [[ARG20]], [[ARG14]] : index
           // CHECK: loop.for [[THREADLOOP1XIV:%.*]] = [[THREADLOOP1XLB]] to {{%.*}} step [[THREADLOOP1XSTEP]]
           loop.for %iv4 = %c1 to %arg4 step %c1 {
             // CHECK: [[INDEX2:%.*]] = addi [[BLOCKLOOPYIV]], [[THREADLOOP1YIV]] : index
             %10 = addi %iv1, %iv3 : index
             // CHECK: [[INDEX1:%.*]] = addi [[BLOCKLOOPXIV]], [[THREADLOOP1XIV]] : index
             %11 = addi %iv2, %iv4 : index
-            // CHECK: [[VAL1:%.*]] = load [[ARG19]]{{\[}}[[INDEX1]], [[INDEX2]]{{\]}} : memref<?x?xf32>
+            // CHECK: [[VAL1:%.*]] = load [[ARG1]]{{\[}}[[INDEX1]], [[INDEX2]]{{\]}} : memref<?x?xf32>
             %12 = load %arg1[%11, %10] : memref<?x?xf32>
             // CHECK: store [[VAL1]], [[SCRATCHSPACE:%.*]]{{\[}}[[THREADLOOP1XIV]], [[THREADLOOP1YIV]]{{\]}} : memref<?x?xf32>
             store %12, %2[%iv4, %iv3] : memref<?x?xf32>
           }
         }
 
         // TODO: There needs to be a sync here for correctness, but
         // testing only loop partitioning for now.
loop.for %iv1 = %c0 to %0 step %arg3 {
           // CHECK: loop.for [[THREADLOOP2XIV:%.*]] = [[THREADLOOP2XLB]] to {{%.*}} step [[THREADLOOP2XSTEP]]
           loop.for %iv4 = %c1 to %arg4 step %c1 {
             // CHECK: [[INDEX3:%.*]] = addi [[BLOCKLOOPYIV]], [[THREADLOOP2YIV]] : index
             %13 = addi %iv1, %iv3 : index
             // CHECK: [[INDEX4:%.*]] = addi [[BLOCKLOOPXIV]], [[THREADLOOP2XIV]] : index
             %14 = addi %iv2, %iv4 : index
             // CHECK: {{%.*}} = load [[SCRATCHSPACE]]{{\[}}[[THREADLOOP2XIV]], [[THREADLOOP2YIV]]{{\]}} : memref<?x?xf32>
             %15 = load %2[%iv4, %iv3] : memref<?x?xf32>
-            // CHECK: {{%.*}} = load [[ARG22]]{{\[}}[[INDEX3]], [[INDEX4]]{{\]}}
+            // CHECK: {{%.*}} = load [[ARG0]]{{\[}}[[INDEX3]], [[INDEX4]]{{\]}}
             %16 = load %arg0[%13, %14] : memref<?x?xf32>
             %17 = mulf %15, %16 : f32
-            // CHECK: store {{%.*}}, [[ARG23]]{{\[}}[[INDEX3]], [[INDEX4]]{{\]}}
+            // CHECK: store {{%.*}}, [[ARG2]]{{\[}}[[INDEX3]], [[INDEX4]]{{\]}}
             store %17, %arg2[%13, %14] : memref<?x?xf32>
           }
         }
 
         dealloc %2 : memref<?x?xf32>
       }
     }
     return
   }
-}
+}
 

mlir/test/Conversion/LoopsToGPU/linalg_to_gpu.mlir

 // RUN: mlir-opt -convert-loops-to-gpu %s | FileCheck %s
 
 // CHECK-LABEL: @foo
 func @foo(%arg0: memref<?xf32>, %arg1 : index) {
   %c0 = constant 0 : index
   %c42 = constant 42 : index
   %c3 = constant 3 : index
   // CHECK:      subi %{{.*}}, %{{.*}} : index
   // CHECK-NEXT: %[[range_i:.*]] = divi_signed {{.*}}, %{{.*}} : index
   loop.for %i0 = %c0 to %c42 step %c3 {
     // CHECK:      subi %{{.*}}, %{{.*}} : index
     // CHECK-NEXT: %[[range_j:.*]] = divi_signed {{.*}}, %{{.*}} : index
     loop.for %i1 = %c3 to %c42 step %arg1 {
       // CHECK:      gpu.launch
       // CHECK-SAME: blocks
       // CHECK-SAME: threads
-      // CHECK-SAME: args
 
       // Replacements of loop induction variables.  Take a product with the
       // step and add the lower bound.
       // CHECK: %[[prod_i:.*]] = muli %{{.*}}, %{{.*}} : index
       // CHECK: addi %{{.*}}, %[[prod_i]] : index
       // CHECK: %[[prod_j:.*]] = muli %{{.*}}, %{{.*}} : index
       // CHECK: addi %{{.*}}, %[[prod_j]] : index
 
       // CHECK: gpu.terminator
     }
   }
   return
 }

mlir/test/Conversion/LoopsToGPU/step_one.mlir

affine.for %i = 0 to 42 {
     // CHECK-22-NEXT: %{{.*}} = constant 0 : index
     // CHECK-22-NEXT: %{{.*}} = constant 10 : index
     // CHECK-22-NEXT: %{{.*}} = subi %{{.*}}, %{{.*}} : index
     // CHECK-22-NEXT: %{{.*}} = constant 1 : index
     affine.for %j = 0 to 10 {
     // CHECK-11: gpu.launch
     // CHECK-11-SAME: blocks
     // CHECK-11-SAME: threads
-    // CHECK-11-SAME: args
 
       // Remapping of the loop induction variables.
       // CHECK-11:        %[[i:.*]] = addi %{{.*}}, %{{.*}} : index
       // CHECK-11-NEXT:   %[[j:.*]] = addi %{{.*}}, %{{.*}} : index
 
       // This loop is not converted if mapping to 1, 1 dimensions.
       // CHECK-11-NEXT: affine.for %[[ii:.*]] = 2 to 16
       //
// CHECK-11-SAME: threads
         // CHECK-22-NEXT: %{{.*}} = constant 5 : index
         // CHECK-22-NEXT: %{{.*}} = constant 17 : index
         // CHECK-22-NEXT: %{{.*}} = subi %{{.*}}, %{{.*}} : index
         // CHECK-22-NEXT: %{{.*}} = constant 1 : index
         affine.for %jj = 5 to 17 {
         // CHECK-22: gpu.launch
         // CHECK-22-SAME: blocks
         // CHECK-22-SAME: threads
-        // CHECK-22-SAME: args
 
           // Remapping of the loop induction variables in the last mapped loop.
           // CHECK-22:        %[[i:.*]] = addi %{{.*}}, %{{.*}} : index
           // CHECK-22-NEXT:   %[[j:.*]] = addi %{{.*}}, %{{.*}} : index
           // CHECK-22-NEXT:   %[[ii:.*]] = addi %{{.*}}, %{{.*}} : index
           // CHECK-22-NEXT:   %[[jj:.*]] = addi %{{.*}}, %{{.*}} : index
 
           // Using remapped values instead of loop iterators.

mlir/test/Dialect/GPU/canonicalize.mlir

-// RUN: mlir-opt -pass-pipeline='func(canonicalize)' %s | FileCheck %s
-
-// CHECK-LABEL: @propagate_constant
-// CHECK-SAME:  %[[arg1:.*]]: memref
-func @propagate_constant(%arg1: memref<?xf32>) {
-  // The outer constant must be preserved because it still has uses.
-  // CHECK: %[[outer_cst:.*]] = constant 1
-  %c1 = constant 1 : index
-
-  // The constant must be dropped from the args list, but the memref should
-  // remain.
-  // CHECK: gpu.launch
-  // CHECK-SAME: args(%[[inner_arg:.*]] = %[[arg1]]) : memref
-  gpu.launch blocks(%bx, %by, %bz) in (%sbx = %c1, %sby = %c1, %sbz = %c1)
-             threads(%tx, %ty, %tz) in (%stx = %c1, %sty = %c1, %stz = %c1)
-             args(%x = %c1, %y = %arg1) : index, memref<?xf32> {
-    // The constant is propagated into the kernel body and used.
-    // CHECK: %[[inner_cst:.*]] = constant 1
-    // CHECK: "foo"(%[[inner_cst]])
-    "foo"(%x) : (index) -> ()
-
-    // CHECK: "bar"(%[[inner_arg]])
-    "bar"(%y) : (memref<?xf32>) -> ()
-    gpu.terminator
-  }
-  return
-}
-

mlir/test/Dialect/GPU/invalid.mlir

 // RUN: mlir-opt -split-input-file -verify-diagnostics %s
 
 func @not_enough_sizes(%sz : index) {
-  // expected-error@+1 {{expected 6 or more operands}}
+  // expected-error@+1 {{expected 6 operands, but found 5}}
   "gpu.launch"(%sz, %sz, %sz, %sz, %sz) ({
     gpu.return
   }) : (index, index, index, index, index) -> ()
   return
 }
 
 // -----
 
 func @no_region_attrs(%sz : index) {
   // expected-error@+1 {{unexpected number of region arguments}}
  "gpu.launch"(%sz, %sz, %sz, %sz, %sz, %sz) ({
   ^bb1(%bx: index, %by: index, %bz: index,
        %tx: index, %ty: index, %tz: index):
     gpu.return
   }) : (index, index, index, index, index, index) -> ()
   return
 }
 
 // -----
 
-func @isolation_arg(%sz : index) {
- // expected-note@+1 {{required by region isolation constraints}}
- "gpu.launch"(%sz, %sz, %sz, %sz, %sz, %sz) ({
-  ^bb1(%bx: index, %by: index, %bz: index,
-       %tx: index, %ty: index, %tz: index,
-       %szbx: index, %szby: index, %szbz: index,
-       %sztx: index, %szty: index, %sztz: index):
-    // expected-error@+1 {{using value defined outside the region}}
-    "use"(%sz) : (index) -> ()
-    gpu.return
-  }) : (index, index, index, index, index, index) -> ()
-  return
-}
-
-// -----
-
-func @isolation_op(%sz : index) {
- %val = "produce"() : () -> (index)
- // expected-note@+1 {{required by region isolation constraints}}
- "gpu.launch"(%sz, %sz, %sz, %sz, %sz, %sz) ({
-  ^bb1(%bx: index, %by: index, %bz: index,
-       %tx: index, %ty: index, %tz: index,
-       %szbx: index, %szby: index, %szbz: index,
-       %sztx: index, %szty: index, %sztz: index):
-    // expected-error@+1 {{using value defined outside the region}}
-    "use"(%val) : (index) -> ()
-    gpu.return
-  }) : (index, index, index, index, index, index) -> ()
-  return
-}
-
-// -----
-
-func @nested_isolation(%sz : index) {
-  // expected-note@+1 {{required by region isolation constraints}}
-  "gpu.launch"(%sz, %sz, %sz, %sz, %sz, %sz) ({
-  ^bb1(%bx: index, %by: index, %bz: index,
-       %tx: index, %ty: index, %tz: index,
-       %szbx: index, %szby: index, %szbz: index,
-       %sztx: index, %szty: index, %sztz: index):
-    "region"() ({
-      "region"() ({
-        // expected-error@+1 {{using value defined outside the region}}
-        "use"(%sz) : (index) -> ()
-      }) : () -> ()
-    }) : () -> ()
-    gpu.return
-  }) : (index, index, index, index, index, index) -> ()
-  return
-}
-
-// -----
-
 func @launch_requires_gpu_return(%sz : index) {
   // @expected-note@+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@+1 {{expected 'gpu.terminator' or a terminator with successors}}
     return
   }
   return

 module {
   gpu.module @gpu_funcs {
     // expected-error @+1 {{'gpu.func' op expected void return type for kernel function}}
     gpu.func @kernel() -> index kernel {
       %0 = constant 0 : index
       gpu.return
     }
   }
-}
+}
 

mlir/test/Dialect/GPU/ops.mlir

gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %sz, %grid_y = %sz, %grid_z = %sz)
       // CHECK: gpu.terminator
       gpu.terminator
     }
     return
   }
 
   // CHECK-LABEL:func @args(%{{.*}}: index, %{{.*}}: index, %{{.*}}: f32, %{{.*}}: memref<?xf32, 1>) {
   func @args(%blk : index, %thrd : index, %float : f32, %data : memref<?xf32,1>) {
-    // CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) args(%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) : f32, memref<?xf32, 1>
-    gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %blk, %grid_y = %blk, %grid_z = %blk)
-               threads(%tx, %ty, %tz) in (%block_x = %thrd, %block_y = %thrd, %block_z = %thrd)
-               args(%kernel_arg0 = %float, %kernel_arg1 = %data) : f32, memref<?xf32, 1> {
-      // CHECK: gpu.terminator
-      gpu.terminator
-    }
-    return
-  }
-
-  // It is possible to use values passed into the region as arguments.
-  // CHECK-LABEL: func @passing_values
-  func @passing_values(%blk : index, %thrd : index, %float : f32, %data : memref<?xf32,1>) {
-    // CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) args(%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) : f32, memref<?xf32, 1>
+    // CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}})
     gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %blk, %grid_y = %blk, %grid_z = %blk)
-               threads(%tx, %ty, %tz) in (%block_x = %thrd, %block_y = %thrd, %block_z = %thrd)
-               args(%kernel_arg0 = %float, %kernel_arg1 = %data) : f32, memref<?xf32, 1> {
-      // CHECK: "use"(%{{.*}})
-      "use"(%kernel_arg0): (f32) -> ()
-      // CHECK: gpu.terminator
-      gpu.terminator
-    }
-    return
-  }
-
-  // It is possible to use values defined in nested regions as long as they don't
-  // cross kernel launch region boundaries.
-  // CHECK-LABEL: func @nested_isolation
-  func @nested_isolation(%sz : index) {
-    gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %sz, %grid_y = %sz, %grid_z = %sz)
-               threads(%tx, %ty, %tz) in (%block_x = %sz, %block_y = %sz, %block_z = %sz) {
-      "region"() ({
-        // CHECK: %{{.*}} = "produce"()
-        %val = "produce"() : () -> (index)
-        "region"() ({
-          // CHECK: "use"(%{{.*}})
-          "use"(%val) : (index) -> ()
-        }) : () -> ()
-      }) : () -> ()
+               threads(%tx, %ty, %tz) in (%block_x = %thrd, %block_y = %thrd, %block_z = %thrd) {
+      "use"(%float) : (f32) -> ()
+      "use"(%data) : (memref<?xf32,1>) -> ()
       // CHECK: gpu.terminator
       gpu.terminator
     }
     return
   }
 
   gpu.module @kernels {
     gpu.func @kernel_1(%arg0 : f32, %arg1 : memref<?xf32, 1>) attributes {gpu.kernel} {

mlir/test/Dialect/GPU/outlining.mlir

-// RUN: mlir-opt -gpu-kernel-outlining -split-input-file -verify-diagnostics %s | FileCheck %s
+// RUN: mlir-opt -gpu-kernel-outlining -split-input-file -verify-diagnostics %s | FileCheck %s -dump-input-on-failure
 
 // CHECK: module attributes {gpu.container_module}
 
 // CHECK-LABEL: func @launch()
 func @launch() {
   // CHECK: %[[ARG0:.*]] = "op"() : () -> f32
   %0 = "op"() : () -> (f32)
   // CHECK: %[[ARG1:.*]] = "op"() : () -> memref<?xf32, 1>
func @launch() {
   // CHECK: %[[BDIMZ:.*]] = constant 28
   %bDimZ = constant 28 : index
 
   // CHECK: "gpu.launch_func"(%[[GDIMX]], %[[GDIMY]], %[[GDIMZ]], %[[BDIMX]], %[[BDIMY]], %[[BDIMZ]], %[[ARG0]], %[[ARG1]]) {kernel = "launch_kernel", kernel_module = @launch_kernel} : (index, index, index, index, index, index, f32, memref<?xf32, 1>) -> ()
   // CHECK-NOT: gpu.launch blocks
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %gDimX, %grid_y = %gDimY,
                                        %grid_z = %gDimZ)
              threads(%tx, %ty, %tz) in (%block_x = %bDimX, %block_y = %bDimY,
-                                        %block_z = %bDimZ)
-             args(%arg0 = %0, %arg1 = %1) : f32, memref<?xf32, 1> {
-    "use"(%arg0): (f32) -> ()
+                                        %block_z = %bDimZ) {
+    "use"(%0): (f32) -> ()
     "some_op"(%bx, %block_x) : (index, index) -> ()
-    %42 = load %arg1[%tx] : memref<?xf32, 1>
+    %42 = load %1[%tx] : memref<?xf32, 1>
     gpu.terminator
   }
   return
 }
 
 
 // CHECK-LABEL: module @launch_kernel
 // CHECK-NEXT: gpu.func @launch_kernel
func @extra_constants(%arg0 : memref<?xf32>) {
   // CHECK: %[[CST:.*]] = constant 8 : index
   %cst = constant 8 : index
   %cst2 = constant 2 : index
   %cst3 = dim %arg0, 0 : memref<?xf32>
   // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %{{.*}}) {kernel = "extra_constants_kernel", kernel_module = @extra_constants_kernel} : (index, index, index, index, index, index, memref<?xf32>) -> ()
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,
                                        %grid_z = %cst)
              threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,
-                                        %block_z = %cst)
-             args(%kernel_arg0 = %cst2, %kernel_arg1 = %arg0, %kernel_arg2 = %cst3) : index, memref<?xf32>, index {
-    "use"(%kernel_arg0, %kernel_arg1, %kernel_arg2) : (index, memref<?xf32>, index) -> ()
+                                        %block_z = %cst) {
+    "use"(%cst2, %arg0, %cst3) : (index, memref<?xf32>, index) -> ()
     gpu.terminator
   }
   return
 }
 
 // CHECK-LABEL: func @extra_constants_kernel(%{{.*}}: memref<?xf32>)
 // CHECK: constant
 // CHECK: constant

mlir/test/mlir-cuda-runner/all-reduce-op.mlir

 // RUN: mlir-cuda-runner %s --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
 
 // CHECK-COUNT-8: [{{(5356, ){12}5356}}]
 func @main() {
   %arg = alloc() : memref<2x4x13xf32>
   %dst = memref_cast %arg : memref<2x4x13xf32> to memref<?x?x?xf32>
   %one = constant 1 : index
   %sx = dim %dst, 2 : memref<?x?x?xf32>
   %sy = dim %dst, 1 : memref<?x?x?xf32>
   %sz = dim %dst, 0 : memref<?x?x?xf32>
   call @mcuMemHostRegisterMemRef3dFloat(%dst) : (memref<?x?x?xf32>) -> ()
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
-             threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %sy, %block_z = %sz)
-             args(%kernel_dst = %dst) : memref<?x?x?xf32> {
+             threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %sy, %block_z = %sz) {
     %t0 = muli %tz, %block_y : index
     %t1 = addi %ty, %t0 : index
     %t2 = muli %t1, %block_x : index
     %idx = addi %tx, %t2 : index
     %t3 = index_cast %idx : index to i32
     %val = sitofp %t3 : i32 to f32
     %sum = "gpu.all_reduce"(%val) ({}) { op = "add" } : (f32) -> (f32)
-    store %sum, %kernel_dst[%tz, %ty, %tx] : memref<?x?x?xf32>
+    store %sum, %dst[%tz, %ty, %tx] : memref<?x?x?xf32>
     gpu.terminator
   }
   %U = memref_cast %dst : memref<?x?x?xf32> to memref<*xf32>
   call @print_memref_f32(%U) : (memref<*xf32>) -> ()
   return
 }
 
 func @mcuMemHostRegisterMemRef3dFloat(%ptr : memref<?x?x?xf32>)
 func @print_memref_f32(%ptr : memref<*xf32>)

mlir/test/mlir-cuda-runner/all-reduce-region.mlir

 // RUN: mlir-cuda-runner %s --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
 
 // CHECK: [{{(35, ){34}35}}]
 func @main() {
   %arg = alloc() : memref<35xf32>
   %dst = memref_cast %arg : memref<35xf32> to memref<?xf32>
   %one = constant 1 : index
   %sx = dim %dst, 0 : memref<?xf32>
   call @mcuMemHostRegisterMemRef1dFloat(%dst) : (memref<?xf32>) -> ()
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
-             threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %one, %block_z = %one)
-             args(%kernel_dst = %dst) : memref<?xf32> {
+             threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %one, %block_z = %one) {
     %val = index_cast %tx : index to i32
     %xor = "gpu.all_reduce"(%val) ({
     ^bb(%lhs : i32, %rhs : i32):
       %xor = xor %lhs, %rhs : i32
       "gpu.yield"(%xor) : (i32) -> ()
     }) : (i32) -> (i32)
     %res = sitofp %xor : i32 to f32
-    store %res, %kernel_dst[%tx] : memref<?xf32>
+    store %res, %dst[%tx] : memref<?xf32>
     gpu.terminator
   }
   %U = memref_cast %dst : memref<?xf32> to memref<*xf32>
   call @print_memref_f32(%U) : (memref<*xf32>) -> ()
   return
 }
 
 func @mcuMemHostRegisterMemRef1dFloat(%ptr : memref<?xf32>)
 func @print_memref_f32(memref<*xf32>)

mlir/test/mlir-cuda-runner/gpu-to-cubin.mlir

 // RUN: mlir-cuda-runner %s --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
 
 func @other_func(%arg0 : f32, %arg1 : memref<?xf32>) {
   %cst = constant 1 : index
   %cst2 = dim %arg1, 0 : memref<?xf32>
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst, %grid_z = %cst)
-             threads(%tx, %ty, %tz) in (%block_x = %cst2, %block_y = %cst, %block_z = %cst)
-             args(%kernel_arg0 = %arg0, %kernel_arg1 = %arg1) : f32, memref<?xf32> {
-    store %kernel_arg0, %kernel_arg1[%tx] : memref<?xf32>
+             threads(%tx, %ty, %tz) in (%block_x = %cst2, %block_y = %cst, %block_z = %cst) {
+    store %arg0, %arg1[%tx] : memref<?xf32>
     gpu.terminator
   }
   return
 }
 
 // CHECK: [1, 1, 1, 1, 1]
 func @main() {
   %arg0 = alloc() : memref<5xf32>

mlir/test/mlir-cuda-runner/shuffle.mlir

 // RUN: mlir-cuda-runner %s --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
 
 // CHECK: [4, 5, 6, 7, 0, 1, 2, 3, 12, -1, -1, -1, 8]
 func @main() {
   %arg = alloc() : memref<13xf32>
   %dst = memref_cast %arg : memref<13xf32> to memref<?xf32>
   %one = constant 1 : index
   %sx = dim %dst, 0 : memref<?xf32>
   call @mcuMemHostRegisterMemRef1dFloat(%dst) : (memref<?xf32>) -> ()
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
-             threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %one, %block_z = %one)
-             args(%kernel_dst = %dst) : memref<?xf32> {
+             threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %one, %block_z = %one) {
     %t0 = index_cast %tx : index to i32
     %val = sitofp %t0 : i32 to f32
     %width = index_cast %block_x : index to i32
     %offset = constant 4 : i32
     %shfl, %valid = gpu.shuffle %val, %offset, %width xor : f32
     cond_br %valid, ^bb1(%shfl : f32), ^bb0
   ^bb0:
     %m1 = constant -1.0 : f32
     br ^bb1(%m1 : f32)
   ^bb1(%value : f32):
-    store %value, %kernel_dst[%tx] : memref<?xf32>
+    store %value, %dst[%tx] : memref<?xf32>
     gpu.terminator
   }
   %U = memref_cast %dst : memref<?xf32> to memref<*xf32>
   call @print_memref_f32(%U) : (memref<*xf32>) -> ()
   return
 }
 
 func @mcuMemHostRegisterMemRef1dFloat(%ptr : memref<?xf32>)
 func @print_memref_f32(%ptr : memref<*xf32>)