[flang][cuda] Lower attribute for procedure

[clementval]

This PR adds a new attribute to represent the CUDA attribute attached to procedure. This attribute is attached to the func.func operation during lowering.
Other procedures information such as launch_bounds and cluster_dims will be added separately.

[llvmbot]

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

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

Changes

This PR adds a new attribute to represent the CUDA attribute attached to procedure. This attribute is attached to the func.func operation during lowering.
Other procedures information such as launch_bounds and cluster_dims will be added separately.

---

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

5 Files Affected:

• (modified) flang/include/flang/Optimizer/Dialect/FIRAttr.td (+37-17)
• (modified) flang/include/flang/Optimizer/Support/Utils.h (+27)
• (modified) flang/lib/Lower/CallInterface.cpp (+7)
• (modified) flang/lib/Optimizer/Dialect/FIRAttr.cpp (+1-1)
• (added) flang/test/Lower/CUDA/cuda-proc-attribute.cuf (+29)

diff --git a/flang/include/flang/Optimizer/Dialect/FIRAttr.td b/flang/include/flang/Optimizer/Dialect/FIRAttr.td
index 422ad531181ae3..00e293e2f04278 100644
--- a/flang/include/flang/Optimizer/Dialect/FIRAttr.td
+++ b/flang/include/flang/Optimizer/Dialect/FIRAttr.td
@@ -58,19 +58,34 @@ def fir_FortranVariableFlagsAttr : fir_Attr<"FortranVariableFlags"> {
         "::fir::FortranVariableFlagsAttr::get($_builder.getContext(), $0)";
 }
 
-def CUDAconstant : I32EnumAttrCase<"Constant", 0, "constant">;
-def CUDAdevice   : I32EnumAttrCase<"Device", 1, "device">;
-def CUDAmanaged  : I32EnumAttrCase<"Managed", 2, "managed">;
-def CUDApinned   : I32EnumAttrCase<"Pinned", 3, "pinned">;
-def CUDAshared   : I32EnumAttrCase<"Shared", 4, "shared">;
-def CUDAunified  : I32EnumAttrCase<"Unified", 5, "unified">;
-// Texture is omitted since it is obsolete and rejected by semantic.
+def fir_BoxFieldAttr : I32EnumAttr<
+    "BoxFieldAttr", "",
+    [
+      I32EnumAttrCase<"base_addr", 0>,
+      I32EnumAttrCase<"derived_type", 1>
+    ]> {
+  let cppNamespace = "fir";
+}
+
+// mlir::SideEffects::Resource for modelling operations which add debugging information
+def DebuggingResource : Resource<"::fir::DebuggingResource">;
+
+//===----------------------------------------------------------------------===//
+// CUDA Fortran specific attributes
+//===----------------------------------------------------------------------===//
 
 def fir_CUDADataAttribute : I32EnumAttr<
     "CUDADataAttribute",
     "CUDA Fortran variable attributes",
-    [CUDAconstant, CUDAdevice, CUDAmanaged, CUDApinned, CUDAshared,
-     CUDAunified]> {
+    [
+      I32EnumAttrCase<"Constant", 0, "constant">,
+      I32EnumAttrCase<"Device", 1, "device">,
+      I32EnumAttrCase<"Managed", 2, "managed">,
+      I32EnumAttrCase<"Pinned", 3, "pinned">,
+      I32EnumAttrCase<"Shared", 4, "shared">,
+      I32EnumAttrCase<"Unified", 5, "unified">,
+      // Texture is omitted since it is obsolete and rejected by semantic.
+    ]> {
   let genSpecializedAttr = 0;
   let cppNamespace = "::fir";
 }
@@ -80,17 +95,22 @@ def fir_CUDADataAttributeAttr :
   let assemblyFormat = [{ ```<` $value `>` }];
 }
 
-def fir_BoxFieldAttr : I32EnumAttr<
-    "BoxFieldAttr", "",
+def fir_CUDAProcAttribute : I32EnumAttr<
+    "CUDAProcAttribute", "CUDA Fortran procedure attributes",
     [
-      I32EnumAttrCase<"base_addr", 0>,
-      I32EnumAttrCase<"derived_type", 1>
+      I32EnumAttrCase<"Host", 0, "host">,
+      I32EnumAttrCase<"Device", 1, "device">,
+      I32EnumAttrCase<"HostDevice", 2, "host_device">,
+      I32EnumAttrCase<"Global", 3, "global">,
+      I32EnumAttrCase<"GridGlobal", 4, "grid_global">,
     ]> {
-  let cppNamespace = "fir";
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::fir";
 }
 
-
-// mlir::SideEffects::Resource for modelling operations which add debugging information
-def DebuggingResource : Resource<"::fir::DebuggingResource">;
+def fir_CUDAProcAttributeAttr :
+    EnumAttr<fir_Dialect, fir_CUDAProcAttribute, "cuda_proc"> {
+  let assemblyFormat = [{ ```<` $value `>` }];
+}
 
 #endif // FIR_DIALECT_FIR_ATTRS
diff --git a/flang/include/flang/Optimizer/Support/Utils.h b/flang/include/flang/Optimizer/Support/Utils.h
index 84c550afc0777a..4e06bf8bafdc83 100644
--- a/flang/include/flang/Optimizer/Support/Utils.h
+++ b/flang/include/flang/Optimizer/Support/Utils.h
@@ -303,6 +303,33 @@ getCUDADataAttribute(mlir::MLIRContext *mlirContext,
   return {};
 }
 
