[flang][NFC] Characterize allocation based on MemAlloc effect instead of pattern matching

flang/lib/Optimizer/Analysis/AliasAnalysis.cpp

@@ -22,11 +22,73 @@
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
 
 using namespace mlir;
 
 #define DEBUG_TYPE "fir-alias-analysis"
 
+//===----------------------------------------------------------------------===//
+// AliasAnalysis: alias helpers
+//===----------------------------------------------------------------------===//
+
+static bool tryClassifyAllocateFromEffects(mlir::Operation *op,
+    mlir::Value candidate, bool allowValueScoped, bool allowOpScoped,
+    mlir::Value &v, mlir::Operation *&defOp,
+    fir::AliasAnalysis::SourceKind &type) {
+  auto iface = llvm::dyn_cast<mlir::MemoryEffectOpInterface>(op);
+  if (!iface)
+    return false;
+
+  llvm::SmallVector<mlir::MemoryEffects::EffectInstance, 4> effects;
+  iface.getEffects(effects);
+
+  if (allowValueScoped) {
+    for (mlir::MemoryEffects::EffectInstance &e : effects) {
+      if (mlir::isa<mlir::MemoryEffects::Allocate>(e.getEffect()) &&
+          e.getValue() && e.getValue() == candidate) {
+        v = candidate;
+        defOp = op;
+        type = fir::AliasAnalysis::SourceKind::Allocate;
+        return true;
+      }
+    }
+  }
+
+  if (!allowOpScoped)
+    return false;
+
+  bool hasOpScopedAlloc = llvm::any_of(
+      effects, [](const mlir::MemoryEffects::EffectInstance &e) {
+        return !e.getValue() &&
+               mlir::isa<mlir::MemoryEffects::Allocate>(e.getEffect());
+      });
+  if (!hasOpScopedAlloc)
+    return false;
+
+  bool opIsViewLike =
+      (bool)mlir::dyn_cast_or_null<mlir::ViewLikeOpInterface>(op);
+  auto isMemoryRefLikeType = [](mlir::Type type) {
+    return fir::isa_ref_type(type) || mlir::isa<mlir::BaseMemRefType>(type) ||
+           mlir::isa<mlir::LLVM::LLVMPointerType>(type);
+  };
+  bool hasMemOperands = llvm::any_of(op->getOperands(), [&](mlir::Value o) {
+    return isMemoryRefLikeType(o.getType());
+  });
+  if (opIsViewLike || hasMemOperands)
+    return false;
+
+  for (mlir::Value res : op->getResults()) {
+    if (res == candidate && isMemoryRefLikeType(res.getType())) {
+      v = candidate;
+      defOp = op;
+      type = fir::AliasAnalysis::SourceKind::Allocate;
+      return true;
+    }
+  }
+  return false;
+}
+
 //===----------------------------------------------------------------------===//
 // AliasAnalysis: alias
 //===----------------------------------------------------------------------===//
@@ -533,30 +595,47 @@ AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v,
   mlir::SymbolRefAttr global;
   Source::Attributes attributes;
   mlir::Operation *instantiationPoint{nullptr};
