[flang][acc] Implement MappableType's generatePrivateInit (#148302)

razvanlupusoru · web-flow · commit b54cfa46a7eb · 2025-07-14T10:53:54.000-07:00
The recipe body generation was moved from lowering into FIR's
implementation of MappableType API. And now since all Fortran variable
types implement this type, lowering of OpenACC was updated to use this
API directly. No test changes were needed - all of the private,
firstprivate, and recipe tests get the same body as before.
diff --git a/flang/include/flang/Optimizer/OpenACC/FIROpenACCTypeInterfaces.h b/flang/include/flang/Optimizer/OpenACC/FIROpenACCTypeInterfaces.h
@@ -52,6 +52,13 @@ struct OpenACCMappableModel
 
   mlir::acc::VariableTypeCategory getTypeCategory(mlir::Type type,
                                                   mlir::Value var) const;
+
+  mlir::Value generatePrivateInit(mlir::Type type, mlir::OpBuilder &builder,
+                                  mlir::Location loc,
+                                  mlir::TypedValue<mlir::acc::MappableType> var,
+                                  llvm::StringRef varName,
+                                  mlir::ValueRange extents,
+                                  mlir::Value initVal) const;
 };
 
 } // namespace fir::acc
diff --git a/flang/lib/Lower/CMakeLists.txt b/flang/lib/Lower/CMakeLists.txt
@@ -51,6 +51,7 @@ add_flang_library(FortranLower
   FIRDialect
   FIRDialectSupport
   FIRBuilder
+  FIROpenACCSupport
   FIRSupport
   FIRTransforms
   HLFIRDialect
diff --git a/flang/lib/Lower/OpenACC.cpp b/flang/lib/Lower/OpenACC.cpp
@@ -959,119 +959,6 @@ static mlir::Value getReductionInitValue(fir::FirOpBuilder &builder,
   llvm::report_fatal_error("Unsupported OpenACC reduction type");
 }
 
-template <typename RecipeOp>
-static void genPrivateLikeInitRegion(fir::FirOpBuilder &builder,
-                                     RecipeOp recipe, mlir::Type argTy,
-                                     mlir::Location loc,
-                                     mlir::Value initValue) {
-  mlir::Value retVal = recipe.getInitRegion().front().getArgument(0);
-  mlir::Type unwrappedTy = fir::unwrapRefType(argTy);
-
-  llvm::StringRef initName;
-  if constexpr (std::is_same_v<RecipeOp, mlir::acc::ReductionRecipeOp>)
-    initName = accReductionInitName;
-  else
-    initName = accPrivateInitName;
-
-  auto getDeclareOpForType = [&](mlir::Type ty) -> hlfir::DeclareOp {
-    auto alloca = builder.create<fir::AllocaOp>(loc, ty);
-    return builder.create<hlfir::DeclareOp>(
-        loc, alloca, initName, /*shape=*/nullptr, llvm::ArrayRef<mlir::Value>{},
-        /*dummy_scope=*/nullptr, fir::FortranVariableFlagsAttr{});
-  };
-
-  if (fir::isa_trivial(unwrappedTy)) {
-    auto declareOp = getDeclareOpForType(unwrappedTy);
-    if (initValue) {
-      auto convert = builder.createConvert(loc, unwrappedTy, initValue);
-      builder.create<fir::StoreOp>(loc, convert, declareOp.getBase());
-    }
-    retVal = declareOp.getBase();
-  } else if (auto seqTy =
-                 mlir::dyn_cast_or_null<fir::SequenceType>(unwrappedTy)) {
-    if (fir::isa_trivial(seqTy.getEleTy())) {
-      mlir::Value shape;
-      llvm::SmallVector<mlir::Value> extents;
-      if (seqTy.hasDynamicExtents()) {
-        // Extents are passed as block arguments. First argument is the
-        // original value.
-        for (unsigned i = 1; i < recipe.getInitRegion().getArguments().size();
-             ++i)
-          extents.push_back(recipe.getInitRegion().getArgument(i));
-        shape = builder.create<fir::ShapeOp>(loc, extents);
-      } else {
-        shape = genShapeOp(builder, seqTy, loc);
-      }
-      auto alloca = builder.create<fir::AllocaOp>(
-          loc, seqTy, /*typeparams=*/mlir::ValueRange{}, extents);
-      auto declareOp = builder.create<hlfir::DeclareOp>(
-          loc, alloca, initName, shape, llvm::ArrayRef<mlir::Value>{},
-          /*dummy_scope=*/nullptr, fir::FortranVariableFlagsAttr{});
-
-      if (initValue) {
-        mlir::Type idxTy = builder.getIndexType();
-        mlir::Type refTy = fir::ReferenceType::get(seqTy.getEleTy());
-        llvm::SmallVector<fir::DoLoopOp> loops;
-        llvm::SmallVector<mlir::Value> ivs;
-
-        if (seqTy.hasDynamicExtents()) {
-          builder.create<hlfir::AssignOp>(loc, initValue, declareOp.getBase());
-        } else {
-          for (auto ext : seqTy.getShape()) {
-            auto lb = builder.createIntegerConstant(loc, idxTy, 0);
-            auto ub = builder.createIntegerConstant(loc, idxTy, ext - 1);
-            auto step = builder.createIntegerConstant(loc, idxTy, 1);
-            auto loop = builder.create<fir::DoLoopOp>(loc, lb, ub, step,
-                                                      /*unordered=*/false);
-            builder.setInsertionPointToStart(loop.getBody());
-            loops.push_back(loop);
-            ivs.push_back(loop.getInductionVar());
-          }
-          auto coord = builder.create<fir::CoordinateOp>(
-              loc, refTy, declareOp.getBase(), ivs);
-          builder.create<fir::StoreOp>(loc, initValue, coord);
-          builder.setInsertionPointAfter(loops[0]);
-        }
-      }
-      retVal = declareOp.getBase();
-    }
-  } else if (auto boxTy =
-                 mlir::dyn_cast_or_null<fir::BaseBoxType>(unwrappedTy)) {
-    mlir::Type innerTy = fir::unwrapRefType(boxTy.getEleTy());
-    if (fir::isa_trivial(innerTy)) {
-      retVal = getDeclareOpForType(unwrappedTy).getBase();
-    } else if (mlir::isa<fir::SequenceType>(innerTy)) {
-      fir::FirOpBuilder firBuilder{builder, recipe.getOperation()};
-      hlfir::Entity source = hlfir::Entity{retVal};
-      auto [temp, cleanup] = hlfir::createTempFromMold(loc, firBuilder, source);
-      if (fir::isa_ref_type(argTy)) {
-        // When the temp is created - it is not a reference - thus we can
-        // end up with a type inconsistency. Therefore ensure storage is created
-        // for it.
-        retVal = getDeclareOpForType(unwrappedTy).getBase();
-        mlir::Value storeDst = retVal;
-        if (fir::unwrapRefType(retVal.getType()) != temp.getType()) {
-          // `createTempFromMold` makes the unfortunate choice to lose the
-          // `fir.heap` and `fir.ptr` types when wrapping with a box. Namely,
-          // when wrapping a `fir.heap<fir.array>`, it will create instead a
-          // `fir.box<fir.array>`. Cast here to deal with this inconsistency.
-          storeDst = firBuilder.createConvert(
-              loc, firBuilder.getRefType(temp.getType()), retVal);
-        }
-        builder.create<fir::StoreOp>(loc, temp, storeDst);
-      } else {
-        retVal = temp;
-      }
-    } else {
-      TODO(loc, "Unsupported boxed type for OpenACC private-like recipe");
-    }
-    if (initValue) {
-      builder.create<hlfir::AssignOp>(loc, initValue, retVal);
-    }
-  }
-  builder.create<mlir::acc::YieldOp>(loc, retVal);
-}
-
 template <typename RecipeOp>
 static RecipeOp genRecipeOp(
     fir::FirOpBuilder &builder, mlir::ModuleOp mod, llvm::StringRef recipeName,
@@ -1100,15 +987,35 @@ static RecipeOp genRecipeOp(
       }
     }
   }
