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

[via flang-commits]

Author: Valentin Clement (バレンタイン クレメン)
Date: 2024-02-13T10:02:52-08:00
New Revision: 5e3c7e3aa48356a62a4b70d5d9d3e4ddd055a390

URL: https://github.com/llvm/llvm-project/commit/5e3c7e3aa48356a62a4b70d5d9d3e4ddd055a390
DIFF: https://github.com/llvm/llvm-project/commit/5e3c7e3aa48356a62a4b70d5d9d3e4ddd055a390.diff

LOG: [flang][cuda] Lower cluster_dims values (#81636)

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.

Added: 
    

Modified: 
    flang/include/flang/Optimizer/Dialect/FIRAttr.td
    flang/include/flang/Optimizer/Dialect/FIROpsSupport.h
    flang/lib/Lower/CallInterface.cpp
    flang/lib/Optimizer/Dialect/FIRAttr.cpp
    flang/test/Lower/CUDA/cuda-proc-attribute.cuf

Removed: 
    


################################################################################
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..6b71aabf7fdc89 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,20 @@ 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>}