[flang-commits] [flang] [flang][cuda] Add cuf.launch_config attribute recording a kernel directive's launch configuration (PR #229980)

[Zhen Wang via flang-commits]

https://github.com/wangzpgi updated https://github.com/llvm/llvm-project/pull/229980

>From d6233474a432634934cbeaf74eae55d8b77dbcf3 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Wed, 7 Oct 2026 15:16:16 -0700
Subject: [PATCH 1/2] Add cuf.launch_config attribute recording a kernel
 directive's launch configuration

---
 .../Dialect/CUF/Attributes/CUFAttr.h          |  6 +++++
 .../Dialect/CUF/Attributes/CUFAttr.td         | 23 +++++++++++++++++++
 .../Dialect/CUF/Attributes/CUFAttr.cpp        |  2 +-
 .../Fir/CUDA/cuda-launch-config-attr.mlir     | 21 +++++++++++++++++
 4 files changed, 51 insertions(+), 1 deletion(-)
 create mode 100644 flang/test/Fir/CUDA/cuda-launch-config-attr.mlir

diff --git a/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h b/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
index 4082abd2431c2..8cf372c808576 100644
--- a/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
+++ b/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
@@ -64,6 +64,12 @@ static constexpr llvm::StringRef getClusterDimsAttrName() {
   return "cuf.cluster_dims";
 }
 
+/// Attribute to carry a kernel directive's launch configuration as written
+/// (a cuf::LaunchConfigAttr).
+static constexpr llvm::StringRef getLaunchConfigAttrName() {
+  return "cuf.launch_config";
+}
+
 inline cuf::DataAttributeAttr
 getDataAttribute(mlir::MLIRContext *mlirContext,
                  std::optional<Fortran::common::CUDADataAttr> cudaAttr) {
diff --git a/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.td b/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.td
index 84fcecc63433a..bc56fe1719049 100644
--- a/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.td
+++ b/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.td
@@ -77,6 +77,29 @@ def cuf_ClusterDimsAttr : cuf_Attr<"ClusterDims"> {
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
+def cuf_LaunchConfigAttr : cuf_Attr<"LaunchConfig"> {
+  let mnemonic = "launch_config";
+  let summary = "Launch configuration of a kernel directive as written";
+  let description = [{
+    Records how a `cuf.kernel` directive spelled its launch configuration so
+    that later stages can report it: the number of dimensions the directive
+    names, and each grid dimension as written. A grid dimension left to the
+    compiler is spelled `*`, a constant by its value, and a variable by its
+    name.
+
+    ```mlir
+    #cuf.launch_config<dims = 2 : i64, grid = ["*", "*"]>
+    ```
+  }];
+
+  let parameters = (ins
+    "mlir::IntegerAttr":$dims,
+    "mlir::ArrayAttr":$grid
+  );
+
+  let assemblyFormat = "`<` struct(params) `>`";
+}
+
 def cuf_DataTransferKind : I32EnumAttr<
     "DataTransferKind", "CUDA Fortran data transfer kind",
     [
diff --git a/flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp b/flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp
index a8860eda73626..80baf92c00f2c 100644
--- a/flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp
+++ b/flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp
@@ -26,7 +26,7 @@ namespace cuf {
 
 void CUFDialect::registerAttributes() {
   addAttributes<ClusterDimsAttr, DataAttributeAttr, DataTransferKindAttr,
-                LaunchBoundsAttr, ProcAttributeAttr>();
+                LaunchBoundsAttr, LaunchConfigAttr, ProcAttributeAttr>();
 }
 
 cuf::DataAttributeAttr getDataAttr(mlir::Operation *op) {
diff --git a/flang/test/Fir/CUDA/cuda-launch-config-attr.mlir b/flang/test/Fir/CUDA/cuda-launch-config-attr.mlir
new file mode 100644
index 0000000000000..7be4706972dbd
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-launch-config-attr.mlir
@@ -0,0 +1,21 @@
+// RUN: fir-opt %s | fir-opt | FileCheck %s
+
+// Round-trip the launch configuration attribute a kernel directive records.
+
+// CHECK-LABEL: func.func private @one_dim()
+// CHECK-SAME: cuf.launch_config = #cuf.launch_config<dims = 1 : i64, grid = ["*"]>
+func.func private @one_dim() attributes {cuf.launch_config = #cuf.launch_config<dims = 1 : i64, grid = ["*"]>} {
+  return
+}
+
+// CHECK-LABEL: func.func private @two_dims_constant_grid()
+// CHECK-SAME: cuf.launch_config = #cuf.launch_config<dims = 2 : i64, grid = ["64", "*"]>
+func.func private @two_dims_constant_grid() attributes {cuf.launch_config = #cuf.launch_config<dims = 2 : i64, grid = ["64", "*"]>} {
+  return
+}
+
+// CHECK-LABEL: func.func private @three_dims_variable_grid()
+// CHECK-SAME: cuf.launch_config = #cuf.launch_config<dims = 3 : i64, grid = ["ng", "*", "*"]>
+func.func private @three_dims_variable_grid() attributes {cuf.launch_config = #cuf.launch_config<dims = 3 : i64, grid = ["ng", "*", "*"]>} {
+  return
+}

>From 7844e0de4deef1589b112be8cd5d6e8ef605b8ec Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Wed, 7 Oct 2026 21:13:43 -0700
Subject: [PATCH 2/2] Use LaunchConfigAttr::name instead of a separate
 attribute name helper

---
 .../flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h        | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h b/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
index 8cf372c808576..4082abd2431c2 100644
--- a/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
+++ b/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
@@ -64,12 +64,6 @@ static constexpr llvm::StringRef getClusterDimsAttrName() {
   return "cuf.cluster_dims";
 }
 
-/// Attribute to carry a kernel directive's launch configuration as written
-/// (a cuf::LaunchConfigAttr).
-static constexpr llvm::StringRef getLaunchConfigAttrName() {
-  return "cuf.launch_config";
-}
-
 inline cuf::DataAttributeAttr
 getDataAttribute(mlir::MLIRContext *mlirContext,
                  std::optional<Fortran::common::CUDADataAttr> cudaAttr) {