-  builder.createBlock(&recipe.getInitRegion(), recipe.getInitRegion().end(),
-                      argsTy, argsLoc);
+  auto initBlock = builder.createBlock(
+      &recipe.getInitRegion(), recipe.getInitRegion().end(), argsTy, argsLoc);
   builder.setInsertionPointToEnd(&recipe.getInitRegion().back());
   mlir::Value initValue;
   if constexpr (std::is_same_v<RecipeOp, mlir::acc::ReductionRecipeOp>) {
     assert(op != mlir::acc::ReductionOperator::AccNone);
     initValue = getReductionInitValue(builder, loc, fir::unwrapRefType(ty), op);
   }
-  genPrivateLikeInitRegion<RecipeOp>(builder, recipe, ty, loc, initValue);
+
+  // Since we reuse the same recipe for all variables of the same type - we
+  // cannot use the actual variable name. Thus use a temporary name.
+  llvm::StringRef initName;
+  if constexpr (std::is_same_v<RecipeOp, mlir::acc::ReductionRecipeOp>)
+    initName = accReductionInitName;
+  else
+    initName = accPrivateInitName;
+
+  auto mappableTy = mlir::dyn_cast<mlir::acc::MappableType>(ty);
+  assert(mappableTy &&
+         "Expected that all variable types are considered mappable");
+  auto retVal = mappableTy.generatePrivateInit(
+      builder, loc,
+      mlir::cast<mlir::TypedValue<mlir::acc::MappableType>>(
+          initBlock->getArgument(0)),
+      initName,
+      initBlock->getArguments().take_back(initBlock->getArguments().size() - 1),
+      initValue);
+  builder.create<mlir::acc::YieldOp>(loc, retVal ? retVal
+                                                 : initBlock->getArgument(0));
   return recipe;
 }
 
