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

[via flang-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: Zhen Wang (wangzpgi)

<details>
<summary>Changes</summary>

Add a `cuf.launch_config` attribute to the CUF dialect, together with a `getLaunchConfigAttrName()` helper for the attribute name.

The attribute records how a `!$cuf kernel do` directive spelled its launch configuration: 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 = ["*", "*"]>
```


---
Full diff: https://github.com/llvm/llvm-project/pull/229980.diff


4 Files Affected:

- (modified) flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h (+6) 
- (modified) flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.td (+23) 
- (modified) flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp (+1-1) 
- (added) flang/test/Fir/CUDA/cuda-launch-config-attr.mlir (+21) 


``````````diff
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
+}

``````````

</details>


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