[flang] Follow memory source through more operations (#66713)

Add support for fir.box_addr, fir.array_corr, fir.coordinate, fir.embox,
fir.rebox and fir.load.

1) Through the use of boolean `followBoxAddr` determine whether the
analysis should apply to the address of the box or the address wrapped
by the box.
2) Some asserts have been removed to allow for more SourceKinds though
the flow, in a particular SourceKind::Direct
3) getSource was a public method but the returned type (SourceKind) was
not public making it impossible to be called publicly
4) About 12 tests have been added to check for real Fortran scenarios 
5) More tests will be added with HLFIR
6) A few TODOs have been identified and will need to be addressed in
follow-up patches. I felt that more changes would increase the
complexity of the patch.

Author: Renaud-K

flang/include/flang/Optimizer/Analysis/AliasAnalysis.h

@@ -20,7 +20,7 @@ namespace fir {
 //===----------------------------------------------------------------------===//
 // AliasAnalysis
 //===----------------------------------------------------------------------===//
-class AliasAnalysis {
+struct AliasAnalysis {
   // Structures to describe the memory source of a value.
 
   /// Kind of the memory source referenced by a value.
@@ -36,11 +36,15 @@ class AliasAnalysis {
              /// Represents memory allocated outside of a function
              /// and passed to the function via host association tuple.
              HostAssoc,
+             /// Represents direct memory access whose source cannot be further
+             /// determined
+             Direct,
              /// Represents memory allocated by unknown means and
              /// with the memory address defined by a memory reading
              /// operation (e.g. fir::LoadOp).
              Indirect,
-             /// Represents memory allocated by unknown means.
+             /// Starting point to the analysis whereby nothing is known about
+             /// the source
              Unknown);
 
   /// Attributes of the memory source object.
@@ -81,7 +85,6 @@ class AliasAnalysis {
   friend llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                                        const AliasAnalysis::Source &op);
 
-public:
   /// Given two values, return their aliasing behavior.
   mlir::AliasResult alias(mlir::Value lhs, mlir::Value rhs);
 

flang/lib/Optimizer/Analysis/AliasAnalysis.cpp

@@ -36,6 +36,21 @@ static bool isDummyArgument(mlir::Value v) {
   return blockArg.getOwner()->isEntryBlock();
 }
 
+/// Temporary function to skip through all the no op operations
+/// TODO: Generalize support of fir.load
+static mlir::Value getOriginalDef(mlir::Value v) {
+  mlir::Operation *defOp;
+  bool breakFromLoop = false;
+  while (!breakFromLoop && (defOp = v.getDefiningOp())) {
+    llvm::TypeSwitch<Operation *>(defOp)
+        .Case<fir::ConvertOp>([&](fir::ConvertOp op) { v = op.getValue(); })
+        .Case<fir::DeclareOp, hlfir::DeclareOp>(
+            [&](auto op) { v = op.getMemref(); })
+        .Default([&](auto op) { breakFromLoop = true; });
+  }
+  return v;
+}
+
 namespace fir {
 
 void AliasAnalysis::Source::print(llvm::raw_ostream &os) const {
@@ -82,10 +97,26 @@ AliasResult AliasAnalysis::alias(Value lhs, Value rhs) {
 
   // Indirect case currently not handled. Conservatively assume
   // it aliases with everything
-  if (lhsSrc.kind == SourceKind::Indirect ||
-      lhsSrc.kind == SourceKind::Unknown ||
-      rhsSrc.kind == SourceKind::Indirect || rhsSrc.kind == SourceKind::Unknown)
+  if (lhsSrc.kind > SourceKind::Direct || rhsSrc.kind > SourceKind::Direct) {
     return AliasResult::MayAlias;
+  }
+
+  // SourceKind::Direct is set for the addresses wrapped in a global boxes.
+  // ie: fir.global @_QMpointersEp : !fir.box<!fir.ptr<f32>>
+  // Though nothing is known about them, they would only alias with targets or
+  // pointers
+  bool directSourceToNonTargetOrPointer = false;
+  if (lhsSrc.u != rhsSrc.u) {
+    if ((lhsSrc.kind == SourceKind::Direct && !rhsSrc.isTargetOrPointer()) ||
+        (rhsSrc.kind == SourceKind::Direct && !lhsSrc.isTargetOrPointer()))
+      directSourceToNonTargetOrPointer = true;
+  }
+
+  if (lhsSrc.kind == SourceKind::Direct ||
+      rhsSrc.kind == SourceKind::Direct) {
+    if (!directSourceToNonTargetOrPointer)
+      return AliasResult::MayAlias;
+  }
 
   if (lhsSrc.kind == rhsSrc.kind) {
     if (lhsSrc.u == rhsSrc.u) {
@@ -103,9 +134,6 @@ AliasResult AliasAnalysis::alias(Value lhs, Value rhs) {
         lhsSrc.kind == SourceKind::Global)
       return AliasResult::NoAlias;
 
-    assert(lhsSrc.kind == SourceKind::Argument &&
-           "unexpected memory source kind");
-
     // Dummy TARGET/POINTER arguments may alias.
     if (lhsSrc.isTargetOrPointer() && rhsSrc.isTargetOrPointer())
       return AliasResult::MayAlias;
@@ -122,7 +150,7 @@ AliasResult AliasAnalysis::alias(Value lhs, Value rhs) {
     return AliasResult::NoAlias;
   }
 
-  assert(lhsSrc.kind != rhsSrc.kind && "memory source kinds must be the same");
+  assert(lhsSrc.kind != rhsSrc.kind && "memory source kinds must be different");
 
   Source *src1, *src2;
   if (lhsSrc.kind < rhsSrc.kind) {
@@ -133,18 +161,6 @@ AliasResult AliasAnalysis::alias(Value lhs, Value rhs) {
     src2 = &lhsSrc;
   }
 
-  assert(src2->kind <= SourceKind::HostAssoc &&
-         "unexpected memory source kind");
-  if (src1->kind == SourceKind::Allocate)
-    return AliasResult::NoAlias;
-
-  assert(((src1->kind == SourceKind::Global &&
-           (src2->kind == SourceKind::Argument ||
-            src2->kind == SourceKind::HostAssoc)) ||
-          (src1->kind == SourceKind::Argument &&
-           src2->kind == SourceKind::HostAssoc)) &&
-         "unexpected memory source kinds");
-
   if (src1->kind == SourceKind::Argument &&
       src2->kind == SourceKind::HostAssoc) {
     // Treat the host entity as TARGET for the purpose of disambiguating
@@ -229,6 +245,7 @@ AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v) {
   mlir::Type ty;
   bool breakFromLoop{false};
   bool approximateSource{false};
+  bool followBoxAddr{false};
   mlir::SymbolRefAttr global;
   Source::Attributes attributes;
   while (defOp && !breakFromLoop) {
@@ -244,22 +261,74 @@ AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v) {
           v = op->getOperand(0);
           defOp = v.getDefiningOp();
         })
+        .Case<fir::BoxAddrOp>([&](auto op) {
+          v = op->getOperand(0);
+          defOp = v.getDefiningOp();
+          if (mlir::isa<fir::BaseBoxType>(v.getType()))
+            followBoxAddr = true;
+        })
+        .Case<fir::ArrayCoorOp, fir::CoordinateOp>([&](auto op) {
+          v = op->getOperand(0);
+          defOp = v.getDefiningOp();
+          if (mlir::isa<fir::BaseBoxType>(v.getType()))
+            followBoxAddr = true;
+          approximateSource = true;
+        })
+        .Case<fir::EmboxOp, fir::ReboxOp>([&](auto op) {
+          if (followBoxAddr) {
+            v = op->getOperand(0);
+            defOp = v.getDefiningOp();
+          } else
+            breakFromLoop = true;
+        })
         .Case<fir::LoadOp>([&](auto op) {
-          // No further tracking for addresses loaded from memory (e.g. a box)
-          // right now.
+          if (followBoxAddr && mlir::isa<fir::BaseBoxType>(op.getType())) {
+            // For now, support the load of an argument or fir.address_of
+            // TODO: generalize to all operations (in particular fir.alloca and
+            // fir.allocmem)
+            auto def = getOriginalDef(op.getMemref());
+            if (isDummyArgument(def) ||
+                def.template getDefiningOp<fir::AddrOfOp>()) {
+              v = def;
+              defOp = v.getDefiningOp();
+              return;
+            }
+          }
+          // No further tracking for addresses loaded from memory for now.
           type = SourceKind::Indirect;
           breakFromLoop = true;
         })
         .Case<fir::AddrOfOp>([&](auto op) {
           // Address of a global scope object.
-          type = SourceKind::Global;
           ty = v.getType();
+
+          // When the global is a
+          // fir.global @_QMpointersEp : !fir.box<!fir.ptr<f32>>
+          //   or
+          // fir.global @_QMpointersEp : !fir.box<!fir.heap<f32>>
+          //
+          // and when following through the wrapped address, capture
+          // the fact that there is nothing known about it. Therefore setting
+          // the source to Direct.
+          //
+          // When not following the wrapped address, then consider the address
+          // of the box, which has nothing to do with the wrapped address and
+          // lies in the global memory space.
+          if (followBoxAddr &&
+              mlir::isa<fir::BaseBoxType>(fir::unwrapRefType(ty)))
+            type = SourceKind::Direct;
+          else
+            type = SourceKind::Global;
+
           if (fir::valueHasFirAttribute(v,
                                         fir::GlobalOp::getTargetAttrNameStr()))
             attributes.set(Attribute::Target);
 
+          // TODO: Take followBoxAddr into account when setting the pointer
+          // attribute
           if (Source::isPointerReference(ty))
             attributes.set(Attribute::Pointer);
+
           global = llvm::cast<fir::AddrOfOp>(op).getSymbol();
           breakFromLoop = true;
         })
@@ -278,7 +347,7 @@ AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v) {
             breakFromLoop = true;
             return;
           }
-
+          // TODO: Look for the fortran attributes present on the operation
           // Track further through the operand
           v = op.getMemref();
           defOp = v.getDefiningOp();

flang/test/Analysis/AliasAnalysis/alias-analysis-2.fir

@@ -11,8 +11,11 @@
 // p1.addr and p2.addr could both be wrapped inside boxes
 // CHECK-DAG: p1.addr#0 <-> boxp1.addr#0: MayAlias
 // CHECK-DAG: p2.addr#0 <-> boxp1.addr#0: MayAlias
-// CHECK-DAG: p1.addr#0 <-> arg2.addr#0: MayAlias
-// CHECK-DAG: p2.addr#0 <-> arg2.addr#0: MayAlias
+
+// TODO: To really see aliasing, we should be looking at a load of p1.addr
+// p1.addr is just a local address holding the address to the data 
+// CHECK-DAG: p1.addr#0 <-> arg2.addr#0: NoAlias
+// CHECK-DAG: p2.addr#0 <-> arg2.addr#0: NoAlias
 
 // p1.addr and p2.addr are the result of an allocation
 // They cannot physically alias with an argument
@@ -41,7 +44,7 @@
 // pointer arguments
 // CHECK-DAG: arg2.addr#0 <-> func.region0#0: MayAlias
 // CHECK-DAG: arg2.addr#0 <-> func.region0#1: MayAlias
-// CHECK-DAG: arg2.addr#0 <-> func.region0#2: MayAlias
+// CHECK-DAG: arg2.addr#0 <-> func.region0#2: MustAlias
 // CHECK-DAG: boxp1.addr#0 <-> arg2.addr#0: MayAlias
 
 func.func @_QFPtest(%arg0: !fir.ref<f32> {fir.bindc_name = "v1", fir.target}, %arg1: !fir.ref<f32> {fir.bindc_name = "v2", fir.target}, %arg2: !fir.ref<!fir.box<!fir.ptr<f32>>> ) attributes {test.ptr = "func"} {
@@ -99,7 +102,7 @@ func.func @_QFPtest(%arg0: !fir.ref<f32> {fir.bindc_name = "v1", fir.target}, %a
 // CHECK-DAG: func.region0#0 <-> func.region0#1: MayAlias
 // CHECK-DAG: func.region0#0 <-> func.region0#2: MayAlias
 
-func.func @_QFPtest2(%arg0: !fir.ref<f32> {fir.bindc_name = "v1", fir.target}, %arg1: !fir.ref<!fir.box<!fir.ptr<f32>>>, %arg2: !fir.ref<!fir.ptr<f32>> ) attributes {test.ptr = "func"} {
+func.func @_QFPtest2(%arg0: !fir.ref<f32> {fir.bindc_name = "v1", fir.target}, %arg1: !fir.ref<!fir.box<!fir.ptr<f32>>>, %arg2: !fir.ref<!fir.box<!fir.ptr<f32>>> ) attributes {test.ptr = "func"} {
   return
 }
 
@@ -131,24 +134,29 @@ func.func @_QFPtest2(%arg0: !fir.ref<f32> {fir.bindc_name = "v1", fir.target}, %
 // CHECK-DAG: p#0 <-> func.region0#0: MayAlias
 // CHECK-DAG: p#0 <-> func.region0#1: NoAlias
 
-// FIXME: p and p1 are pointers, they cannot alias with a wrapped address.
-//        Only the addresses they wrap could alias with the address wrapped by the box
-// CHECK-DAG: p#0 <-> box.addr#0: MayAlias
+// p could be pointing to var2 
+// var2, being a target, could also be passed as argument arg0 
+
+// This was the wrong question to ask. We are asking if the address of box _QMpointersEp can 
+// alias with the wrapped scalar _QFEvar2. We meant box_addr of _QMpointersEp 
+// CHECK-DAG: p#0 <-> box.addr#0: NoAlias
+
+// TODO: Still need to handle more gracefully the difference between !fir.ref<!fir.box<>> and !fir.box<> 
 // CHECK-DAG: box.addr#0 <-> func.region0#0: MayAlias
 
 // var2, although it is a target, cannot alias with p
 // A modification of p would only make them point to a new target but not modify it
 // CHECK-DAG: var2#0 <-> p#0: NoAlias
 // It can alias with p1, if p1 is a pointer component
 // CHECK-DAG: var2#0 <-> func.region0#0: MayAlias
-// It can alias with a box.addr
-// CHECK-DAG: var2#0 <-> box.addr#0: MayAlias
+// It is the same as box.addr
+// CHECK-DAG: var2#0 <-> box.addr#0: MustAlias
 
 // A global may not alias with a dummy
 // CHECK-DAG: var2#0 <-> func.region0#1: NoAlias
 
-// FIXME: a pointer may only alias with a target but arg1 is a regular dummy
-// CHECK-DAG: box.addr#0 <-> func.region0#1: MayAlias
+// A pointer may only alias with a target but arg1 is a regular dummy
+// CHECK-DAG: box.addr#0 <-> func.region0#1: NoAlias
 
 // Dummy argument do not alias
 // CHECK-DAG: func.region0#0 <-> func.region0#1: NoAlias