[Mlir-commits] [mlir] d271fc0 - [mlir][gpu] Split ops sinking from gpu-kernel-outlining pass into separate pass
Ivan Butygin
llvmlistbot at llvm.org
Wed Feb 16 23:51:42 PST 2022
Author: Ivan Butygin
Date: 2022-02-17T10:34:20+03:00
New Revision: d271fc04d5b97b12e6b797c6067d3c96a8d7470e
URL: https://github.com/llvm/llvm-project/commit/d271fc04d5b97b12e6b797c6067d3c96a8d7470e
DIFF: https://github.com/llvm/llvm-project/commit/d271fc04d5b97b12e6b797c6067d3c96a8d7470e.diff
LOG: [mlir][gpu] Split ops sinking from gpu-kernel-outlining pass into separate pass
Previously `gpu-kernel-outlining` pass was also doing index computation sinking into gpu.launch before actual outlining.
Split ops sinking from `gpu-kernel-outlining` pass into separate pass, so users can use theirs own sinking pass before outlining.
To achieve old behavior users will need to call both passes: `-gpu-launch-sink-index-computations -gpu-kernel-outlining`.
Differential Revision: https://reviews.llvm.org/D119932
Added:
mlir/test/Dialect/GPU/sink-ops.mlir
Modified:
mlir/include/mlir/Dialect/GPU/Passes.h
mlir/include/mlir/Dialect/GPU/Passes.td
mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
mlir/lib/Dialect/GPU/Transforms/PassDetail.h
mlir/test/Dialect/GPU/outlining.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/GPU/Passes.h b/mlir/include/mlir/Dialect/GPU/Passes.h
index c9c6f8668b4d3..729363ace255b 100644
--- a/mlir/include/mlir/Dialect/GPU/Passes.h
+++ b/mlir/include/mlir/Dialect/GPU/Passes.h
@@ -23,6 +23,10 @@ class Module;
} // namespace llvm
namespace mlir {
+/// Pass that moves ops which are likely an index computation into gpu.launch
+/// body.
+std::unique_ptr<Pass> createGpuLauchSinkIndexComputationsPass();
+
/// Replaces `gpu.launch` with `gpu.launch_func` by moving the region into
/// a separate kernel function.
std::unique_ptr<OperationPass<ModuleOp>>
diff --git a/mlir/include/mlir/Dialect/GPU/Passes.td b/mlir/include/mlir/Dialect/GPU/Passes.td
index eaabc6ea36012..0380e1db5e01e 100644
--- a/mlir/include/mlir/Dialect/GPU/Passes.td
+++ b/mlir/include/mlir/Dialect/GPU/Passes.td
@@ -11,6 +11,12 @@
include "mlir/Pass/PassBase.td"
+def GpuLaunchSinkIndexComputations : Pass<"gpu-launch-sink-index-computations"> {
+ let summary = "Sink index computations into gpu.launch body";
+ let constructor = "mlir::createGpuLauchSinkIndexComputationsPass()";
+ let dependentDialects = ["mlir::gpu::GPUDialect"];
+}
+
def GpuKernelOutlining : Pass<"gpu-kernel-outlining", "ModuleOp"> {
let summary = "Outline gpu.launch bodies to kernel functions";
let constructor = "mlir::createGpuKernelOutliningPass()";
diff --git a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
index c8b9b9b2bc9a2..3b0f51444e027 100644
--- a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
@@ -59,7 +59,7 @@ static void injectGpuIndexOperations(Location loc, Region &launchFuncOpBody,
/// Identifies operations that are beneficial to sink into kernels. These
/// operations may not have side-effects, as otherwise sinking (and hence
/// duplicating them) is not legal.
-static bool isLikelyAnIndexComputatio(Operation *op) {
+static bool isLikelyAnIndexComputation(Operation *op) {
return isa<arith::ConstantOp, ConstantOp, memref::DimOp, arith::SelectOp,
arith::CmpIOp>(op);
}
@@ -232,6 +232,26 @@ static void convertToLaunchFuncOp(gpu::LaunchOp launchOp,
}
namespace {
+/// Pass that moves ops which are likely an index computation into gpu.launch
+/// body.
+class GpuLaunchSinkIndexComputationsPass
+ : public GpuLaunchSinkIndexComputationsBase<
+ GpuLaunchSinkIndexComputationsPass> {
+public:
+ void runOnOperation() override {
+ Operation *op = getOperation();
+ if (op->walk([](gpu::LaunchOp launch) {
+ // Pull in instructions that can be sunk
+ if (failed(sinkOperationsIntoLaunchOp(launch,
+ isLikelyAnIndexComputation)))
+ return WalkResult::interrupt();
+
+ return WalkResult::advance();
+ }).wasInterrupted())
+ signalPassFailure();
+ }
+};
+
/// 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
@@ -280,9 +300,6 @@ class GpuKernelOutliningPass
std::string kernelFnName =
Twine(op->getParentOfType<FuncOp>().getName(), "_kernel").str();
- // Pull in instructions that can be sunk
- if (failed(sinkOperationsIntoLaunchOp(op, isLikelyAnIndexComputatio)))
- return WalkResult::interrupt();
gpu::GPUFuncOp outlinedFunc =
outlineKernelFuncImpl(op, kernelFnName, operands);
@@ -360,6 +377,10 @@ class GpuKernelOutliningPass
} // namespace
+std::unique_ptr<Pass> mlir::createGpuLauchSinkIndexComputationsPass() {
+ return std::make_unique<GpuLaunchSinkIndexComputationsPass>();
+}
+
std::unique_ptr<OperationPass<ModuleOp>>
mlir::createGpuKernelOutliningPass(StringRef dataLayoutStr) {
return std::make_unique<GpuKernelOutliningPass>(dataLayoutStr);
diff --git a/mlir/lib/Dialect/GPU/Transforms/PassDetail.h b/mlir/lib/Dialect/GPU/Transforms/PassDetail.h
index 44e99e00fa155..faa9d3cf7231a 100644
--- a/mlir/lib/Dialect/GPU/Transforms/PassDetail.h
+++ b/mlir/lib/Dialect/GPU/Transforms/PassDetail.h
@@ -11,6 +11,7 @@
#include "mlir/Dialect/Async/IR/Async.h"
#include "mlir/Dialect/DLTI/DLTI.h"
+#include "mlir/Dialect/GPU/GPUDialect.h"
#include "mlir/Pass/Pass.h"
namespace mlir {
diff --git a/mlir/test/Dialect/GPU/outlining.mlir b/mlir/test/Dialect/GPU/outlining.mlir
index 4b15e5b449066..fc418ca442c46 100644
--- a/mlir/test/Dialect/GPU/outlining.mlir
+++ b/mlir/test/Dialect/GPU/outlining.mlir
@@ -1,5 +1,5 @@
-// RUN: mlir-opt -allow-unregistered-dialect -gpu-kernel-outlining -split-input-file -verify-diagnostics %s | FileCheck %s
-// RUN: mlir-opt -allow-unregistered-dialect -gpu-kernel-outlining=data-layout-str='#dlti.dl_spec<#dlti.dl_entry<index,32:i32>>' -split-input-file %s | FileCheck --check-prefix CHECK-DL %s
+// RUN: mlir-opt -allow-unregistered-dialect -gpu-launch-sink-index-computations -gpu-kernel-outlining -split-input-file -verify-diagnostics %s | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect -gpu-launch-sink-index-computations -gpu-kernel-outlining=data-layout-str='#dlti.dl_spec<#dlti.dl_entry<index,32:i32>>' -split-input-file %s | FileCheck --check-prefix CHECK-DL %s
// CHECK: module attributes {gpu.container_module}
diff --git a/mlir/test/Dialect/GPU/sink-ops.mlir b/mlir/test/Dialect/GPU/sink-ops.mlir
new file mode 100644
index 0000000000000..e2b4c238b9ce2
--- /dev/null
+++ b/mlir/test/Dialect/GPU/sink-ops.mlir
@@ -0,0 +1,100 @@
+// RUN: mlir-opt -allow-unregistered-dialect -gpu-launch-sink-index-computations -split-input-file -verify-diagnostics %s | FileCheck %s
+
+
+// CHECK-LABEL: @extra_constants
+// CHECK-SAME: %[[ARG0:.*]]: memref<?xf32>
+func @extra_constants(%arg0: memref<?xf32>) {
+ %cst = arith.constant 8 : index
+ %cst2 = arith.constant 2 : index
+ %c0 = arith.constant 0 : index
+ %cst3 = memref.dim %arg0, %c0 : memref<?xf32>
+ // CHECK: 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) {
+ // CHECK-NOT: arith.constant 8
+ // CHECK: %[[CST2:.*]] = arith.constant 2
+ // CHECK-NEXT: %[[CST0:.*]] = arith.constant 0
+ // CHECK-NEXT: %[[DIM:.*]] = memref.dim %[[ARG0]], %[[CST0]]
+ // CHECK-NEXT: "use"(%[[CST2]], %[[ARG0]], %[[DIM]]) : (index, memref<?xf32>, index) -> ()
+ // CHECK-NEXT: gpu.terminator
+ "use"(%cst2, %arg0, %cst3) : (index, memref<?xf32>, index) -> ()
+ gpu.terminator
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @extra_constants_not_inlined
+// CHECK-SAME: %[[ARG0:.*]]: memref<?xf32>
+func @extra_constants_not_inlined(%arg0: memref<?xf32>) {
+ %cst = arith.constant 8 : index
+ %cst2 = arith.constant 2 : index
+ %c0 = arith.constant 0 : index
+ // CHECK: %[[CST_X:.*]] = "secret_constant"()
+ %cst3 = "secret_constant"() : () -> index
+ // CHECK: 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) {
+ // CHECK-NOT: arith.constant 8
+ // CHECK-NOT: "secret_constant"()
+ // CHECK: %[[CST2:.*]] = arith.constant 2
+ // CHECK-NEXT: "use"(%[[CST2]], %[[ARG0]], %[[CST_X]]) : (index, memref<?xf32>, index) -> ()
+ // CHECK-NEXT: gpu.terminator
+ "use"(%cst2, %arg0, %cst3) : (index, memref<?xf32>, index) -> ()
+ gpu.terminator
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @multiple_uses
+func @multiple_uses(%arg0 : memref<?xf32>) {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ // CHECK: gpu.launch blocks
+ gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1,
+ %grid_z = %c1)
+ threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1,
+ %block_z = %c1) {
+ // CHECK: %[[C2:.*]] = arith.constant 2
+ // CHECK-NEXT: "use1"(%[[C2]], %[[C2]])
+ // CHECK-NEXT: "use2"(%[[C2]])
+ // CHECK-NEXT: gpu.terminator
+ "use1"(%c2, %c2) : (index, index) -> ()
+ "use2"(%c2) : (index) -> ()
+ gpu.terminator
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @multiple_uses2
+func @multiple_uses2(%arg0 : memref<*xf32>) {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %d = memref.dim %arg0, %c2 : memref<*xf32>
+ // CHECK: gpu.launch blocks
+ gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1,
+ %grid_z = %c1)
+ threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1,
+ %block_z = %c1) {
+ // CHECK: %[[C2:.*]] = arith.constant 2 : index
+ // CHECK: %[[D:.*]] = memref.dim %[[ARG:.*]], %[[C2]]
+ // CHECK: "use1"(%[[D]])
+ // CHECK: "use2"(%[[C2]], %[[C2]])
+ // CHECK: "use3"(%[[ARG]])
+ // CHECK: gpu.terminator
+ "use1"(%d) : (index) -> ()
+ "use2"(%c2, %c2) : (index, index) -> ()
+ "use3"(%arg0) : (memref<*xf32>) -> ()
+ gpu.terminator
+ }
+ return
+}
More information about the Mlir-commits
mailing list