Addressing some late review comments on kernel inlining.

Just formatting and better lit tests, no functional change.

PiperOrigin-RevId: 267942907

1 files changed, 15 insertions, 15 deletions

diff --git a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
index f464b091fc8..26449f6e6f1 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