+  // Helper to conservatively classify a candidate value as coming from a
+  // dummy argument or as indirect when no allocation or global can be proven.
+  auto classifyFallbackFrom = [&](mlir::Value candidate) {
+    if (isDummyArgument(candidate)) {
+      defOp = nullptr;
+      v = candidate;
+    } else {
+      type = SourceKind::Indirect;
+    }
+  };
+
+  // Helper to detect memory-ref-like types.
+  auto isMemoryRefLikeType = [](mlir::Type t) {
+    return fir::isa_ref_type(t) || mlir::isa<mlir::BaseMemRefType>(t) ||
+           mlir::isa<mlir::LLVM::LLVMPointerType>(t);
+  };
+
   while (defOp && !breakFromLoop) {
     ty = defOp->getResultTypes()[0];
+
+    // Effect-based detection (op-scoped heuristic only at this level).
+    if (tryClassifyAllocateFromEffects(defOp, v,
+                                       /*allowValueScoped=*/false,
+                                       /*allowOpScoped=*/true,
+                                       v, defOp, type))
+      break;
+
     llvm::TypeSwitch<Operation *>(defOp)
         .Case<hlfir::AsExprOp>([&](auto op) {
           v = op.getVar();
           defOp = v.getDefiningOp();
         })
         .Case<hlfir::AssociateOp>([&](auto op) {
+          // Do not pattern-match Allocate. Trace through the source.
           mlir::Value source = op.getSource();
-          if (fir::isa_trivial(source.getType())) {
-            // Trivial values will always use distinct temp memory,
-            // so we can classify this as Allocate and stop.
-            type = SourceKind::Allocate;
-            breakFromLoop = true;
-          } else {
-            // AssociateOp may reuse the expression storage,
-            // so we have to trace further.
-            v = source;
-            defOp = v.getDefiningOp();
-          }
+          v = source;
+          defOp = v.getDefiningOp();
         })
         .Case<fir::AllocaOp, fir::AllocMemOp>([&](auto op) {
-          // Unique memory allocation.
-          type = SourceKind::Allocate;
+          // Do not pattern-match allocations by op name; rely on memory
+          // effects classification above. Nothing to do here.
           breakFromLoop = true;
         })
         .Case<fir::ConvertOp>([&](auto op) {
@@ -627,18 +706,14 @@ AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v,
             if (global) {
               type = SourceKind::Global;
             } else {
-              auto def = llvm::cast<mlir::Value>(boxSrc.origin.u);
-              // TODO: Add support to fir.allocmem
-              if (auto allocOp = def.template getDefiningOp<fir::AllocaOp>()) {
-                v = def;
-                defOp = v.getDefiningOp();
-                type = SourceKind::Allocate;
-              } else if (isDummyArgument(def)) {
-                defOp = nullptr;
-                v = def;
-              } else {
-                type = SourceKind::Indirect;
-              }
+              mlir::Value def = llvm::cast<mlir::Value>(boxSrc.origin.u);
+              bool classified = false;
+              if (auto defDefOp = def.getDefiningOp())
+                classified = tryClassifyAllocateFromEffects(
+                    defDefOp, def,
+                    /*allowValueScoped=*/true, /*allowOpScoped=*/true,
+                    v, defOp, type);
+              if (!classified) classifyFallbackFrom(def);
             }
             breakFromLoop = true;
             return;

flang/test/Analysis/AliasAnalysis/cuf-alloc-source-kind.mlir

@@ -0,0 +1,22 @@
+// REQUIRES: asserts
+// RUN: fir-opt %s -pass-pipeline='builtin.module(func.func(test-fir-alias-analysis))' -debug-only=fir-alias-analysis --mlir-disable-threading 2>&1 | FileCheck %s
+
+// Verify that a CUF allocation is recognized as SourceKind::Allocate by
+// fir::AliasAnalysis::getSource.
+
+module {
+  func.func @_QQmain() attributes {fir.bindc_name = "TEST"} {
+    // Allocate two independent device arrays and tag the results; with
+    // op-scoped MemAlloc handling in AA, these should be classified as
+    // Allocate and not alias.
+    %a = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a1", data_attr = #cuf.cuda<device>, uniq_name = "_QFEa1", test.ptr = "cuf_alloc_a"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+    %b = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a2", data_attr = #cuf.cuda<device>, uniq_name = "_QFEa2", test.ptr = "cuf_alloc_b"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+    return
+  }
+}
+
+// CHECK-LABEL: Testing : "_QQmain"
+// Distinct allocations should not alias.
+// CHECK: cuf_alloc_a#0 <-> cuf_alloc_b#0: NoAlias
+
+