[flang-commits] [flang] [flang][HLFIR] Relax InlineElementals to support more than two users (PR #186916)

[Tom Eccles via flang-commits]

================
@@ -31,29 +33,291 @@ namespace hlfir {
 #include "flang/Optimizer/HLFIR/Passes.h.inc"
 } // namespace hlfir
 
+/// Collects all memory values (buffers/references) that the elemental body
+/// reads from. Use MemoryEffectOpInterface for a fail-safe implementation.
+static mlir::LogicalResult
+getReadDependencies(hlfir::ElementalOp elemental,
+                    llvm::SmallVectorImpl<mlir::Value> &deps) {
+  llvm::SmallPtrSet<mlir::Value, 8> seen;
+
+  mlir::WalkResult walkResult =
+      elemental.getRegion().walk([&](mlir::Operation *op) {
+        if (mlir::isMemoryEffectFree(op))
+          return mlir::WalkResult::advance();
+
+        if (auto memInterface =
+                mlir::dyn_cast<mlir::MemoryEffectOpInterface>(op)) {
+          llvm::SmallVector<mlir::MemoryEffects::EffectInstance, 4> effects;
+          memInterface.getEffects(effects);
+          bool hasUnspecifiedRead = false;
+
+          for (const auto &effect : effects) {
+            if (mlir::isa<mlir::MemoryEffects::Read>(effect.getEffect())) {
+              if (mlir::Value val = effect.getValue()) {
+                if (seen.insert(val).second)
+                  deps.push_back(val);
+              } else {
+                // Read effect on an unspecified resource (e.g., global state).
+                hasUnspecifiedRead = true;
+              }
+            }
+          }
+
+          // If the op has a read effect but the specific value is unknown,
+          // conservatively capture all potential reference operands.
+          if (hasUnspecifiedRead) {
+            // If there are no operands to track, we can't reason about
+            // the dependency.
+            if (op->getNumOperands() == 0)
+              return mlir::WalkResult::interrupt();
+            for (mlir::Value operand : op->getOperands()) {
+              if (operand.getParentRegion() != &elemental.getRegion()) {
+                if (mlir::isa<fir::ReferenceType, fir::PointerType,
+                              fir::HeapType, fir::BoxType>(operand.getType())) {
+                  if (seen.insert(operand).second)
+                    deps.push_back(operand);
+                }
+              }
+            }
+          }
+          return mlir::WalkResult::advance();
+        }
+
+        // Fail-safe: For operations without the interface, conservatively
+        // assume we cannot reason about the dependency.
+        return mlir::WalkResult::interrupt();
+      });
+
+  return mlir::success(!walkResult.wasInterrupted());
+}
+
+/// Checks if an operation 'op' potentially modifies any memory location that
+/// the elemental reads from (captured in 'deps').
+static bool isConflictingWrite(mlir::Operation *op,
+                               const llvm::SmallVectorImpl<mlir::Value> &deps,
+                               mlir::AliasAnalysis &aa) {
+  // Use walk to handle nested regions (fir.if, fir.do_loop, etc.) recursively.
+  mlir::WalkResult result = op->walk([&](mlir::Operation *nestedOp) {
+    // Operations explicitly marked as having no memory effects are safe.
+    if (mlir::isMemoryEffectFree(nestedOp))
+      return mlir::WalkResult::advance();
+
+    // Explicitly allow safe HLFIR/FIR metadata/lifetime operations.
+    if (mlir::isa<hlfir::DeclareOp, hlfir::AssociateOp, hlfir::EndAssociateOp,
+                  fir::AllocaOp, hlfir::NoReassocOp>(nestedOp))
+      return mlir::WalkResult::advance();
+
+    // Check for explicit memory effects via the interface.
+    if (auto memInterface =
+            mlir::dyn_cast<mlir::MemoryEffectOpInterface>(nestedOp)) {
+      llvm::SmallVector<mlir::MemoryEffects::EffectInstance, 4> effects;
+      memInterface.getEffects(effects);
+
+      for (const auto &effect : effects) {
+        // Analyze effects that modify memory or release resources.
+        if (mlir::isa<mlir::MemoryEffects::Write, mlir::MemoryEffects::Free>(
+                effect.getEffect())) {
+          mlir::Value accessedValue = effect.getValue();
+          // Fail-safe: Assuming conflict for Unknown resource (e.g. external
+          // call).
+          if (!accessedValue)
+            return mlir::WalkResult::interrupt();
+
+          // Perform alias analysis against all read dependencies.
+          for (mlir::Value dep : deps) {
+            if (!aa.alias(accessedValue, dep).isNo())
+              return mlir::WalkResult::interrupt();
+          }
+        }
+      }
+    } else if (nestedOp->getNumRegions() == 0) {
+      // Conservative Fallback: If an operation doesn't have interface and
+      // has no regions (e.g. a fir.call), assume it can modify anything.
+      return mlir::WalkResult::interrupt();
+    }
+
+    return mlir::WalkResult::advance();
+  });
+
+  // Conflict found as walk interrupted.
+  return result.wasInterrupted();
+}
+
+bool isSafeToInline(hlfir::ElementalOp producer, hlfir::ApplyOp applySite,
+                    mlir::AliasAnalysis &aa, mlir::DominanceInfo &domInfo) {
+  if (!domInfo.properlyDominates(producer.getOperation(),
+                                 applySite.getOperation()))
+    return false;
+
+  llvm::SmallVector<mlir::Value> deps;
+  if (mlir::failed(getReadDependencies(producer, deps)))
+    return false;
+
+  mlir::Operation *func = producer->getParentOfType<mlir::func::FuncOp>();
+  if (!func)
+    return false;
+
+  // Check for conflicting writes between the producer and the apply site.
+  mlir::WalkResult result = func->walk([&](mlir::Operation *op) {
+    if (op == producer.getOperation() || op == applySite.getOperation())
+      return mlir::WalkResult::advance();
+
+    // Analyze operations in the execution path from producer to applySite.
+    if (domInfo.properlyDominates(producer.getOperation(), op) &&
+        domInfo.dominates(op, applySite.getOperation())) {
+      // If 'op' contains the applySite (like a loop shell), skip it to avoid
+      // false positives. Its internal operations will be visited individually.
+      if (op->getBlock() == applySite.getOperation()->getBlock()) {
+        if (isConflictingWrite(op, deps, aa))
+          return mlir::WalkResult::interrupt();
+      } else if (!op->isAncestor(applySite.getOperation())) {
+        // Check operations in sibling blocks or preceding control-flow paths.
+        if (isConflictingWrite(op, deps, aa))
+          return mlir::WalkResult::interrupt();
+      }
+    }
+    return mlir::WalkResult::advance();
+  });
+
+  return !result.wasInterrupted();
+}
+
 /// If the elemental has only two uses and those two are an apply operation and
 /// a destroy operation, return those two, otherwise return {}
----------------
tblah wrote:

nit: update comment

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