diff --git a/flang/lib/Optimizer/OpenACC/FIROpenACCTypeInterfaces.cpp b/flang/lib/Optimizer/OpenACC/FIROpenACCTypeInterfaces.cpp
@@ -25,6 +25,7 @@
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Support/LLVM.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 namespace fir::acc {
@@ -525,4 +526,141 @@ OpenACCPointerLikeModel<fir::LLVMPointerType>::getPointeeTypeCategory(
   return categorizePointee(pointer, varPtr, varType);
 }
 
+static fir::ShapeOp genShapeOp(mlir::OpBuilder &builder,
+                               fir::SequenceType seqTy, mlir::Location loc) {
+  llvm::SmallVector<mlir::Value> extents;
+  mlir::Type idxTy = builder.getIndexType();
+  for (auto extent : seqTy.getShape())
+    extents.push_back(builder.create<mlir::arith::ConstantOp>(
+        loc, idxTy, builder.getIntegerAttr(idxTy, extent)));
+  return builder.create<fir::ShapeOp>(loc, extents);
+}
+
+template <typename Ty>
+mlir::Value OpenACCMappableModel<Ty>::generatePrivateInit(
+    mlir::Type type, mlir::OpBuilder &builder, mlir::Location loc,
+    mlir::TypedValue<mlir::acc::MappableType> var, llvm::StringRef varName,
+    mlir::ValueRange extents, mlir::Value initVal) const {
+  mlir::Value retVal;
+  mlir::Type unwrappedTy = fir::unwrapRefType(type);
+  mlir::ModuleOp mod = builder.getInsertionBlock()
+                           ->getParent()
+                           ->getParentOfType<mlir::ModuleOp>();
+  fir::FirOpBuilder firBuilder(builder, mod);
+
+  auto getDeclareOpForType = [&](mlir::Type ty) -> hlfir::DeclareOp {
+    auto alloca = firBuilder.create<fir::AllocaOp>(loc, ty);
+    return firBuilder.create<hlfir::DeclareOp>(
+        loc, alloca, varName, /*shape=*/nullptr, llvm::ArrayRef<mlir::Value>{},
+        /*dummy_scope=*/nullptr, fir::FortranVariableFlagsAttr{});
+  };
+
+  if (fir::isa_trivial(unwrappedTy)) {
+    auto declareOp = getDeclareOpForType(unwrappedTy);
+    if (initVal) {
+      auto convert = firBuilder.createConvert(loc, unwrappedTy, initVal);
+      firBuilder.create<fir::StoreOp>(loc, convert, declareOp.getBase());
+    }
+    retVal = declareOp.getBase();
+  } else if (auto seqTy =
+                 mlir::dyn_cast_or_null<fir::SequenceType>(unwrappedTy)) {
+    if (fir::isa_trivial(seqTy.getEleTy())) {
+      mlir::Value shape;
+      if (seqTy.hasDynamicExtents()) {
+        shape = firBuilder.create<fir::ShapeOp>(loc, llvm::to_vector(extents));
+      } else {
+        shape = genShapeOp(firBuilder, seqTy, loc);
+      }
+      auto alloca = firBuilder.create<fir::AllocaOp>(
+          loc, seqTy, /*typeparams=*/mlir::ValueRange{}, extents);
+      auto declareOp = firBuilder.create<hlfir::DeclareOp>(
+          loc, alloca, varName, shape, llvm::ArrayRef<mlir::Value>{},
+          /*dummy_scope=*/nullptr, fir::FortranVariableFlagsAttr{});
+
+      if (initVal) {
+        mlir::Type idxTy = firBuilder.getIndexType();
+        mlir::Type refTy = fir::ReferenceType::get(seqTy.getEleTy());
+        llvm::SmallVector<fir::DoLoopOp> loops;
+        llvm::SmallVector<mlir::Value> ivs;
+
+        if (seqTy.hasDynamicExtents()) {
+          firBuilder.create<hlfir::AssignOp>(loc, initVal, declareOp.getBase());
+        } else {
+          for (auto ext : seqTy.getShape()) {
+            auto lb = firBuilder.createIntegerConstant(loc, idxTy, 0);
+            auto ub = firBuilder.createIntegerConstant(loc, idxTy, ext - 1);
+            auto step = firBuilder.createIntegerConstant(loc, idxTy, 1);
+            auto loop = firBuilder.create<fir::DoLoopOp>(loc, lb, ub, step,
+                                                         /*unordered=*/false);
+            firBuilder.setInsertionPointToStart(loop.getBody());
+            loops.push_back(loop);
+            ivs.push_back(loop.getInductionVar());
+          }
+          auto coord = firBuilder.create<fir::CoordinateOp>(
+              loc, refTy, declareOp.getBase(), ivs);
+          firBuilder.create<fir::StoreOp>(loc, initVal, coord);
+          firBuilder.setInsertionPointAfter(loops[0]);
+        }
+      }
+      retVal = declareOp.getBase();
+    }
+  } else if (auto boxTy =
+                 mlir::dyn_cast_or_null<fir::BaseBoxType>(unwrappedTy)) {
+    mlir::Type innerTy = fir::unwrapRefType(boxTy.getEleTy());
+    if (fir::isa_trivial(innerTy)) {
+      retVal = getDeclareOpForType(unwrappedTy).getBase();
+    } else if (mlir::isa<fir::SequenceType>(innerTy)) {
+      hlfir::Entity source = hlfir::Entity{var};
+      auto [temp, cleanup] = hlfir::createTempFromMold(loc, firBuilder, source);
+      if (fir::isa_ref_type(type)) {
+        // When the temp is created - it is not a reference - thus we can
+        // end up with a type inconsistency. Therefore ensure storage is created
+        // for it.
+        retVal = getDeclareOpForType(unwrappedTy).getBase();
+        mlir::Value storeDst = retVal;
+        if (fir::unwrapRefType(retVal.getType()) != temp.getType()) {
+          // `createTempFromMold` makes the unfortunate choice to lose the
+          // `fir.heap` and `fir.ptr` types when wrapping with a box. Namely,
+          // when wrapping a `fir.heap<fir.array>`, it will create instead a
+          // `fir.box<fir.array>`. Cast here to deal with this inconsistency.
+          storeDst = firBuilder.createConvert(
+              loc, firBuilder.getRefType(temp.getType()), retVal);
+        }
+        builder.create<fir::StoreOp>(loc, temp, storeDst);
+      } else {
+        retVal = temp;
+      }
+    } else {
+      TODO(loc, "Unsupported boxed type for OpenACC private-like recipe");
+    }
+    if (initVal) {
+      builder.create<hlfir::AssignOp>(loc, initVal, retVal);
+    }
+  }
+  return retVal;
+}
+
+template mlir::Value
+OpenACCMappableModel<fir::BaseBoxType>::generatePrivateInit(
+    mlir::Type type, mlir::OpBuilder &builder, mlir::Location loc,
+    mlir::TypedValue<mlir::acc::MappableType> var, llvm::StringRef varName,
+    mlir::ValueRange extents, mlir::Value initVal) const;
+
+template mlir::Value
+OpenACCMappableModel<fir::ReferenceType>::generatePrivateInit(
+    mlir::Type type, mlir::OpBuilder &builder, mlir::Location loc,
+    mlir::TypedValue<mlir::acc::MappableType> var, llvm::StringRef varName,
+    mlir::ValueRange extents, mlir::Value initVal) const;
+
+template mlir::Value OpenACCMappableModel<fir::HeapType>::generatePrivateInit(
+    mlir::Type type, mlir::OpBuilder &builder, mlir::Location loc,
+    mlir::TypedValue<mlir::acc::MappableType> var, llvm::StringRef varName,
+    mlir::ValueRange extents, mlir::Value initVal) const;
+
+template mlir::Value
+OpenACCMappableModel<fir::PointerType>::generatePrivateInit(
+    mlir::Type type, mlir::OpBuilder &builder, mlir::Location loc,
+    mlir::TypedValue<mlir::acc::MappableType> var, llvm::StringRef varName,
+    mlir::ValueRange extents, mlir::Value initVal) const;
+
 } // namespace fir::acc
