Add 'gpu.terminator' operation.

[herhut]

Summary

The 'gpu.terminator' operation is used as the terminator for the
regions of gpu.launch. This is to disambugaute them from the
return operation on 'gpu.func' functions.

This is a breaking change and users of the gpu dialect will need
to adapt their code when producting 'gpu.launch' operations.

[herhut]

PTAL. I will do another breaking change that removes the use of region arguments in a followup.

[merge_guards_bot]

Unit tests: pass. 62255 tests passed, 0 failed and 827 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]

Thanks!

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

24

Nit: we may want to skip this terminator while printing, but I won't insist.

[herhut]

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

24

I was wondering the same. The only reason I left it printed was that I don't like a basic block without terminator. Confused my mental parser. So I only make them implicit for single-block operations. We could drop it if there is only a single block? Or is that too incoherent?

[ftynse]

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

24

It sounds like too much code for sugaring the IR. Let's keep it visible all the time, having it consistent between single- and multiple-block regions is a good argument.

[mehdi_amini]

The 'gpu.terminator' operation is used as the terminator for the regions of gpu.launch. This is to disambugaute them from the return operation on 'gpu.func' functions.

(Nit: typo on disambugaute)

[mehdi_amini]

D71961 is going in the direction of generalizing std.return ; what is the guideline for adding a new op vs re-using std.return?

[ftynse]

In D73620#1847103, @mehdi_amini wrote:

D71961 is going in the direction of generalizing std.return ; what is the guideline for adding a new op vs re-using std.return?

This diff introduces a terminator that is explicitly "not a return", due to the confusion that appeared in a previous change wrt its semantics. In particular, it is used inside gpu.launch potentially inside std.func, it is never expected to have operands. When the generalization of std.return, I would consider replacing gpu.return with std.return if compatible, but not gpu.terminator.

[mehdi_amini]

In D73620#1847083, @mehdi_amini wrote:

The 'gpu.terminator' operation is used as the terminator for the regions of gpu.launch. This is to disambugaute them from the return operation on 'gpu.func' functions.

(Nit: typo on disambugaute)

Did you fix this?

[mehdi_amini]

In D73620#1848957, @ftynse wrote:

In D73620#1847103, @mehdi_amini wrote:

D71961 is going in the direction of generalizing std.return ; what is the guideline for adding a new op vs re-using std.return?

This diff introduces a terminator that is explicitly "not a return", due to the confusion that appeared in a previous change wrt its semantics. In particular, it is used inside gpu.launch potentially inside std.func, it is never expected to have operands. When the generalization of std.return, I would consider replacing gpu.return with std.return if compatible, but not gpu.terminator.

I don't think I agree with this. The discussion on generalizing std.return goes against this as far as I understand it.

Why was this committed while we're discussing here?

[herhut]

In D73620#1850048, @mehdi_amini wrote:

Why was this committed while we're discussing here?

I misunderstood the comment it seems. I though the idea was that gpu.return should be replaced by a generalized return. From reading the discussion on 'std.return' I assumed this was meant to become a generalized way of returning from a function-like structure.

This change prepares this by splitting terminators that are function like, especially in the sense that their return operands correspond to the parent operation's result value. For the return form a launchOp, this is not the case. Currently, launchOp does not yield any value but that will change soon when we start making things asynchronous. So it does not contradict the work on a generic return.

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

let extraClassDeclaration = [{
     static constexpr unsigned kNumConfigRegionAttributes = 12;
   }];
 
   let parser = [{ return parseLaunchOp(parser, result); }];
   let printer = [{ printLaunchOp(p, *this); }];
   let verifier = [{ return ::verify(*this); }];
 }
 