+inline fir::CUDAProcAttributeAttr getCUDAProcAttribute(
+    mlir::MLIRContext *mlirContext,
+    std::optional<Fortran::common::CUDASubprogramAttrs> cudaAttr) {
+  if (cudaAttr) {
+    fir::CUDAProcAttribute attr;
+    switch (*cudaAttr) {
+    case Fortran::common::CUDASubprogramAttrs::Host:
+      attr = fir::CUDAProcAttribute::Host;
+      break;
+    case Fortran::common::CUDASubprogramAttrs::Device:
+      attr = fir::CUDAProcAttribute::Device;
+      break;
+    case Fortran::common::CUDASubprogramAttrs::HostDevice:
+      attr = fir::CUDAProcAttribute::HostDevice;
+      break;
+    case Fortran::common::CUDASubprogramAttrs::Global:
+      attr = fir::CUDAProcAttribute::Global;
+      break;
+    case Fortran::common::CUDASubprogramAttrs::Grid_Global:
+      attr = fir::CUDAProcAttribute::GridGlobal;
+      break;
+    }
+    return fir::CUDAProcAttributeAttr::get(mlirContext, attr);
+  }
+  return {};
+}
+
 } // namespace fir
 
 #endif // FORTRAN_OPTIMIZER_SUPPORT_UTILS_H
diff --git a/flang/lib/Lower/CallInterface.cpp b/flang/lib/Lower/CallInterface.cpp
index 9c32b7100db253..41597c1b15386e 100644
--- a/flang/lib/Lower/CallInterface.cpp
+++ b/flang/lib/Lower/CallInterface.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "flang/Lower/CallInterface.h"
+#include "flang/Common/Fortran.h"
 #include "flang/Evaluate/fold.h"
 #include "flang/Lower/Bridge.h"
 #include "flang/Lower/Mangler.h"
@@ -559,6 +560,12 @@ void Fortran::lower::CallInterface<T>::declare() {
           func.setArgAttrs(placeHolder.index(), placeHolder.value().attributes);
       side().setFuncAttrs(func);
     }
+    if (characteristic && characteristic->cudaSubprogramAttrs) {
+      func.getOperation()->setAttr(
+          fir::getCUDAAttrName(),
+          fir::getCUDAProcAttribute(func.getContext(),
+                                    *characteristic->cudaSubprogramAttrs));
+    }
   }
 }
 
diff --git a/flang/lib/Optimizer/Dialect/FIRAttr.cpp b/flang/lib/Optimizer/Dialect/FIRAttr.cpp
index 218fa5028e748d..8df7a6c5cfc5d5 100644
--- a/flang/lib/Optimizer/Dialect/FIRAttr.cpp
+++ b/flang/lib/Optimizer/Dialect/FIRAttr.cpp
@@ -298,5 +298,5 @@ void fir::printFirAttribute(FIROpsDialect *dialect, mlir::Attribute attr,
 void FIROpsDialect::registerAttributes() {
   addAttributes<ClosedIntervalAttr, ExactTypeAttr, FortranVariableFlagsAttr,
                 LowerBoundAttr, PointIntervalAttr, RealAttr, SubclassAttr,
-                UpperBoundAttr, CUDADataAttributeAttr>();
+                UpperBoundAttr, CUDADataAttributeAttr, CUDAProcAttributeAttr>();
 }
diff --git a/flang/test/Lower/CUDA/cuda-proc-attribute.cuf b/flang/test/Lower/CUDA/cuda-proc-attribute.cuf
new file mode 100644
index 00000000000000..8991f8dac9f48e
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-proc-attribute.cuf
@@ -0,0 +1,29 @@
+! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
+! RUN: bbc -emit-hlfir -fcuda %s -o - | fir-opt -convert-hlfir-to-fir | FileCheck %s --check-prefix=FIR
+
+! Test lowering of CUDA attribute on procedures.
+
+attributes(host) subroutine sub_host(); end
+! CHECK: func.func @_QPsub_host() attributes {fir.cuda_attr = #fir.cuda_proc<host>}
+
+attributes(device) subroutine sub_device(); end
+! CHECK: func.func @_QPsub_device() attributes {fir.cuda_attr = #fir.cuda_proc<device>}
+
+attributes(host) attributes(device) subroutine sub_host_device; end
+! CHECK: func.func @_QPsub_host_device() attributes {fir.cuda_attr = #fir.cuda_proc<host_device>}
+
+attributes(global) subroutine sub_global(); end
+! CHECK: func.func @_QPsub_global() attributes {fir.cuda_attr = #fir.cuda_proc<global>}
+
+attributes(grid_global) subroutine sub_grid_global(); end
+! CHECK: func.func @_QPsub_grid_global() attributes {fir.cuda_attr = #fir.cuda_proc<grid_global>}
+
+attributes(host) integer function fct_host(); end
+! CHECK: func.func @_QPfct_host() -> i32 attributes {fir.cuda_attr = #fir.cuda_proc<host>}
+
+attributes(device) integer function fct_device(); end
+! CHECK: func.func @_QPfct_device() -> i32 attributes {fir.cuda_attr = #fir.cuda_proc<device>}
+
+attributes(host) attributes(device) integer function fct_host_device; end
+! CHECK: func.func @_QPfct_host_device() -> i32 attributes {fir.cuda_attr = #fir.cuda_proc<host_device>}
+

[vzakhari]

Looks good!