From 318ff019cf703ece10df791b2da8307f3248baee Mon Sep 17 00:00:00 2001
From: Stephan Herhut <herhut@google.com>
Date: Mon, 9 Sep 2019 01:15:10 -0700
Subject: [PATCH] Addressing some late review comments on kernel inlining.

Just formatting and better lit tests, no functional change.

PiperOrigin-RevId: 267942907
---
 .../lib/Dialect/GPU/Transforms/KernelOutlining.cpp | 30 +++++++++++-----------
 mlir/test/Dialect/GPU/outlining.mlir               |  4 ++-
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
index f464b09..26449f6 100644
--- a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
@@ -74,22 +74,22 @@ static gpu::LaunchFuncOp inlineConstants(FuncOp kernelFunc,
     firstBlock.getArgument(i)->replaceAllUsesWith(newConstant->getResult(0));
     firstBlock.eraseArgument(i);
   }
-  if (newLaunchArgs.size() != launch.getNumKernelOperands()) {
-    std::reverse(newLaunchArgs.begin(), newLaunchArgs.end());
-    OpBuilder LaunchBuilder(launch);
-    SmallVector<Type, 8> newArgumentTypes;
-    newArgumentTypes.reserve(firstBlock.getNumArguments());
-    for (auto value : firstBlock.getArguments()) {
-      newArgumentTypes.push_back(value->getType());
-    }
-    kernelFunc.setType(LaunchBuilder.getFunctionType(newArgumentTypes, {}));
-    auto newLaunch = LaunchBuilder.create<gpu::LaunchFuncOp>(
-        launch.getLoc(), kernelFunc, launch.getGridSizeOperandValues(),
-        launch.getBlockSizeOperandValues(), newLaunchArgs);
-    launch.erase();
-    return newLaunch;
+  if (newLaunchArgs.size() == launch.getNumKernelOperands())
+    return launch;
+
+  std::reverse(newLaunchArgs.begin(), newLaunchArgs.end());
+  OpBuilder LaunchBuilder(launch);
+  SmallVector<Type, 8> newArgumentTypes;
+  newArgumentTypes.reserve(firstBlock.getNumArguments());
+  for (auto value : firstBlock.getArguments()) {
+    newArgumentTypes.push_back(value->getType());
   }
-  return launch;
+  kernelFunc.setType(LaunchBuilder.getFunctionType(newArgumentTypes, {}));
+  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
diff --git a/mlir/test/Dialect/GPU/outlining.mlir b/mlir/test/Dialect/GPU/outlining.mlir
index eeb09e4..07499a3 100644
--- a/mlir/test/Dialect/GPU/outlining.mlir
+++ b/mlir/test/Dialect/GPU/outlining.mlir
@@ -1,5 +1,6 @@
 // RUN: mlir-opt -gpu-kernel-outlining -split-input-file %s | FileCheck %s
 
+// CHECK-LABEL: func @launch()
 func @launch() {
   // CHECK: %[[ARG0:.*]] = "op"() : () -> f32
   %0 = "op"() : () -> (f32)
@@ -33,7 +34,8 @@ func @launch() {
   return
 }
 
-// CHECK: func @launch_kernel(%[[ARG0:.*]]: f32, %[[ARG1:.*]]: memref<?xf32, 1>)
+// CHECK-LABEL: func @launch_kernel
+// CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: memref<?xf32, 1>)
 // CHECK-NEXT: attributes {gpu.kernel}
 // CHECK-NEXT: %[[BID:.*]] = "gpu.block_id"() {dimension = "x"} : () -> index
 // CHECK-NEXT: = "gpu.block_id"() {dimension = "y"} : () -> index
-- 
2.7.4