[flang-commits] [flang] [flang][cuda] Lower cluster_dims values (PR #81636)

[via flang-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

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

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

<details>
<summary>Changes</summary>

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
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>}

``````````

</details>


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