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

[purnima shrivastava]

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

Preserve argument attributes when values captured by `gpu.launch` are converted into outlined `gpu.func` arguments. Previously, attributes such as `llvm.align` and `llvm.noalias` were dropped during `gpu-kernel-outlining`. This change propagates attributes from actual function arguments and adds a regression test.

Fixes #221441.


>From 84cf40cbc556862cc12abd659230913b594758c9 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        | 21 +++++++++++++++++++
 mlir/test/Dialect/GPU/outlining.mlir          | 15 +++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
index 0ee3d425f66be..697a8d34a00a0 100644
--- a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
@@ -19,6 +19,7 @@
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Utils/GPUUtils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/IRMapping.h"
@@ -199,6 +200,26 @@ 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}