[flang][cuda] Lower attribute for local variable (#81076)

This is a first simple patch to introduce a new FIR attribute to carry the CUDA variable attribute information to hlfir.declare and fir.declare operations. It currently lowers this information for local variables. The texture attribute is omitted since it is rejected by semantic and will not make its way to MLIR. This new attribute is added as optional attribute to the hlfir.declare and fir.declare operations.
2024-02-08 10:03:08 -08:00 · 2024-02-08 10:03:08 -08:00 · abc4f74df7
commit abc4f74df7
parent 758fd59d01
12 changed files with 121 additions and 21 deletions
--- a/flang/include/flang/Lower/ConvertVariable.h
+++ b/flang/include/flang/Lower/ConvertVariable.h
@ -137,6 +137,12 @@ translateSymbolAttributes(mlir::MLIRContext *mlirContext,
                          fir::FortranVariableFlagsEnum extraFlags =
                              fir::FortranVariableFlagsEnum::None);

+/// Translate the CUDA Fortran attributes of \p sym into the FIR CUDA attribute
+/// representation.
+fir::CUDAAttributeAttr
+translateSymbolCUDAAttribute(mlir::MLIRContext *mlirContext,
+                             const Fortran::semantics::Symbol &sym);
+
 /// Map a symbol to a given fir::ExtendedValue. This will generate an
 /// hlfir.declare when lowering to HLFIR and map the hlfir.declare result to the
 /// symbol.
--- a/flang/include/flang/Optimizer/Builder/HLFIRTools.h
+++ b/flang/include/flang/Optimizer/Builder/HLFIRTools.h
@ -233,11 +233,11 @@ translateToExtendedValue(mlir::Location loc, fir::FirOpBuilder &builder,
                         fir::FortranVariableOpInterface fortranVariable);

 /// Generate declaration for a fir::ExtendedValue in memory.
-fir::FortranVariableOpInterface genDeclare(mlir::Location loc,
-                                           fir::FirOpBuilder &builder,
-                                           const fir::ExtendedValue &exv,
-                                           llvm::StringRef name,
-                                           fir::FortranVariableFlagsAttr flags);
+fir::FortranVariableOpInterface
+genDeclare(mlir::Location loc, fir::FirOpBuilder &builder,
+           const fir::ExtendedValue &exv, llvm::StringRef name,
+           fir::FortranVariableFlagsAttr flags,
+           fir::CUDAAttributeAttr cudaAttr = {});

 /// Generate an hlfir.associate to build a variable from an expression value.
 /// The type of the variable must be provided so that scalar logicals are
--- a/flang/include/flang/Optimizer/Dialect/FIRAttr.td
+++ b/flang/include/flang/Optimizer/Dialect/FIRAttr.td
@ -55,7 +55,28 @@ def fir_FortranVariableFlagsAttr : fir_Attr<"FortranVariableFlags"> {
  let returnType = "::fir::FortranVariableFlagsEnum";
  let convertFromStorage = "$_self.getFlags()";
  let constBuilderCall =
-          "::fir::FortranVariableFlagsAttr::get($_builder.getContext(), $0)";
+        "::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_CUDAAttribute : I32EnumAttr<
+    "CUDAAttribute",
+    "CUDA Fortran variable attributes",
+    [CUDAconstant, CUDAdevice, CUDAmanaged, CUDApinned, CUDAshared,
+     CUDAunified]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::fir";
+}
+
+def fir_CUDAAttributeAttr : EnumAttr<fir_Dialect, fir_CUDAAttribute, "cuda"> {
+  let assemblyFormat = [{ ```<` $value `>` }];
 }

 def fir_BoxFieldAttr : I32EnumAttr<
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@ -3027,7 +3027,8 @@ def fir_DeclareOp : fir_Op<"declare", [AttrSizedOperandSegments,
    Optional<AnyShapeOrShiftType>:$shape,
    Variadic<AnyIntegerType>:$typeparams,
    Builtin_StringAttr:$uniq_name,
-    OptionalAttr<fir_FortranVariableFlagsAttr>:$fortran_attrs
+    OptionalAttr<fir_FortranVariableFlagsAttr>:$fortran_attrs,
+    OptionalAttr<fir_CUDAAttributeAttr>:$cuda_attr
  );

  let results = (outs AnyRefOrBox);
--- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
+++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
@ -88,7 +88,8 @@ def hlfir_DeclareOp : hlfir_Op<"declare", [AttrSizedOperandSegments,
    Optional<AnyShapeOrShiftType>:$shape,
    Variadic<AnyIntegerType>:$typeparams,
    Builtin_StringAttr:$uniq_name,
-    OptionalAttr<fir_FortranVariableFlagsAttr>:$fortran_attrs
+    OptionalAttr<fir_FortranVariableFlagsAttr>:$fortran_attrs,
+    OptionalAttr<fir_CUDAAttributeAttr>:$cuda_attr
  );

  let results = (outs AnyFortranVariable, AnyRefOrBoxLike);
@ -101,7 +102,8 @@ def hlfir_DeclareOp : hlfir_Op<"declare", [AttrSizedOperandSegments,
  let builders = [
    OpBuilder<(ins "mlir::Value":$memref, "llvm::StringRef":$uniq_name,
      CArg<"mlir::Value", "{}">:$shape, CArg<"mlir::ValueRange", "{}">:$typeparams,
-      CArg<"fir::FortranVariableFlagsAttr", "{}">:$fortran_attrs)>];
+      CArg<"fir::FortranVariableFlagsAttr", "{}">:$fortran_attrs,
+      CArg<"fir::CUDAAttributeAttr", "{}">:$cuda_attr)>];

  let extraClassDeclaration = [{
    /// Get the variable original base (same as input). It lacks
--- a/flang/lib/Lower/ConvertVariable.cpp
+++ b/flang/lib/Lower/ConvertVariable.cpp
@ -1579,6 +1579,38 @@ fir::FortranVariableFlagsAttr Fortran::lower::translateSymbolAttributes(
  return fir::FortranVariableFlagsAttr::get(mlirContext, flags);
 }

+fir::CUDAAttributeAttr Fortran::lower::translateSymbolCUDAAttribute(
+    mlir::MLIRContext *mlirContext, const Fortran::semantics::Symbol &sym) {
+  std::optional<Fortran::common::CUDADataAttr> cudaAttr =
+      Fortran::semantics::GetCUDADataAttr(&sym);
+  if (cudaAttr) {
+    fir::CUDAAttribute attr;
+    switch (*cudaAttr) {
+    case Fortran::common::CUDADataAttr::Constant:
+      attr = fir::CUDAAttribute::Constant;
+      break;
+    case Fortran::common::CUDADataAttr::Device:
+      attr = fir::CUDAAttribute::Device;
+      break;
+    case Fortran::common::CUDADataAttr::Managed:
+      attr = fir::CUDAAttribute::Managed;
+      break;
+    case Fortran::common::CUDADataAttr::Pinned:
+      attr = fir::CUDAAttribute::Pinned;
+      break;
+    case Fortran::common::CUDADataAttr::Shared:
+      attr = fir::CUDAAttribute::Shared;
+      break;
+    case Fortran::common::CUDADataAttr::Texture:
+      // Obsolete attribute
+      break;
+    }
+
+    return fir::CUDAAttributeAttr::get(mlirContext, attr);
+  }
+  return {};
+}
+
 /// Map a symbol to its FIR address and evaluated specification expressions.
 /// Not for symbols lowered to fir.box.
 /// Will optionally create fir.declare.
@ -1618,6 +1650,8 @@ static void genDeclareSymbol(Fortran::lower::AbstractConverter &converter,
    auto name = converter.mangleName(sym);
    fir::FortranVariableFlagsAttr attributes =
        Fortran::lower::translateSymbolAttributes(builder.getContext(), sym);
+    fir::CUDAAttributeAttr cudaAttr =
+        Fortran::lower::translateSymbolCUDAAttribute(builder.getContext(), sym);

    if (isCrayPointee) {
      mlir::Type baseType =
@ -1664,7 +1698,7 @@ static void genDeclareSymbol(Fortran::lower::AbstractConverter &converter,
      return;
    }
    auto newBase = builder.create<hlfir::DeclareOp>(
-        loc, base, name, shapeOrShift, lenParams, attributes);
+        loc, base, name, shapeOrShift, lenParams, attributes, cudaAttr);
    symMap.addVariableDefinition(sym, newBase, force);
    return;
  }
@ -1709,9 +1743,12 @@ void Fortran::lower::genDeclareSymbol(
    fir::FortranVariableFlagsAttr attributes =
        Fortran::lower::translateSymbolAttributes(
            builder.getContext(), sym.GetUltimate(), extraFlags);
+    fir::CUDAAttributeAttr cudaAttr =
+        Fortran::lower::translateSymbolCUDAAttribute(builder.getContext(),
+                                                     sym.GetUltimate());
    auto name = converter.mangleName(sym);
    hlfir::EntityWithAttributes declare =
-        hlfir::genDeclare(loc, builder, exv, name, attributes);
+        hlfir::genDeclare(loc, builder, exv, name, attributes, cudaAttr);
    symMap.addVariableDefinition(sym, declare.getIfVariableInterface(), force);
    return;
  }
--- a/flang/lib/Optimizer/Builder/HLFIRTools.cpp
+++ b/flang/lib/Optimizer/Builder/HLFIRTools.cpp
@ -198,7 +198,8 @@ mlir::Value hlfir::Entity::getFirBase() const {
 fir::FortranVariableOpInterface
 hlfir::genDeclare(mlir::Location loc, fir::FirOpBuilder &builder,
                  const fir::ExtendedValue &exv, llvm::StringRef name,
-                  fir::FortranVariableFlagsAttr flags) {
+                  fir::FortranVariableFlagsAttr flags,
+                  fir::CUDAAttributeAttr cudaAttr) {

  mlir::Value base = fir::getBase(exv);
  assert(fir::conformsWithPassByRef(base.getType()) &&
@ -228,7 +229,7 @@ hlfir::genDeclare(mlir::Location loc, fir::FirOpBuilder &builder,
      },
      [](const auto &) {});
  auto declareOp = builder.create<hlfir::DeclareOp>(
-      loc, base, name, shapeOrShift, lenParams, flags);
+      loc, base, name, shapeOrShift, lenParams, flags, cudaAttr);
  return mlir::cast<fir::FortranVariableOpInterface>(declareOp.getOperation());
 }

--- a/flang/lib/Optimizer/Dialect/FIRAttr.cpp
+++ b/flang/lib/Optimizer/Dialect/FIRAttr.cpp
@ -14,6 +14,7 @@
 #include "flang/Optimizer/Dialect/FIRDialect.h"
 #include "flang/Optimizer/Dialect/Support/KindMapping.h"
 #include "mlir/IR/AttributeSupport.h"
+#include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "llvm/ADT/SmallString.h"
@ -297,5 +298,5 @@ void fir::printFirAttribute(FIROpsDialect *dialect, mlir::Attribute attr,
 void FIROpsDialect::registerAttributes() {
  addAttributes<ClosedIntervalAttr, ExactTypeAttr, FortranVariableFlagsAttr,
                LowerBoundAttr, PointIntervalAttr, RealAttr, SubclassAttr,
-                UpperBoundAttr>();
+                UpperBoundAttr, CUDAAttributeAttr>();
 }
--- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
+++ b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
@ -123,14 +123,15 @@ void hlfir::DeclareOp::build(mlir::OpBuilder &builder,
                             mlir::OperationState &result, mlir::Value memref,
                             llvm::StringRef uniq_name, mlir::Value shape,
                             mlir::ValueRange typeparams,
-                             fir::FortranVariableFlagsAttr fortran_attrs) {
+                             fir::FortranVariableFlagsAttr fortran_attrs,
+                             fir::CUDAAttributeAttr cuda_attr) {
  auto nameAttr = builder.getStringAttr(uniq_name);
  mlir::Type inputType = memref.getType();
  bool hasExplicitLbs = hasExplicitLowerBounds(shape);
  mlir::Type hlfirVariableType =
      getHLFIRVariableType(inputType, hasExplicitLbs);
  build(builder, result, {hlfirVariableType, inputType}, memref, shape,
-        typeparams, nameAttr, fortran_attrs);
+        typeparams, nameAttr, fortran_attrs, cuda_attr);
 }

 mlir::LogicalResult hlfir::DeclareOp::verify() {
--- a/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
@ -320,12 +320,16 @@ public:
    mlir::Location loc = declareOp->getLoc();
    mlir::Value memref = declareOp.getMemref();
    fir::FortranVariableFlagsAttr fortranAttrs;
+    fir::CUDAAttributeAttr cudaAttr;
    if (auto attrs = declareOp.getFortranAttrs())
      fortranAttrs =
          fir::FortranVariableFlagsAttr::get(rewriter.getContext(), *attrs);
+    if (auto attr = declareOp.getCudaAttr())
+      cudaAttr = fir::CUDAAttributeAttr::get(rewriter.getContext(), *attr);
    auto firDeclareOp = rewriter.create<fir::DeclareOp>(
        loc, memref.getType(), memref, declareOp.getShape(),
-        declareOp.getTypeparams(), declareOp.getUniqName(), fortranAttrs);
+        declareOp.getTypeparams(), declareOp.getUniqName(), fortranAttrs,
+        cudaAttr);

    // Propagate other attributes from hlfir.declare to fir.declare.
    // OpenACC's acc.declare is one example. Right now, the propagation
--- a/flang/test/Lower/CUDA/cuda-data-attribute.cuf
+++ b/flang/test/Lower/CUDA/cuda-data-attribute.cuf
@ -0,0 +1,22 @@
+! 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 local variables.
+
+subroutine local_var_attrs
+  real, constant :: rc
+  real, device :: rd
+  real, allocatable, managed :: rm
+  real, allocatable, pinned :: rp
+end subroutine
+
+! CHECK-LABEL: func.func @_QPlocal_var_attrs()
+! CHECK: %{{.*}}:2 = hlfir.declare %{{.*}} {cuda_attr = #fir.cuda<constant>, uniq_name = "_QFlocal_var_attrsErc"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
+! CHECK: %{{.*}}:2 = hlfir.declare %{{.*}} {cuda_attr = #fir.cuda<device>, uniq_name = "_QFlocal_var_attrsErd"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
+! CHECK: %{{.*}}:2 = hlfir.declare %{{.*}} {cuda_attr = #fir.cuda<managed>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFlocal_var_attrsErm"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> (!fir.ref<!fir.box<!fir.heap<f32>>>, !fir.ref<!fir.box<!fir.heap<f32>>>)
+! CHECK: %{{.*}}:2 = hlfir.declare %{{.*}} {cuda_attr = #fir.cuda<pinned>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFlocal_var_attrsErp"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> (!fir.ref<!fir.box<!fir.heap<f32>>>, !fir.ref<!fir.box<!fir.heap<f32>>>)
+
+! FIR: %{{.*}} = fir.declare %{{.*}} {cuda_attr = #fir.cuda<constant>, uniq_name = "_QFlocal_var_attrsErc"} : (!fir.ref<f32>) -> !fir.ref<f32>
+! FIR: %{{.*}} = fir.declare %{{.*}} {cuda_attr = #fir.cuda<device>, uniq_name = "_QFlocal_var_attrsErd"} : (!fir.ref<f32>) -> !fir.ref<f32>
+! FIR: %{{.*}} = fir.declare %{{.*}} {cuda_attr = #fir.cuda<managed>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFlocal_var_attrsErm"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> !fir.ref<!fir.box<!fir.heap<f32>>>
+! FIR: %{{.*}} = fir.declare %{{.*}} {cuda_attr = #fir.cuda<pinned>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFlocal_var_attrsErp"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> !fir.ref<!fir.box<!fir.heap<f32>>>
--- a/flang/unittests/Optimizer/FortranVariableTest.cpp
+++ b/flang/unittests/Optimizer/FortranVariableTest.cpp
@ -49,7 +49,8 @@ TEST_F(FortranVariableTest, SimpleScalar) {
  auto name = mlir::StringAttr::get(&context, "x");
  auto declare = builder->create<fir::DeclareOp>(loc, addr.getType(), addr,
      /*shape=*/mlir::Value{}, /*typeParams=*/std::nullopt, name,
-      /*fortran_attrs=*/fir::FortranVariableFlagsAttr{});
+      /*fortran_attrs=*/fir::FortranVariableFlagsAttr{},
+      /*cuda_attr=*/fir::CUDAAttributeAttr{});

  fir::FortranVariableOpInterface fortranVariable = declare;
  EXPECT_FALSE(fortranVariable.isArray());
@ -74,7 +75,8 @@ TEST_F(FortranVariableTest, CharacterScalar) {
  auto name = mlir::StringAttr::get(&context, "x");
  auto declare = builder->create<fir::DeclareOp>(loc, addr.getType(), addr,
      /*shape=*/mlir::Value{}, typeParams, name,
-      /*fortran_attrs=*/fir::FortranVariableFlagsAttr{});
+      /*fortran_attrs=*/fir::FortranVariableFlagsAttr{},
+      /*cuda_attr=*/fir::CUDAAttributeAttr{});

  fir::FortranVariableOpInterface fortranVariable = declare;
  EXPECT_FALSE(fortranVariable.isArray());
@ -104,7 +106,8 @@ TEST_F(FortranVariableTest, SimpleArray) {
  auto name = mlir::StringAttr::get(&context, "x");
  auto declare = builder->create<fir::DeclareOp>(loc, addr.getType(), addr,
      shape, /*typeParams*/ std::nullopt, name,
-      /*fortran_attrs=*/fir::FortranVariableFlagsAttr{});
+      /*fortran_attrs=*/fir::FortranVariableFlagsAttr{},
+      /*cuda_attr=*/fir::CUDAAttributeAttr{});

  fir::FortranVariableOpInterface fortranVariable = declare;
  EXPECT_TRUE(fortranVariable.isArray());
@ -134,7 +137,8 @@ TEST_F(FortranVariableTest, CharacterArray) {
  auto name = mlir::StringAttr::get(&context, "x");
  auto declare = builder->create<fir::DeclareOp>(loc, addr.getType(), addr,
      shape, typeParams, name,
-      /*fortran_attrs=*/fir::FortranVariableFlagsAttr{});
+      /*fortran_attrs=*/fir::FortranVariableFlagsAttr{},
+      /*cuda_attr=*/fir::CUDAAttributeAttr{});

  fir::FortranVariableOpInterface fortranVariable = declare;
  EXPECT_TRUE(fortranVariable.isArray());