func @kernel_1(%arg0 : f32, %arg1 : !llvm<"float*">)
attributes { nvvm.kernel: true } {
- // Operations that produce block/thread IDs and dimensions will be injected
- // when outlining the `gpu.launch` body to a function called by
- // `gpu.launch_func`.
- // TODO(tjoerg): Implement gpu.launch body outlining.
+ // Operations that produce block/thread IDs and dimensions are injected when
+ // outlining the `gpu.launch` body to a function called by `gpu.launch_func`.
%tIdX = "gpu.thread_id"() {dimension: "x"} : () -> (index)
%tIdY = "gpu.thread_id"() {dimension: "y"} : () -> (index)
%tIdZ = "gpu.thread_id"() {dimension: "z"} : () -> (index)
KernelDim3 getGridSize();
/// Get the SSA values corresponding to kernel block size.
KernelDim3 getBlockSize();
+ /// Append the operand values passed as kernel arguments to `out`.
+ void getKernelOperandValues(SmallVectorImpl<Value *> *out);
+ /// Append the operand types passed as kernel arguments to `out`.
+ void getKernelOperandTypes(SmallVectorImpl<Type> *out);
LogicalResult verify();
--- /dev/null
+//===- Passes.h - Pass Entrypoints ------------------------------*- C++ -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This header file defines prototypes that expose pass constructors.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_GPU_PASSES_H_
+#define MLIR_GPU_PASSES_H_
+
+namespace mlir {
+
+class ModulePassBase;
+
+ModulePassBase *createGpuKernelOutliningPass();
+
+} // namespace mlir
+
+#endif // MLIR_GPU_PASSES_H_
return KernelDim3{args[9], args[10], args[11]};
}
+void LaunchOp::getKernelOperandValues(SmallVectorImpl<Value *> *out) {
+ out->reserve(getNumOperands() - kNumConfigOperands + out->size());
+ for (int i = kNumConfigOperands; i < getNumOperands(); ++i) {
+ out->push_back(getOperand(i));
+ }
+}
+
+void LaunchOp::getKernelOperandTypes(SmallVectorImpl<Type> *out) {
+ out->reserve(getNumOperands() - kNumConfigOperands + out->size());
+ for (int i = kNumConfigOperands; i < getNumOperands(); ++i) {
+ out->push_back(getOperand(i)->getType());
+ }
+}
+
LogicalResult LaunchOp::verify() {
// Kernel launch takes kNumConfigOperands leading operands for grid/block
// sizes and transforms them into kNumConfigRegionAttributes region arguments
--- /dev/null
+//===- KernelOutlining.cpp - Implementation of GPU kernel outling ---------===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file implements the GPU dialect kernel outlining pass.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/GPU/GPUDialect.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/StandardOps/Ops.h"
+
+using namespace mlir;
+
+namespace {
+
+template <typename OpTy>
+void createForAllDimensions(FuncBuilder &builder, Location loc,
+ SmallVectorImpl<Value *> *values) {
+ for (string dim : {"x", "y", "z"}) {
+ Value *v = builder.create<OpTy>(loc, builder.getIndexType(),
+ builder.getStringAttr(dim));
+ values->push_back(v);
+ }
+}
+
+// Add operations generating block/thread ids and gird/block dimensions at the
+// beginning of `kernelFunc` and replace uses of the respective function args.
+void injectGpuIndexOperations(Module &module, Location loc,
+ Function *kernelFunc) {
+ Builder builder(&module);
+ FuncBuilder funcBuilder(kernelFunc);
+ SmallVector<Value *, 12> indexOps;
+ createForAllDimensions<gpu::BlockId>(funcBuilder, loc, &indexOps);
+ createForAllDimensions<gpu::ThreadId>(funcBuilder, loc, &indexOps);
+ createForAllDimensions<gpu::GridDim>(funcBuilder, loc, &indexOps);
+ createForAllDimensions<gpu::BlockDim>(funcBuilder, loc, &indexOps);
+ // Replace the leading 12 function args with the respective thread/block index
+ // operations. Iterate backwards since args are erased and indices change.
+ for (int i = 11; i >= 0; --i) {
+ auto &firstBlock = kernelFunc->getBody().getBlocks().front();
+ firstBlock.getArgument(i)->replaceAllUsesWith(indexOps[i]);
+ firstBlock.eraseArgument(i);
+ }
+}
+
+// Outline the `gpu.launch` operation body into a kernel function.
+Function *outlineKernelFunc(Module &module, gpu::LaunchOp &launchOp) {
+ Location loc = launchOp.getLoc();
+ SmallVector<Type, 4> kernelOperandTypes;
+ launchOp.getKernelOperandTypes(&kernelOperandTypes);
+ FunctionType type =
+ FunctionType::get(kernelOperandTypes, {}, module.getContext());
+ string kernelFuncName =
+ Twine(launchOp.getOperation()->getFunction()->getName(), "_kernel").str();
+ mlir::BlockAndValueMapping mapper;
+ Function *outlinedFunc = new mlir::Function(loc, kernelFuncName, type);
+ outlinedFunc->getBody().takeBody(launchOp.getBody());
+ injectGpuIndexOperations(module, loc, outlinedFunc);
+ module.getFunctions().push_back(outlinedFunc);
+ return outlinedFunc;
+}
+
+// Replace `gpu.launch` operations with an `gpu.launch_func` operation launching
+// `kernelFunc`.
+void convertToLaunchFuncOp(gpu::LaunchOp &launchOp, Function *kernelFunc) {
+ Location loc = launchOp.getLoc();
+ FuncBuilder funcBuilder(launchOp);
+ SmallVector<Value *, 4> kernelOperandValues;
+ launchOp.getKernelOperandValues(&kernelOperandValues);
+ // TODO(tjoerg): Pass KernelDims rather than individual values.
+ funcBuilder.create<gpu::LaunchFuncOp>(
+ loc, kernelFunc, launchOp.getOperand(0), launchOp.getOperand(1),
+ launchOp.getOperand(2), launchOp.getOperand(3), launchOp.getOperand(4),
+ launchOp.getOperand(5), kernelOperandValues);
+ launchOp.erase();
+}
+
+} // namespace
+
+class GpuKernelOutliningPass : public ModulePass<GpuKernelOutliningPass> {
+public:
+ void runOnModule() override {
+ for (auto &func : getModule()) {
+ func.walk<mlir::gpu::LaunchOp>([&](mlir::gpu::LaunchOp op) {
+ Function *outlinedFunc = outlineKernelFunc(getModule(), op);
+ convertToLaunchFuncOp(op, outlinedFunc);
+ });
+ }
+ }
+};
+
+ModulePassBase *createGpuKernelOutliningPass() {
+ return new GpuKernelOutliningPass();
+}
+
+static PassRegistration<GpuKernelOutliningPass>
+ pass("gpu-kernel-outlining",
+ "Outline gpu.launch bodies to kernel functions.");
--- /dev/null
+// RUN: mlir-opt -gpu-kernel-outlining -split-input-file %s | FileCheck %s
+
+func @launch() {
+ %0 = "op"() : () -> (f32)
+ %1 = "op"() : () -> (memref<?xf32, 1>)
+ %cst = constant 8 : index
+
+ // CHECK: "gpu.launch_func"(%c8, %c8, %c8, %c8, %c8, %c8, %0, %1) {kernel: @launch_kernel : (f32, memref<?xf32, 1>) -> ()} : (index, index, index, index, index, index, f32, memref<?xf32, 1>) -> ()
+ // CHECK-NOT: gpu.launch blocks
+ 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(%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>
+ return
+ }
+ return
+}
+
+// CHECK: func @launch_kernel(%arg0: f32, %arg1: memref<?xf32, 1>)
+// CHECK-NEXT: %0 = "gpu.block_id"() {dimension: "x"} : () -> index
+// CHECK-NEXT: %1 = "gpu.block_id"() {dimension: "y"} : () -> index
+// CHECK-NEXT: %2 = "gpu.block_id"() {dimension: "z"} : () -> index
+// CHECK-NEXT: %3 = "gpu.thread_id"() {dimension: "x"} : () -> index
+// CHECK-NEXT: %4 = "gpu.thread_id"() {dimension: "y"} : () -> index
+// CHECK-NEXT: %5 = "gpu.thread_id"() {dimension: "z"} : () -> index
+// CHECK-NEXT: %6 = "gpu.grid_dim"() {dimension: "x"} : () -> index
+// CHECK-NEXT: %7 = "gpu.grid_dim"() {dimension: "y"} : () -> index
+// CHECK-NEXT: %8 = "gpu.grid_dim"() {dimension: "z"} : () -> index
+// CHECK-NEXT: %9 = "gpu.block_dim"() {dimension: "x"} : () -> index
+// CHECK-NEXT: %10 = "gpu.block_dim"() {dimension: "y"} : () -> index
+// CHECK-NEXT: %11 = "gpu.block_dim"() {dimension: "z"} : () -> index
+// CHECK-NEXT: "use"(%arg0) : (f32) -> ()
+// CHECK-NEXT: "some_op"(%0, %9) : (index, index) -> ()
+// CHECK-NEXT: %12 = load %arg1[%3] : memref<?xf32, 1>
+
+// -----
+
+func @multiple_launches() {
+ %cst = constant 8 : index
+ // CHECK: "gpu.launch_func"(%c8, %c8, %c8, %c8, %c8, %c8) {kernel: @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) {
+ return
+ }
+ // CHECK: "gpu.launch_func"(%c8, %c8, %c8, %c8, %c8, %c8) {kernel: @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) {
+ return
+ }
+ return
+}
+
+// CHECK: func @multiple_launches_kernel()
+// CHECK: func @multiple_launches_kernel_0()
projects / platform / upstream / llvm.git / commitdiff