[flang-commits] [flang] [Flang][OpenMP] Fix crash and IR errors for user-defined reduction on allocatable variables (PR #186765)

[via flang-commits]

================
@@ -734,6 +735,150 @@ bool ReductionProcessor::doReductionByRef(mlir::Value reductionVar) {
   return doReductionByRef(reductionVar.getType());
 }
 
+static bool baseInitReadsOrig(mlir::Block &initBlock) {
+  bool reads = false;
+  initBlock.walk([&](hlfir::DeclareOp declOp) {
+    std::optional<llvm::StringRef> name = declOp.getUniqName();
+    if (name && *name == "omp_orig" &&
+        (!declOp.getResult(0).use_empty() || !declOp.getResult(1).use_empty()))
+      reads = true;
+  });
+  return reads;
+}
+
+template <typename OpType>
+static OpType getOrCreateBoxedUserReduction(
+    lower::AbstractConverter &converter, semantics::SemanticsContext *semaCtx,
+    const semantics::Symbol &redSym, mlir::Type redType, mlir::Location loc) {
+  fir::FirOpBuilder &builder = converter.getFirOpBuilder();
+  mlir::ModuleOp module = builder.getModule();
+
+  // Boxed op name encodes the full boxed type (..._byref_box_heap_i32),
+  // distinct from the base element op.
+  std::string boxedName = ReductionProcessor::getScopedUserReductionName(
+      converter, redSym, redType, /*isByRef=*/true);
+  if (auto existing = module.lookupSymbol<OpType>(boxedName))
+    return existing;
+
+  // Base op is named from the scalar element type (identical for scalar and
+  // array); materialize on demand for separate compilation.
+  mlir::Type elementType = unwrapSeqOrBoxedType(fir::unwrapRefType(redType));
+  const bool baseByRef = ReductionProcessor::doReductionByRef(elementType);
+  std::string baseName = ReductionProcessor::getScopedUserReductionName(
+      converter, redSym, elementType, baseByRef);
+  auto baseDecl = module.lookupSymbol<OpType>(baseName);
+  if (!baseDecl && semaCtx && redSym.owner().symbol() &&
+      redSym.owner().symbol()->test(semantics::Symbol::Flag::ModFile)) {
+    Fortran::lower::materializeUserReduction(converter, *semaCtx, redSym,
+                                             baseName, elementType, baseByRef);
+    baseDecl = module.lookupSymbol<OpType>(baseName);
+  }
+  if (!baseDecl)
+    return {};
+
+  // Only a by-value base op can be wrapped (its init/combiner yield scalars).
+  mlir::Region &baseInitRegion = baseDecl.getInitializerRegion();
+  mlir::Region &baseCombinerRegion = baseDecl.getReductionRegion();
+  mlir::Block &baseInitBlock = baseInitRegion.front();
+  if (fir::isa_ref_type(baseInitBlock.getArgument(0).getType()) ||
+      fir::isa_ref_type(baseCombinerRegion.front().getArgument(0).getType()))
+    return {};
+  const bool initReadsOrig = baseInitReadsOrig(baseInitBlock);
+
+  // Clone the region's entry block (up to its terminator) mapping block args to
+  // \p args; return the yielded value.
+  auto cloneBody = [](fir::FirOpBuilder &b, mlir::Region &region,
+                      mlir::ValueRange args) -> mlir::Value {
+    mlir::Block &block = region.front();
+    mlir::IRMapping mapper;
+    for (auto [blockArg, val] : llvm::zip(block.getArguments(), args))
+      mapper.map(blockArg, val);
+    for (mlir::Operation &op : block.without_terminator())
+      b.clone(op, mapper);
+    mlir::Operation *term = block.getTerminator();
+    if (term && term->getNumOperands() > 0)
+      return mapper.lookupOrDefault(term->getOperand(0));
+    return {};
+  };
+
+  auto genInitValueCB = [&](fir::FirOpBuilder &b, mlir::Location loc,
+                            mlir::Type ty, mlir::Value moldArg,
+                            mlir::Value /*privArg*/) -> mlir::Value {
+    mlir::Type scalarTy = unwrapSeqOrBoxedType(ty);
+    llvm::SmallVector<mlir::Value> initArgs;
+    if (initReadsOrig) {
+      auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(fir::unwrapRefType(ty));
+      auto seqTy = boxTy ? mlir::dyn_cast_or_null<fir::SequenceType>(
+                               fir::unwrapRefType(boxTy.getEleTy()))
+                         : nullptr;
+      if (seqTy)
+        TODO(loc, "OpenMP user-defined reduction on an allocatable/pointer "
+                  "array whose initializer reads omp_orig");
+      // Scalar: feed the original element to the cloned initializer's omp_orig.
+      mlir::Value box = fir::LoadOp::create(b, loc, moldArg);
+      mlir::Value addr = fir::BoxAddrOp::create(b, loc, box);
+      initArgs.push_back(fir::LoadOp::create(b, loc, addr));
+    } else {
+      // orig-independent: a placeholder feeds the base's dead temporary stores.
+      initArgs.push_back(fir::UndefOp::create(b, loc, scalarTy));
+    }
+    return cloneBody(b, baseInitRegion, initArgs);
----------------
MattPD wrote:

Non-blocking. For an allocatable array, this returns a single scalar init value, which the shared boxed-array init then broadcasts to every element with `hlfir.assign`. For a constant identity like `omp_priv = 0` that is correct and matches the intrinsic reduction path. An initializer with observable per-element evaluation cannot safely be broadcast, though: you already gate an array `omp_priv = omp_orig` behind a clean TODO in this callback, but a non-`omp_orig`, non-constant initializer such as `omp_priv = f()` for some function `f` falls through to a single call broadcast across the whole array. I checked, and `f` is invoked once, not once per element. Is broadcasting intended here (treating the value as the reduction identity), or would routing these initializers through the same TODO be more consistent?

https://github.com/llvm/llvm-project/pull/186765