[flang][cuda] Lower cluster_dims values

[clementval]

This PR adds a new attribute to carry over the information from cluster_dims. The new attribute CUDAClusterDimsAttr holds 3 integer attributes and is added to func.func operation.

[llvmbot]

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

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

This PR adds a new attribute to carry over the information from cluster_dims. The new attribute CUDAClusterDimsAttr holds 3 integer attributes and is added to func.func operation.

---

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

5 Files Affected:

• (modified) flang/include/flang/Optimizer/Dialect/FIRAttr.td (+12)
• (modified) flang/include/flang/Optimizer/Dialect/FIROpsSupport.h (+5)
• (modified) flang/lib/Lower/CallInterface.cpp (+11-1)
• (modified) flang/lib/Optimizer/Dialect/FIRAttr.cpp (+1-1)
• (modified) flang/test/Lower/CUDA/cuda-proc-attribute.cuf (+3)

diff --git a/flang/include/flang/Optimizer/Dialect/FIRAttr.td b/flang/include/flang/Optimizer/Dialect/FIRAttr.td
index 3602c67de1412a..66d6cd471116b0 100644
--- a/flang/include/flang/Optimizer/Dialect/FIRAttr.td
+++ b/flang/include/flang/Optimizer/Dialect/FIRAttr.td
@@ -125,4 +125,16 @@ def fir_CUDALaunchBoundsAttr : fir_Attr<"CUDALaunchBounds"> {
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
+def fir_CUDAClusterDimsAttr : fir_Attr<"CUDAClusterDims"> {
+  let mnemonic = "cluster_dims";
+
+  let parameters = (ins
+    "mlir::IntegerAttr":$x,
+    "mlir::IntegerAttr":$y,
+    "mlir::IntegerAttr":$z
+  );
+
+  let assemblyFormat = "`<` struct(params) `>`";
+}
+
 #endif // FIR_DIALECT_FIR_ATTRS
diff --git a/flang/include/flang/Optimizer/Dialect/FIROpsSupport.h b/flang/include/flang/Optimizer/Dialect/FIROpsSupport.h
index 29fa57cd7a0d8a..e8226b6df58ca2 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROpsSupport.h
+++ b/flang/include/flang/Optimizer/Dialect/FIROpsSupport.h
@@ -80,6 +80,11 @@ static constexpr llvm::StringRef getCUDALaunchBoundsAttrName() {
   return "fir.cuda_launch_bounds";
 }
 
+/// Attribute to carry CUDA cluster_dims values.
+static constexpr llvm::StringRef getCUDAClusterDimsAttrName() {
+  return "fir.cuda_cluster_dims";
+}
+
 /// Attribute to mark that a function argument is a character dummy procedure.
 /// Character dummy procedure have special ABI constraints.
 static constexpr llvm::StringRef getCharacterProcedureDummyAttrName() {
diff --git a/flang/lib/Lower/CallInterface.cpp b/flang/lib/Lower/CallInterface.cpp
index f990e0b7ce4dcf..6c71e884307d7b 100644
--- a/flang/lib/Lower/CallInterface.cpp
+++ b/flang/lib/Lower/CallInterface.cpp
@@ -540,10 +540,10 @@ setCUDAAttributes(mlir::func::FuncOp func,
     if (auto details =
             sym->GetUltimate()
                 .detailsIf<Fortran::semantics::SubprogramDetails>()) {
+      mlir::Type i64Ty = mlir::IntegerType::get(func.getContext(), 64);
       if (!details->cudaLaunchBounds().empty()) {
         assert(details->cudaLaunchBounds().size() >= 2 &&
                "expect at least 2 values");
-        mlir::Type i64Ty = mlir::IntegerType::get(func.getContext(), 64);
         auto maxTPBAttr =
             mlir::IntegerAttr::get(i64Ty, details->cudaLaunchBounds()[0]);
         auto minBPMAttr =
@@ -557,6 +557,16 @@ setCUDAAttributes(mlir::func::FuncOp func,
             fir::CUDALaunchBoundsAttr::get(func.getContext(), maxTPBAttr,
                                            minBPMAttr, ubAttr));
       }
+
+      if (!details->cudaClusterDims().empty()) {
+        assert(details->cudaClusterDims().size() == 3 && "expect 3 values");
+        auto xAttr = mlir::IntegerAttr::get(i64Ty, details->cudaClusterDims()[0]);
+        auto yAttr = mlir::IntegerAttr::get(i64Ty, details->cudaClusterDims()[1]);
+        auto zAttr = mlir::IntegerAttr::get(i64Ty, details->cudaClusterDims()[2]);
+        func.getOperation()->setAttr(
+            fir::getCUDAClusterDimsAttrName(),
+            fir::CUDAClusterDimsAttr::get(func.getContext(), xAttr, yAttr, zAttr));
+      }
     }
   }
 }
diff --git a/flang/lib/Optimizer/Dialect/FIRAttr.cpp b/flang/lib/Optimizer/Dialect/FIRAttr.cpp
index 8d780e03dcbe73..0cf8dfb9f784c3 100644
--- a/flang/lib/Optimizer/Dialect/FIRAttr.cpp
+++ b/flang/lib/Optimizer/Dialect/FIRAttr.cpp
@@ -299,5 +299,5 @@ void FIROpsDialect::registerAttributes() {
   addAttributes<ClosedIntervalAttr, ExactTypeAttr, FortranVariableFlagsAttr,
                 LowerBoundAttr, PointIntervalAttr, RealAttr, SubclassAttr,
                 UpperBoundAttr, CUDADataAttributeAttr, CUDAProcAttributeAttr,
-                CUDALaunchBoundsAttr>();
+                CUDALaunchBoundsAttr, CUDAClusterDimsAttr>();
 }
diff --git a/flang/test/Lower/CUDA/cuda-proc-attribute.cuf b/flang/test/Lower/CUDA/cuda-proc-attribute.cuf
index 9eb2b85aaf0b83..d9765f6cd2fe8c 100644
--- a/flang/test/Lower/CUDA/cuda-proc-attribute.cuf
+++ b/flang/test/Lower/CUDA/cuda-proc-attribute.cuf
@@ -38,3 +38,6 @@ attributes(global) launch_bounds(1, 2) subroutine sub_lbounds1(); end
 
 attributes(global) launch_bounds(1, 2, 3) subroutine sub_lbounds2(); end
 ! CHECK: func.func @_QPsub_lbounds2() attributes {fir.cuda_attr = #fir.cuda_proc<global>, fir.cuda_launch_bounds = #fir.launch_bounds<maxTPB = 1 : i64, minBPM = 2 : i64, upperBoundClusterSize = 3 : i64>}
+
+attributes(global) cluster_dims(1, 2, 3) subroutine sub_clusterdims1(); end
+! CHECK: func.func @_QPsub_clusterdims1() attributes {fir.cuda_attr = #fir.cuda_proc<global>, fir.cuda_cluster_dims = #fir.cluster_dims<x = 1 : i64, y = 2 : i64, z = 3 : i64>}

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[vzakhari]

Looks good!