[mlir][gpu] Split ops sinking from gpu-kernel-outlining pass into separate pass

[Hardcode84]

Summary

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.

[Hardcode84]

fix missing return

[Hardcode84]

fix Passes.td

[herhut]

Thanks for splitting this out.

Consider adding an explanation to the commit message to let users know what they have to change to get the old behavior.

mlir/include/mlir/Dialect/GPU/Passes.td

14

As this is now hardwired to sinking index computations, maybe call it gpu-launch-sink-index-computations?

[Hardcode84]

rename pass

[akuegel]

mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp did use the pass, was it intentional to change behavior, or should it now also have the new pass before outlining?

[Hardcode84]

mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp did use the pass, was it intentional to change behavior, or should it now also have the new pass before outlining?

It should be safe to add this new pass before outlining to exactly preserve old behavior but I'm don't know if it really needed there.

mlir/include/mlir/Dialect/GPU/Passes.h

 
 namespace llvm {
 class TargetMachine;
 class LLVMContext;
 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>>
 createGpuKernelOutliningPass(StringRef dataLayoutStr = StringRef());
 
 /// Rewrites a function region so that GPU ops execute asynchronously.
 std::unique_ptr<OperationPass<FuncOp>> createGpuAsyncRegionPass();
 

mlir/include/mlir/Dialect/GPU/Passes.td

 //===-- Passes.td - GPU pass definition file ---------------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_GPU_PASSES
 #define MLIR_DIALECT_GPU_PASSES
 
 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()";
   let dependentDialects = ["mlir::DLTIDialect"];
 }
 
 def GpuAsyncRegionPass : Pass<"gpu-async-region", "FuncOp"> {
   let summary = "Make GPU ops async";
   let constructor = "mlir::createGpuAsyncRegionPass()";
   let dependentDialects = ["async::AsyncDialect"];
 }
 
 #endif // MLIR_DIALECT_GPU_PASSES

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

static void injectGpuIndexOperations(Location loc, Region &launchFuncOpBody,
   // operations. Iterate backwards since args are erased and indices change.
   for (const auto &indexOp : enumerate(indexOps))
     map.map(firstBlock.getArgument(indexOp.index()), indexOp.value());
 }
 
 /// 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);
 }
 
 /// For a given operation `op`, computes whether it is beneficial to sink the
 /// operation into the kernel. An operation can be sunk if doing so does not
 /// introduce new kernel arguments. Whether a value is already available in the
 /// kernel (and hence does not introduce new arguments) is checked by
static void convertToLaunchFuncOp(gpu::LaunchOp launchOp,
   builder.create<gpu::LaunchFuncOp>(
       launchOp.getLoc(), kernelFunc, launchOp.getGridSizeOperandValues(),
       launchOp.getBlockSizeOperandValues(), launchOp.dynamicSharedMemorySize(),
       operands);
   launchOp.erase();
 }
 
 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
 /// 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
void runOnOperation() override {
     for (auto func : getOperation().getOps<FuncOp>()) {
       // Insert just after the function.
       Block::iterator insertPt(func->getNextNode());
       auto funcWalkResult = func.walk([&](gpu::LaunchOp op) {
         SetVector<Value> operands;
         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);
 
         // 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);
Option<std::string> dataLayoutStr{
       llvm::cl::desc("String containing the data layout specification to be "
                      "attached to the GPU kernel module")};
 
   DataLayoutSpecInterface dataLayoutSpec;
 };
 
 } // 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);
 }

mlir/lib/Dialect/GPU/Transforms/PassDetail.h

 //===- PassDetail.h - GPU Pass class details --------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef DIALECT_GPU_TRANSFORMS_PASSDETAIL_H_
 #define DIALECT_GPU_TRANSFORMS_PASSDETAIL_H_
 
 #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 {
 
 #define GEN_PASS_CLASSES
 #include "mlir/Dialect/GPU/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // DIALECT_GPU_TRANSFORMS_PASSDETAIL_H_

mlir/test/Dialect/GPU/outlining.mlir

-// 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}
 
 // CHECK-LABEL: func @launch()
 func @launch() {
   // CHECK: %[[ARG0:.*]] = "op"() : () -> f32
   %0 = "op"() : () -> (f32)
   // CHECK: %[[ARG1:.*]] = "op"() : () -> memref<?xf32, 1>

mlir/test/Dialect/GPU/sink-ops.mlir

+// 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
+}