[Mlir-commits] [mlir] [mlir][gpu] Preserve argument attributes during kernel outlining (PR #221465)

[purnima shrivastava]

https://github.com/purnima-nlp updated https://github.com/llvm/llvm-project/pull/221465

>From 55b388d909ff3f205db17673675955c45f0017d6 Mon Sep 17 00:00:00 2001
From: Purnima Shrivastava <purnimashrivastava05 at .com>
Date: Sat, 5 Sep 2026 22:22:57 +0530
Subject: [PATCH] [mlir][gpu] Preserve argument attributes during kernel
 outlining

---
 .../GPU/Transforms/KernelOutlining.cpp        | 20 +++++++++++++++++++
 mlir/test/Dialect/GPU/outlining.mlir          | 15 ++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
index 0ee3d425f66be..e5b402bb7d332 100644
--- a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
@@ -24,6 +24,7 @@
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/SymbolTable.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include <limits>
@@ -199,6 +200,25 @@ static gpu::GPUFuncOp outlineKernelFuncImpl(gpu::LaunchOp launchOp,
       builder, loc, kernelFnName, type,
       TypeRange(ValueRange(launchOp.getWorkgroupAttributionBBArgs())),
       TypeRange(ValueRange(launchOp.getPrivateAttributions())));
+
+  for (auto [index, operand] : llvm::enumerate(operands)) {
+    auto blockArg = dyn_cast<BlockArgument>(operand);
+    if (!blockArg)
+      continue;
+
+    auto function = dyn_cast_or_null<FunctionOpInterface>(
+        blockArg.getOwner()->getParentOp());
+    if (!function || function.empty())
+      continue;
+
+    // Only propagate attributes from actual function arguments.
+    if (blockArg.getOwner() != &function.front())
+      continue;
+
+    if (DictionaryAttr attrs = function.getArgAttrDict(blockArg.getArgNumber()))
+      outlinedFunc.setArgAttrs(index, attrs);
+  }
+
   outlinedFunc.setKernel(true);
 
   // If we can infer bounds on the grid and/or block sizes from the arguments
diff --git a/mlir/test/Dialect/GPU/outlining.mlir b/mlir/test/Dialect/GPU/outlining.mlir
index f708561bc2f01..9fac049f8f8ef 100644
--- a/mlir/test/Dialect/GPU/outlining.mlir
+++ b/mlir/test/Dialect/GPU/outlining.mlir
@@ -719,3 +719,18 @@ func.func @launch_cooperative() {
   }
   return
 }
+
+// -----
+
+llvm.func @kernel_argument(%arg0: !llvm.ptr {llvm.align = 16 : i64, llvm.noalias}) {
+  %c1 = arith.constant 1 : index
+  gpu.launch blocks(%arg1, %arg2, %arg3) in (%arg7 = %c1, %arg8 = %c1, %arg9 = %c1)
+             threads(%arg4, %arg5, %arg6) in (%arg10 = %c1, %arg11 = %c1, %arg12 = %c1) {
+    "test.consume"(%arg0) : (!llvm.ptr) -> ()
+    gpu.terminator
+  }
+  llvm.return
+}
+
+// CHECK-LABEL: gpu.func @kernel_argument_kernel(
+// CHECK-SAME: %{{.*}}: !llvm.ptr {llvm.align = 16 : i64, llvm.noalias}