-def GPU_ReturnOp : GPU_Op<"return", [Terminator]>, Arguments<(ins)>,
-    Results<(outs)> {
+def GPU_ReturnOp : GPU_Op<"return", [HasParent<"GPUFuncOp">, Terminator]>,
+    Arguments<(ins Variadic<AnyType>:$operands)>, Results<(outs)> {
+  let summary = "Terminator for GPU functions.";
+  let description = [{
+    A terminator operation for 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`.
+  }];
+
+  let builders = [OpBuilder<"Builder *builder, OperationState &result", " // empty">];
+
+  let parser = [{ return parseReturnOp(parser, result); }];
+  let printer = [{ p << getOperationName(); }];
+  let verifier = [{ return ::verify(*this); }];
+}
+
+def GPU_TerminatorOp : GPU_Op<"terminator", [HasParent<"LaunchOp">, 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(); }];

mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp

 
 /// Generate the body of the launch operation.
 template <typename OpTy>
 static LogicalResult
 createLaunchBody(OpBuilder &builder, OpTy rootForOp, gpu::LaunchOp launchOp,
                  unsigned numBlockDims, unsigned numThreadDims) {
   OpBuilder::InsertionGuard bodyInsertionGuard(builder);
   builder.setInsertionPointToEnd(&launchOp.body().front());
-  auto returnOp = builder.create<gpu::ReturnOp>(launchOp.getLoc());
+  auto terminatorOp = builder.create<gpu::TerminatorOp>(launchOp.getLoc());
 
-  rootForOp.getOperation()->moveBefore(returnOp);
+  rootForOp.getOperation()->moveBefore(terminatorOp);
   SmallVector<Value, 3> workgroupID, numWorkGroups;
   packIdAndNumId(launchOp.getBlockIds(), launchOp.getGridSize(), numBlockDims,
                  workgroupID, numWorkGroups);
 
   // Partition the loop for mapping to workgroups.
   auto loopOp = createGPULaunchLoops(rootForOp, workgroupID, numWorkGroups);
 
   // Iterate over the body of the loopOp and get the loops to partition for
void LoopToGpuConverter::createLaunch(OpTy rootForOp, OpTy innermostForOp,
   // 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
   // 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::ReturnOp>(terminatorLoc);
+  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.

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

static LogicalResult verify(LaunchOp op) {
   if (!op.body().empty()) {
     Block &entryBlock = op.body().front();
     if (entryBlock.getNumArguments() !=
         LaunchOp::kNumConfigOperands + op.getNumOperands())
       return op.emitOpError("unexpected number of region arguments");
   }
 
   // Block terminators without successors are expected to exit the kernel region
-  // and must be `gpu.launch`.
+  // and must be `gpu.terminator`.
   for (Block &block : op.body()) {
     if (block.empty())
       continue;
     if (block.back().getNumSuccessors() != 0)
       continue;
-    if (!isa<gpu::ReturnOp>(&block.back())) {
+    if (!isa<gpu::TerminatorOp>(&block.back())) {
       return block.back()
-                 .emitError("expected 'gpu.terminator' or a terminator with "
-                            "successors")
-                 .attachNote(op.getLoc())
-             << "in '" << LaunchOp::getOperationName() << "' body region";
+          .emitError()
+          .append("expected '", gpu::TerminatorOp::getOperationName(),
+                  "' or a terminator with successors")
+          .append("in '", LaunchOp::getOperationName(), "' body region");
     }
   }
 
   return success();
 }
 
 // Pretty-print the kernel grid/block size assignment as
 //   (%iter-x, %iter-y, %iter-z) in
if (failed(impl::parseFunctionSignature(
           isVariadic, resultTypes, resultAttrs)))
     return failure();
 
   if (entryArgs.empty() && !argTypes.empty())
     return parser.emitError(signatureLocation)
            << "gpu.func requires named arguments";
 
   // Construct the function type. More types will be added to the region, but
-  // not to the functiont type.
+  // not to the function type.
   Builder &builder = parser.getBuilder();
   auto type = builder.getFunctionType(argTypes, resultTypes);
   result.addAttribute(GPUFuncOp::getTypeAttrName(), TypeAttr::get(type));
 
   // Parse workgroup memory attributions.
   if (failed(parseAttributions(parser, GPUFuncOp::getWorkgroupKeyword(),
                                entryArgs, argTypes)))
     return failure();

 }
 
 /// Hook for FunctionLike verifier.
 LogicalResult GPUFuncOp::verifyType() {
   Type type = getTypeAttr().getValue();
   if (!type.isa<FunctionType>())
     return emitOpError("requires '" + getTypeAttrName() +
                        "' attribute of function type");
+
+  if (isKernel() && getType().getNumResults() != 0)
+    return emitOpError() << "expected void return type for kernel function";
+
   return success();
 }
 
 static LogicalResult verifyAttributions(Operation *op,
                                         ArrayRef<BlockArgument> attributions,
                                         unsigned memorySpace) {
   for (Value v : attributions) {
     auto type = v.getType().dyn_cast<MemRefType>();
if (failed(verifyAttributions(getOperation(), getWorkgroupAttributions(),
       failed(verifyAttributions(getOperation(), getPrivateAttributions(),
                                 GPUDialect::getPrivateAddressSpace())))
     return failure();
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// ReturnOp
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseReturnOp(OpAsmParser &parser, OperationState &result) {
+  llvm::SmallVector<OpAsmParser::OperandType, 4> operands;
+  llvm::SmallVector<Type, 4> types;
+  if (parser.parseOperandList(operands) ||
+      parser.parseOptionalColonTypeList(types) ||
+      parser.resolveOperands(operands, types, parser.getCurrentLocation(),
+                             result.operands))
+    return failure();
+
+  return success();
+}
+
+static LogicalResult verify(gpu::ReturnOp returnOp) {
+  GPUFuncOp function = returnOp.getParentOfType<GPUFuncOp>();
+
+  FunctionType funType = function.getType();
+
+  if (funType.getNumResults() != returnOp.operands().size())
+    return returnOp.emitOpError()
+        .append("expected ", funType.getNumResults(), " result operands")
+        .attachNote(function.getLoc())
+        .append("return type declared here");
+
+  for (auto pair : llvm::enumerate(
+           llvm::zip(function.getType().getResults(), returnOp.operands()))) {
+    Type type;
+    Value operand;
+    std::tie(type, operand) = pair.value();
+    if (type != operand.getType())
+      return returnOp.emitOpError() << "unexpected type `" << operand.getType()
+                                    << "' for operand #" << pair.index();
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // GPUModuleOp
 //===----------------------------------------------------------------------===//
 
 void GPUModuleOp::build(Builder *builder, OperationState &result,
                         StringRef name) {
   ensureTerminator(*result.addRegion(), *builder, result.location);
   result.attributes.push_back(builder->getNamedAttr(
       ::mlir::SymbolTable::getSymbolAttrName(), builder->getStringAttr(name)));

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

static gpu::LaunchFuncOp inlineBeneficiaryOps(gpu::GPUFuncOp kernelFunc,
   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.return` operations by `std.return` in the generated function.
+// `gpu.terminator` operations by `gpu.return` in the generated function.
 static gpu::GPUFuncOp outlineKernelFunc(gpu::LaunchOp launchOp) {
   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());
   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());
+  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) {

mlir/test/Conversion/LoopsToGPU/linalg_to_gpu.mlir

loop.for %i1 = %c3 to %c42 step %arg1 {
 
       // 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.return
+      // CHECK: gpu.terminator
     }
   }
   return
 }

mlir/test/Conversion/LoopsToGPU/step_one.mlir

// CHECK-11-SAME: args
           // Using remapped values instead of loop iterators.
           // CHECK-11:        {{.*}} = load %{{.*}}[%[[i]], %[[j]], %[[ii]], %[[jj]]] : memref<?x?x?x?xf32>
           // CHECK-22:        {{.*}} = load %{{.*}}[%[[i]], %[[j]], %[[ii]], %[[jj]]] : memref<?x?x?x?xf32>
           %0 = load %A[%i, %j, %ii, %jj] : memref<?x?x?x?xf32>
           // CHECK-11-NEXT:   store {{.*}}, %{{.*}}[%[[i]], %[[j]], %[[ii]], %[[jj]]] : memref<?x?x?x?xf32>
           // CHECK-22-NEXT:   store {{.*}}, %{{.*}}[%[[i]], %[[j]], %[[ii]], %[[jj]]] : memref<?x?x?x?xf32>
           store %0, %B[%i, %j, %ii, %jj] : memref<?x?x?x?xf32>
 
-          // CHECK-11: gpu.return
-          // CHECK-22: gpu.return
+          // CHECK-11: gpu.terminator
+          // CHECK-22: gpu.terminator
         }
       }
     }
   }
   return
 }
 

mlir/test/Dialect/GPU/canonicalize.mlir

gpu.launch blocks(%bx, %by, %bz) in (%sbx = %c1, %sby = %c1, %sbz = %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.return
+    gpu.terminator
   }
   return
 }
 

mlir/test/Dialect/GPU/invalid.mlir

 func @shuffle_unsupported_type(%arg0 : index, %arg1 : i32, %arg2 : i32) {
   // expected-error@+1 {{'gpu.shuffle' op requires value operand type to be f32 or i32}}
   %shfl, %pred = gpu.shuffle %arg0, %arg1, %arg2 xor : index
 }
 
 // -----
 
 module {
-  module @gpu_funcs attributes {gpu.kernel_module} {
+  gpu.module @gpu_funcs {
     // expected-error @+1 {{custom op 'gpu.func' gpu.func requires named arguments}}
     gpu.func @kernel_1(f32, f32) {
     ^bb0(%arg0: f32):
       gpu.return
     }
   }
 }
 

 module {
   module @gpu_funcs attributes {gpu.kernel_module} {
     // expected-error @+1 {{expected memory space 5 in attribution}}
     gpu.func @kernel() private(%0: memref<4xf32>) {
       gpu.return
     }
   }
 }
+
+// -----
+
+module {
+  module @gpu_funcs attributes {gpu.kernel_module} {
+    // expected-error @+1 {{expected memory space 5 in attribution}}
+    gpu.func @kernel() private(%0: memref<4xf32>) {
+      gpu.return
+    }
+  }
+}
+
+// -----
+
+module {
+  gpu.module @gpu_funcs {
+    // expected-note @+1 {{return type declared here}}
+    gpu.func @kernel() {
+      %0 = constant 0 : index
+      // expected-error @+1 {{'gpu.return' op expected 0 result operands}}
+      gpu.return %0 : index
+    }
+  }
+}
+
+// -----
+
+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
+    }
+  }
+}
 No newline at end of file

mlir/test/Dialect/GPU/ops.mlir

 // RUN: mlir-opt %s | FileCheck %s
 
 module attributes {gpu.container_module} {
 
   // CHECK-LABEL:func @no_args(%{{.*}}: index)
   func @no_args(%sz : index) {
     // CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}})
     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) {
-      // CHECK: gpu.return
-      gpu.return
+      // 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.return
-      gpu.return
+      // 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>
     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.return
-      gpu.return
+      // 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) -> ()
         }) : () -> ()
       }) : () -> ()
-      // CHECK: gpu.return
-      gpu.return
+      // CHECK: gpu.terminator
+      gpu.terminator
     }
     return
   }
 
   gpu.module @kernels {
     gpu.func @kernel_1(%arg0 : f32, %arg1 : memref<?xf32, 1>) attributes {gpu.kernel} {
       %tIdX = "gpu.thread_id"() {dimension = "x"} : () -> (index)
       %tIdY = "gpu.thread_id"() {dimension = "y"} : () -> (index)
func @foo() {
     "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1)
     { kernel = "kernel_2", kernel_module = @kernels }
         : (index, index, index, index, index, index, f32, memref<?xf32, 1>) -> ()
 
     return
   }
 
   module @gpu_funcs attributes {gpu.kernel_module} {
-    // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32) -> f32
+    // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32)
     // CHECK:       workgroup
     // CHECK:       private
     // CHECK:       attributes
-    gpu.func @kernel_1(%arg0: f32) -> f32
+    gpu.func @kernel_1(%arg0: f32)
         workgroup(%arg1: memref<42xf32, 3>)
         private(%arg2: memref<2xf32, 5>, %arg3: memref<1xf32, 5>)
         kernel
         attributes {foo="bar"} {
       "use"(%arg1) : (memref<42xf32, 3>) -> ()
       "use"(%arg2) : (memref<2xf32, 5>) -> ()
       "use"(%arg3) : (memref<1xf32, 5>) -> ()
       gpu.return

mlir/test/Dialect/GPU/outlining.mlir

func @launch() {
   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) -> ()
     "some_op"(%bx, %block_x) : (index, index) -> ()
     %42 = load %arg1[%tx] : memref<?xf32, 1>
-    gpu.return
+    gpu.terminator
   }
   return
 }
 
 
 // CHECK-LABEL: module @launch_kernel
 // CHECK-NEXT: gpu.func @launch_kernel
 // CHECK-SAME: (%[[KERNEL_ARG0:.*]]: f32, %[[KERNEL_ARG1:.*]]: memref<?xf32, 1>)

 func @multiple_launches() {
   // CHECK: %[[CST:.*]] = constant 8 : index
   %cst = constant 8 : index
   // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]]) {kernel = "multiple_launches_kernel", kernel_module = @multiple_launches_kernel} : (index, index, index, index, index, index) -> ()
   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) {
-    gpu.return
+    gpu.terminator
   }
   // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]]) {kernel = "multiple_launches_kernel", kernel_module = @multiple_launches_kernel_0} : (index, index, index, index, index, index) -> ()
   gpu.launch blocks(%bx2, %by2, %bz2) in (%grid_x2 = %cst, %grid_y2 = %cst,
                                           %grid_z2 = %cst)
              threads(%tx2, %ty2, %tz2) in (%block_x2 = %cst, %block_y2 = %cst,
                                            %block_z2 = %cst) {
-    gpu.return
+    gpu.terminator
   }
   return
 }
 
 // CHECK: module @multiple_launches_kernel
 // CHECK: func @multiple_launches_kernel
 // CHECK: module @multiple_launches_kernel_0
 // CHECK: func @multiple_launches_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) -> ()
-    gpu.return
+    gpu.terminator
   }
   return
 }
 
 // CHECK-LABEL: func @extra_constants_kernel(%{{.*}}: memref<?xf32>)
 // CHECK: constant
 // CHECK: constant
 
 // -----
 
 llvm.mlir.global internal @global(42 : i64) : !llvm.i64
 
 func @function_call(%arg0 : memref<?xf32>) {
   %cst = constant 8 : index
   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) {
     call @device_function() : () -> ()
     call @device_function() : () -> ()
     %0 = llvm.mlir.addressof @global : !llvm<"i64*">
-    gpu.return
+    gpu.terminator
   }
   return
 }
 
 func @device_function() {
   call @recursive_device_function() : () -> ()
-  gpu.return
+  return
 }
 
 func @recursive_device_function() {
   call @recursive_device_function() : () -> ()
-  gpu.return
+  return
 }
 
 // CHECK: gpu.module @function_call_kernel {
 // CHECK:   gpu.func @function_call_kernel()
 // CHECK:     call @device_function() : () -> ()
 // CHECK:     call @device_function() : () -> ()
 // CHECK:     llvm.mlir.addressof @global : !llvm<"i64*">
+// CHECK:     gpu.return
 //
 // CHECK:   llvm.mlir.global internal @global(42 : i64) : !llvm.i64
 //
 // CHECK:   func @device_function()
 // CHECK:   func @recursive_device_function()
 // CHECK-NOT:   func @device_function

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

gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
     %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>
-    gpu.return
+    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

gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_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>
-    gpu.return
+    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>
-    gpu.return
+    gpu.terminator
   }
   return
 }
 
 // CHECK: [1, 1, 1, 1, 1]
 func @main() {
   %arg0 = alloc() : memref<5xf32>
   %21 = constant 5 : i32

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

gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
     %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>
-    gpu.return
+    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>)