MLIR: alias classes for globals, and alloc-like effect handling

enzyme/Enzyme/MLIR/Analysis/DataFlowAliasAnalysis.cpp

@@ -248,6 +248,51 @@ static bool isNoOp(Operation *op) {
              LLVM::LifetimeEndOp, LLVM::AssumeOp, LLVM::UnreachableOp>(op);
 }
 
+static bool isAliasRelevantEffect(const MemoryEffects::EffectInstance &effect) {
+  Value v = effect.getValue();
+  if (v)
+    return true;
+
+  if (const auto *resource = effect.getResource())
+    return resource->isAddressable();
+
+  return true;
+}
+
+static Value getSingleFreshAllocatedPointerLikeResultIfEffectOnlyAllocLike(
+    Operation *op, ArrayRef<MemoryEffects::EffectInstance> effects) {
+  Value allocatedValue = nullptr;
+  bool sawAllocate = false;
+
+  for (const auto &effect : effects) {
+    Value v = effect.getValue();
+
+    if (isa<MemoryEffects::Allocate>(effect.getEffect()) && v &&
+        v.getDefiningOp() == op && isPointerLike(v.getType())) {
+      sawAllocate = true;
+
+      // Single-result helper: reject multiple allocated pointer-like results.
+      if (!allocatedValue)
+        allocatedValue = v;
+      else
+        return nullptr;
+
+      continue;
+    }
+
+    // Ignore non-alias-relevant effects, e.g. runtime/comm effects on
+    // non-addressable resources.
+    if (!isAliasRelevantEffect(effect))
+      continue;
+
+    // Any other alias-relevant effect means this op is not "pure alloc-like"
+    // for our purposes.
+    return nullptr;
+  }
+
+  return sawAllocate ? allocatedValue : nullptr;
+}
+
 LogicalResult enzyme::PointsToPointerAnalysis::visitOperation(
     Operation *op, const PointsToSets &before, PointsToSets *after) {
   join(after, before);
@@ -267,11 +312,11 @@ LogicalResult enzyme::PointsToPointerAnalysis::visitOperation(
 
   // If the operation allocates fresh memory and doesn't write into it, that
   // memory is known not to point to any known alias class.
-  if (effects.size() == 1 &&
-      isa<MemoryEffects::Allocate>(effects.front().getEffect()) &&
-      effects.front().getValue()) {
+  if (Value allocatedValue =
+          getSingleFreshAllocatedPointerLikeResultIfEffectOnlyAllocLike(
+              op, effects)) {
     const auto *destClasses = getOrCreateFor<AliasClassLattice>(
-        getProgramPointAfter(op), effects.front().getValue());
+        getProgramPointAfter(op), allocatedValue);
     propagateIfChanged(
         after, after->setPointingToEmpty(destClasses->getAliasClassesObject()));
     return success();
@@ -820,17 +865,25 @@ void enzyme::AliasAnalysis::setToEntryState(AliasClassLattice *lattice) {
 // would need additional processing.
 //
 // TODO: turn this into an interface.
-static bool isAliasTransferFullyDescribedByMemoryEffects(Operation *op) {
+static bool isAliasTransferFullyDescribedByMemoryEffects(
+    Operation *op, ArrayRef<MemoryEffects::EffectInstance> effects) {
   if (auto call = dyn_cast<CallOpInterface>(op)) {
     if (auto symbol = dyn_cast<SymbolRefAttr>(call.getCallableForCallee())) {
-      if (symbol.getLeafReference().getValue() == "malloc") {
+      StringRef name = symbol.getLeafReference().getValue();
+      if (name == "malloc" || name == "calloc" || name == "_Znwm")
         return true;
-      }
     }
   }
-  return isa<memref::LoadOp, memref::StoreOp, affine::AffineLoadOp,
-             affine::AffineStoreOp, LLVM::LoadOp, LLVM::StoreOp, enzyme::PushOp,
-             enzyme::PopOp>(op);
+
+  if (isa<memref::LoadOp, memref::StoreOp, affine::AffineLoadOp,
+          affine::AffineStoreOp, LLVM::LoadOp, LLVM::StoreOp, enzyme::PushOp,
+          enzyme::PopOp>(op))
+    return true;
+
+  // Generic alloc-like op: fresh pointer-like result, no alias-relevant
+  // effects other than the allocation itself.
+  return getSingleFreshAllocatedPointerLikeResultIfEffectOnlyAllocLike(
+             op, effects) != nullptr;
 }
 
 void enzyme::AliasAnalysis::transfer(
@@ -842,7 +895,8 @@ void enzyme::AliasAnalysis::transfer(
     // If the effect is global read, record that.
     Value value = effect.getValue();
     if (!value) {
-      globalRead |= isa<MemoryEffects::Read>(effect.getEffect());
+      globalRead |= isa<MemoryEffects::Read>(effect.getEffect()) &&
+                    isAliasRelevantEffect(effect);
       continue;
     }
 
@@ -913,7 +967,8 @@ void enzyme::AliasAnalysis::transfer(
   }
 
   // If it was enough to reason about effects, exit here.
-  if (!effects.empty() && isAliasTransferFullyDescribedByMemoryEffects(op))
+  if (!effects.empty() &&
+      isAliasTransferFullyDescribedByMemoryEffects(op, effects))
     return;
 
   // Conservatively assume all results alias all operands.
@@ -1006,9 +1061,64 @@ LogicalResult getEffectsForExternalCall(
   return failure();
 }
 
+// Whether an `llvm.mlir.addressof` of this symbol has to fall back on the
+// shared entry class, or can be given its own. Symbols that may be defined
+// outside the module can be aliased by an incoming pointer, so they share the
+// entry class; symbols whose definition the module owns get a unique class.
+static bool shouldUseEntryClassForAddressOf(Operation *symbol) {
+  if (isa<LLVM::GlobalOp>(symbol)) {
+    // A global's storage is a distinct object even when the definition is
+    // external: taking its address cannot produce a pointer into some other
+    // global. Give every global its own alias class.
+    //
+    // TODO: this ignores unnamed_addr. A global marked unnamed_addr has no
+    // meaningful identity and may be merged with an equivalent constant, so
+    // strictly it should not get a unique class of its own.
+    return false;
+  }
+
+  if (auto f = dyn_cast<LLVM::LLVMFuncOp>(symbol)) {
+    auto linkage = f.getLinkage();
+    return linkage == LLVM::Linkage::External ||
+           linkage == LLVM::Linkage::ExternWeak ||
+           linkage == LLVM::Linkage::AvailableExternally;
+  }
+
+  if (auto a = dyn_cast<LLVM::AliasOp>(symbol)) {
+    auto linkage = a.getLinkage();
+    return linkage == LLVM::Linkage::External ||
+           linkage == LLVM::Linkage::ExternWeak ||
+           linkage == LLVM::Linkage::AvailableExternally;
+  }
+
+  // Conservative fallback.
+  return true;
+}
+
 LogicalResult enzyme::AliasAnalysis::visitOperation(
     Operation *op, ArrayRef<const AliasClassLattice *> operands,
     ArrayRef<AliasClassLattice *> results) {
+  if (auto addr = dyn_cast<LLVM::AddressOfOp>(op)) {
+    AliasClassLattice *resultLattice = results[0];
+
+    Operation *symbol =
+        SymbolTable::lookupNearestSymbolFrom(op, addr.getGlobalNameAttr());
+
+    DistinctAttr cls = nullptr;
+    if (!symbol || shouldUseEntryClassForAddressOf(symbol)) {
+      cls = entryClass;
+    } else {
+      // Keyed on the symbol, not on this result: two `llvm.mlir.addressof` of
+      // the same global must land in the same alias class.
+      cls = originalClasses.getSymbolClass(symbol, addr.getGlobalNameAttr());
+    }
+
+    propagateIfChanged(resultLattice,
+                       resultLattice->join(AliasClassLattice::single(
+                           resultLattice->getAnchor(), cls)));
+
+    return success();
+  }
 
   // If we don't have memory effect information, don't assume anything about
   // values.

enzyme/Enzyme/MLIR/Analysis/DataFlowAliasAnalysis.h

@@ -69,6 +69,19 @@ class OriginalClasses {
     return aliasClass;
   }
 
+  /// Alias class of the storage designated by a symbol, e.g. an
+  /// `llvm.mlir.global`. Keyed on the symbol operation rather than on a value
+  /// so that every reference to the same symbol gets the same class.
+  DistinctAttr getSymbolClass(Operation *symbol, Attribute referenced) {
+    DistinctAttr &aliasClass = symbolClasses[symbol];
+    if (!aliasClass) {
+      if (!referenced)
+        referenced = UnitAttr::get(symbol->getContext());
+      aliasClass = DistinctAttr::create(referenced);
+    }
+    return aliasClass;
+  }
+
   DistinctAttr getSameOriginalClass(ValueRange values, StringRef debugLabel) {
     if (values.empty())
       return nullptr;
@@ -94,6 +107,7 @@ class OriginalClasses {
 
 private:
   DenseMap<Value, DistinctAttr> originalClasses;
+  DenseMap<Operation *, DistinctAttr> symbolClasses;
 };
 
 //===----------------------------------------------------------------------===//

enzyme/test/MLIR/AliasAnalysis/globals.mlir

@@ -0,0 +1,98 @@
+// RUN: %eopt --test-print-alias-analysis --split-input-file %s 2>&1 | FileCheck %s
+
+// The address of a global is a concrete alias class. Before, `llvm.mlir.addressof`
+// had no transfer function at all, so its lattice stayed undefined and any query
+// against it hit the "incomplete alias analysis" assertion.
+
+// Two distinct globals are distinct objects: taking the address of one can never
+// produce a pointer into the other, whatever their linkage.
+
+// CHECK: "a" and "b": NoAlias
+llvm.mlir.global external @ext1(0 : i64) : i64
+llvm.mlir.global external @ext2(0 : i64) : i64
+func.func @two_external_globals() {
+  %a = llvm.mlir.addressof @ext1 {tag = "a"} : !llvm.ptr
+  %b = llvm.mlir.addressof @ext2 {tag = "b"} : !llvm.ptr
+  return
+}
+
+// -----
+
+// CHECK: "a" and "b": NoAlias
+llvm.mlir.global internal @g1(0 : i64) : i64
+llvm.mlir.global internal @g2(0 : i64) : i64
+func.func @two_internal_globals() {
+  %a = llvm.mlir.addressof @g1 {tag = "a"} : !llvm.ptr
+  %b = llvm.mlir.addressof @g2 {tag = "b"} : !llvm.ptr
+  return
+}
+
+// -----
+
+// Mixed linkage is no different: still two separate objects.
+
+// CHECK: "a" and "b": NoAlias
+llvm.mlir.global external @pub(0 : i64) : i64
+llvm.mlir.global internal @priv(0 : i64) : i64
+func.func @mixed_linkage() {
+  %a = llvm.mlir.addressof @pub {tag = "a"} : !llvm.ptr
+  %b = llvm.mlir.addressof @priv {tag = "b"} : !llvm.ptr
+  return
+}
+
+// -----
+
+// Conversely, two addresses of the *same* global are the same object. The class
+// is keyed on the symbol rather than on the `llvm.mlir.addressof` result, which
+// matters because the LLVM importer emits one `addressof` per use.
+
+// CHECK: "a" and "b": MayAlias
+llvm.mlir.global external @same(0 : i64) : i64
+func.func @same_global_twice() {
+  %a = llvm.mlir.addressof @same {tag = "a"} : !llvm.ptr
+  %b = llvm.mlir.addressof @same {tag = "b"} : !llvm.ptr
+  return
+}
+
+// -----
+
+// A global keeps its own class across functions too.
+
+// CHECK: "a" and "b": MayAlias
+llvm.mlir.global external @shared(0 : i64) : i64
+func.func @user1() {
+  %a = llvm.mlir.addressof @shared {tag = "a"} : !llvm.ptr
+  return
+}
+func.func @user2() {
+  %b = llvm.mlir.addressof @shared {tag = "b"} : !llvm.ptr
+  return
+}
+
+// -----
+
+// An `llvm.mlir.addressof` of a function is not data storage; it stays in the
+// shared entry class when the function may be defined elsewhere.
+
+// CHECK: "a" and "b": MayAlias
+llvm.func external @decl()
+func.func @address_of_external_func() {
+  %a = llvm.mlir.addressof @decl {tag = "a"} : !llvm.ptr
+  %b = llvm.mlir.addressof @decl {tag = "b"} : !llvm.ptr
+  return
+}
+
+// -----
+
+// A global does not alias an unannotated pointer argument: argument classes and
+// global classes are disjoint. Note that this asserts no incoming pointer points
+// into a global, which a caller passing `&g` would violate. The analysis answered
+// NoAlias here before this file existed too, but only because the addressof side
+// was undefined rather than because anything had decided it.
+
+// CHECK: "arg" and "glob": NoAlias
+llvm.mlir.global external @escapes(0 : i64) : i64
+func.func @global_vs_argument(%arg0: !llvm.ptr {enzyme.tag = "arg"}) {
+  %glob = llvm.mlir.addressof @escapes {tag = "glob"} : !llvm.ptr
+  return
+}