Reland "[flang][cuda] Set the allocator of derived type component after allocation"

flang/include/flang/Lower/Cuda.h → flang/include/flang/Lower/CUDA.h

@@ -1,4 +1,4 @@
-//===-- Lower/Cuda.h -- Cuda Fortran utilities ------------------*- C++ -*-===//
+//===-- Lower/CUDA.h -- CUDA Fortran utilities ------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -14,13 +14,23 @@
 #define FORTRAN_LOWER_CUDA_H
 
 #include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Builder/MutableBox.h"
 #include "flang/Optimizer/Dialect/CUF/CUFOps.h"
+#include "flang/Runtime/allocator-registry-consts.h"
 #include "flang/Semantics/tools.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 
+namespace mlir {
+class Value;
+class Location;
+class MLIRContext;
+} // namespace mlir
+
 namespace Fortran::lower {
 
+class AbstractConverter;
+
 static inline unsigned getAllocatorIdx(const Fortran::semantics::Symbol &sym) {
   std::optional<Fortran::common::CUDADataAttr> cudaAttr =
       Fortran::semantics::GetCUDADataAttr(&sym.GetUltimate());
@@ -37,6 +47,21 @@ static inline unsigned getAllocatorIdx(const Fortran::semantics::Symbol &sym) {
   return kDefaultAllocator;
 }
 
+void initializeDeviceComponentAllocator(
+    Fortran::lower::AbstractConverter &converter,
+    const Fortran::semantics::Symbol &sym, const fir::MutableBoxValue &box);
+
+mlir::Type gatherDeviceComponentCoordinatesAndType(
+    fir::FirOpBuilder &builder, mlir::Location loc,
+    const Fortran::semantics::Symbol &sym, fir::RecordType recTy,
+    llvm::SmallVector<mlir::Value> &coordinates);
+
+/// Translate the CUDA Fortran attributes of \p sym into the FIR CUDA attribute
+/// representation.
+cuf::DataAttributeAttr
+translateSymbolCUFDataAttribute(mlir::MLIRContext *mlirContext,
+                                const Fortran::semantics::Symbol &sym);
+
 } // end namespace Fortran::lower
 
 #endif // FORTRAN_LOWER_CUDA_H

flang/include/flang/Lower/ConvertVariable.h

@@ -162,12 +162,6 @@ translateSymbolAttributes(mlir::MLIRContext *mlirContext,
                           fir::FortranVariableFlagsEnum extraFlags =
                               fir::FortranVariableFlagsEnum::None);
 
-/// Translate the CUDA Fortran attributes of \p sym into the FIR CUDA attribute
-/// representation.
-cuf::DataAttributeAttr
-translateSymbolCUFDataAttribute(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.

flang/lib/Lower/Allocatable.cpp

@@ -13,9 +13,9 @@
 #include "flang/Lower/Allocatable.h"
 #include "flang/Evaluate/tools.h"
 #include "flang/Lower/AbstractConverter.h"
+#include "flang/Lower/CUDA.h"
 #include "flang/Lower/ConvertType.h"
 #include "flang/Lower/ConvertVariable.h"
-#include "flang/Lower/Cuda.h"
 #include "flang/Lower/IterationSpace.h"
 #include "flang/Lower/Mangler.h"
 #include "flang/Lower/OpenACC.h"
@@ -445,10 +445,14 @@ class AllocateStmtHelper {
                                        /*mustBeHeap=*/true);
   }
 
-  void postAllocationAction(const Allocation &alloc) {
+  void postAllocationAction(const Allocation &alloc,
+                            const fir::MutableBoxValue &box) {
     if (alloc.getSymbol().test(Fortran::semantics::Symbol::Flag::AccDeclare))
       Fortran::lower::attachDeclarePostAllocAction(converter, builder,
                                                    alloc.getSymbol());
+    if (Fortran::semantics::HasCUDAComponent(alloc.getSymbol()))
+      Fortran::lower::initializeDeviceComponentAllocator(
+          converter, alloc.getSymbol(), box);
   }
 
   void setPinnedToFalse() {
@@ -481,7 +485,7 @@ class AllocateStmtHelper {
       // Pointers must use PointerAllocate so that their deallocations
       // can be validated.
       genInlinedAllocation(alloc, box);
-      postAllocationAction(alloc);
+      postAllocationAction(alloc, box);
       setPinnedToFalse();
       return;
     }
@@ -504,7 +508,7 @@ class AllocateStmtHelper {
           genCudaAllocate(builder, loc, box, errorManager, alloc.getSymbol());
     }
     fir::factory::syncMutableBoxFromIRBox(builder, loc, box);
-    postAllocationAction(alloc);
+    postAllocationAction(alloc, box);
     errorManager.assignStat(builder, loc, stat);
   }
 
@@ -647,7 +651,7 @@ class AllocateStmtHelper {
       setPinnedToFalse();
     }
     fir::factory::syncMutableBoxFromIRBox(builder, loc, box);
-    postAllocationAction(alloc);
+    postAllocationAction(alloc, box);
     errorManager.assignStat(builder, loc, stat);
   }
 

flang/lib/Lower/Bridge.cpp

@@ -13,14 +13,14 @@
 #include "flang/Lower/Bridge.h"
 
 #include "flang/Lower/Allocatable.h"
+#include "flang/Lower/CUDA.h"
 #include "flang/Lower/CallInterface.h"
 #include "flang/Lower/Coarray.h"
 #include "flang/Lower/ConvertCall.h"
 #include "flang/Lower/ConvertExpr.h"
 #include "flang/Lower/ConvertExprToHLFIR.h"
 #include "flang/Lower/ConvertType.h"
 #include "flang/Lower/ConvertVariable.h"
-#include "flang/Lower/Cuda.h"
 #include "flang/Lower/DirectivesCommon.h"
 #include "flang/Lower/HostAssociations.h"
 #include "flang/Lower/IO.h"

flang/lib/Lower/CMakeLists.txt

@@ -15,6 +15,7 @@ add_flang_library(FortranLower
   ConvertProcedureDesignator.cpp
   ConvertType.cpp
   ConvertVariable.cpp
+  CUDA.cpp
   CustomIntrinsicCall.cpp
   HlfirIntrinsics.cpp
   HostAssociations.cpp

flang/lib/Lower/CUDA.cpp

@@ -0,0 +1,155 @@
+//===-- CUDA.cpp -- CUDA Fortran specific lowering ------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Lower/CUDA.h"
+
+#define DEBUG_TYPE "flang-lower-cuda"
+
+void Fortran::lower::initializeDeviceComponentAllocator(
+    Fortran::lower::AbstractConverter &converter,
+    const Fortran::semantics::Symbol &sym, const fir::MutableBoxValue &box) {
+  if (const auto *details{
+          sym.GetUltimate()
+              .detailsIf<Fortran::semantics::ObjectEntityDetails>()}) {
+    const Fortran::semantics::DeclTypeSpec *type{details->type()};
+    const Fortran::semantics::DerivedTypeSpec *derived{type ? type->AsDerived()
+                                                            : nullptr};
+    if (derived) {
+      if (!FindCUDADeviceAllocatableUltimateComponent(*derived))
+        return; // No device components.
+
+      fir::FirOpBuilder &builder = converter.getFirOpBuilder();
+      mlir::Location loc = converter.getCurrentLocation();
+
+      mlir::Type baseTy = fir::unwrapRefType(box.getAddr().getType());
+
+      // Only pointer and allocatable needs post allocation initialization
+      // of components descriptors.
+      if (!fir::isAllocatableType(baseTy) && !fir::isPointerType(baseTy))
+        return;
+
+      // Extract the derived type.
+      mlir::Type ty = fir::getDerivedType(baseTy);
+      auto recTy = mlir::dyn_cast<fir::RecordType>(ty);
+      assert(recTy && "expected fir::RecordType");
+
+      if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(baseTy))
+        baseTy = boxTy.getEleTy();
+      baseTy = fir::unwrapRefType(baseTy);
+
+      Fortran::semantics::UltimateComponentIterator components{*derived};
+      mlir::Value loadedBox = fir::LoadOp::create(builder, loc, box.getAddr());
+      mlir::Value addr;
+      if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(baseTy)) {
+        mlir::Type idxTy = builder.getIndexType();
+        mlir::Value one = builder.createIntegerConstant(loc, idxTy, 1);
+        mlir::Value zero = builder.createIntegerConstant(loc, idxTy, 0);
+        llvm::SmallVector<fir::DoLoopOp> loops;
+        llvm::SmallVector<mlir::Value> indices;
+        llvm::SmallVector<mlir::Value> extents;
+        for (unsigned i = 0; i < seqTy.getDimension(); ++i) {
+          mlir::Value dim = builder.createIntegerConstant(loc, idxTy, i);
+          auto dimInfo = fir::BoxDimsOp::create(builder, loc, idxTy, idxTy,
+                                                idxTy, loadedBox, dim);
+          mlir::Value lbub = mlir::arith::AddIOp::create(
+              builder, loc, dimInfo.getResult(0), dimInfo.getResult(1));
+          mlir::Value ext =
+              mlir::arith::SubIOp::create(builder, loc, lbub, one);
+          mlir::Value cmp = mlir::arith::CmpIOp::create(
+              builder, loc, mlir::arith::CmpIPredicate::sgt, ext, zero);
+          ext = mlir::arith::SelectOp::create(builder, loc, cmp, ext, zero);
+          extents.push_back(ext);
+
+          auto loop = fir::DoLoopOp::create(
+              builder, loc, dimInfo.getResult(0), dimInfo.getResult(1),
+              dimInfo.getResult(2), /*isUnordered=*/true,
+              /*finalCount=*/false, mlir::ValueRange{});
+          loops.push_back(loop);
+          indices.push_back(loop.getInductionVar());
+          builder.setInsertionPointToStart(loop.getBody());
+        }
+        mlir::Value boxAddr = fir::BoxAddrOp::create(builder, loc, loadedBox);
+        auto shape = fir::ShapeOp::create(builder, loc, extents);
+        addr = fir::ArrayCoorOp::create(
+            builder, loc, fir::ReferenceType::get(recTy), boxAddr, shape,
+            /*slice=*/mlir::Value{}, indices, /*typeparms=*/mlir::ValueRange{});
+      } else {
+        addr = fir::BoxAddrOp::create(builder, loc, loadedBox);
+      }
+      for (const auto &compSym : components) {
+        if (Fortran::semantics::IsDeviceAllocatable(compSym)) {
+          llvm::SmallVector<mlir::Value> coord;
+          mlir::Type fieldTy = gatherDeviceComponentCoordinatesAndType(
+              builder, loc, compSym, recTy, coord);
+          assert(coord.size() == 1 && "expect one coordinate");
+          mlir::Value comp = fir::CoordinateOp::create(
+              builder, loc, builder.getRefType(fieldTy), addr, coord[0]);
+          cuf::DataAttributeAttr dataAttr =
+              Fortran::lower::translateSymbolCUFDataAttribute(
+                  builder.getContext(), compSym);
+          cuf::SetAllocatorIndexOp::create(builder, loc, comp, dataAttr);
+        }
+      }
+    }
+  }
+}
+
+mlir::Type Fortran::lower::gatherDeviceComponentCoordinatesAndType(
+    fir::FirOpBuilder &builder, mlir::Location loc,
+    const Fortran::semantics::Symbol &sym, fir::RecordType recTy,
+    llvm::SmallVector<mlir::Value> &coordinates) {
+  unsigned fieldIdx = recTy.getFieldIndex(sym.name().ToString());
+  mlir::Type fieldTy;
+  if (fieldIdx != std::numeric_limits<unsigned>::max()) {
+    // Field found in the base record type.
+    auto fieldName = recTy.getTypeList()[fieldIdx].first;
+    fieldTy = recTy.getTypeList()[fieldIdx].second;
+    mlir::Value fieldIndex = fir::FieldIndexOp::create(
+        builder, loc, fir::FieldType::get(fieldTy.getContext()), fieldName,
+        recTy,
+        /*typeParams=*/mlir::ValueRange{});
+    coordinates.push_back(fieldIndex);
+  } else {
+    // Field not found in base record type, search in potential
+    // record type components.
+    for (auto component : recTy.getTypeList()) {
+      if (auto childRecTy = mlir::dyn_cast<fir::RecordType>(component.second)) {
+        fieldIdx = childRecTy.getFieldIndex(sym.name().ToString());
+        if (fieldIdx != std::numeric_limits<unsigned>::max()) {
+          mlir::Value parentFieldIndex = fir::FieldIndexOp::create(
+              builder, loc, fir::FieldType::get(childRecTy.getContext()),
+              component.first, recTy,
+              /*typeParams=*/mlir::ValueRange{});
+          coordinates.push_back(parentFieldIndex);
+          auto fieldName = childRecTy.getTypeList()[fieldIdx].first;
+          fieldTy = childRecTy.getTypeList()[fieldIdx].second;
+          mlir::Value childFieldIndex = fir::FieldIndexOp::create(
+              builder, loc, fir::FieldType::get(fieldTy.getContext()),
+              fieldName, childRecTy,
+              /*typeParams=*/mlir::ValueRange{});
+          coordinates.push_back(childFieldIndex);
+          break;
+        }
+      }
+    }
+  }
+  if (coordinates.empty())
+    TODO(loc, "device resident component in complex derived-type hierarchy");
+  return fieldTy;
+}
+
+cuf::DataAttributeAttr Fortran::lower::translateSymbolCUFDataAttribute(
+    mlir::MLIRContext *mlirContext, const Fortran::semantics::Symbol &sym) {
+  std::optional<Fortran::common::CUDADataAttr> cudaAttr =
+      Fortran::semantics::GetCUDADataAttr(&sym.GetUltimate());
+  return cuf::getDataAttribute(mlirContext, cudaAttr);
+}