[flang-commits] [flang] [flang] Promote scalar slots reached through fir.convert and fir.declare (PR #219314)

[Caroline Newcombe via flang-commits]

================
@@ -2135,6 +2136,131 @@ llvm::LogicalResult fir::ConvertOp::verify() {
          << getValue().getType() << " / " << getType();
 }
 
+/// Pointee of a reference to a simple scalar, or null. Both fir.ref and rank-0
+/// memref qualify, since the storage of a scalar is cast between those forms.
+static mlir::Type getScalarSlotPointeeType(mlir::Type type) {
+  mlir::Type eleTy;
+  if (auto refTy = mlir::dyn_cast<fir::ReferenceType>(type)) {
+    eleTy = refTy.getEleTy();
+  } else if (auto memrefTy = mlir::dyn_cast<mlir::MemRefType>(type)) {
+    // A rank, layout or memory space lets the cast reinterpret the storage.
+    if (memrefTy.getRank() != 0 || !memrefTy.getLayout().isIdentity() ||
+        memrefTy.getMemorySpace())
+      return {};
+    eleTy = memrefTy.getElementType();
+  } else {
+    return {};
+  }
+  if (!mlir::isa<mlir::IntegerType, mlir::FloatType, mlir::ComplexType,
+                 fir::LogicalType>(eleTy))
+    return {};
+  return eleTy;
+}
+
+/// The value a cast or declare aliases, or null if `value` is not one of those.
+static mlir::Value getAliasedSlotPointer(mlir::Value value) {
+  mlir::Operation *def = value.getDefiningOp();
+  if (auto convert = mlir::dyn_cast_or_null<fir::ConvertOp>(def))
+    return convert.getValue();
+  if (auto declare = mlir::dyn_cast_or_null<fir::DeclareOp>(def))
+    return declare.getMemref();
+  return {};
+}
+
+/// Checks whether every read of the slot reached through `pointer` has a write
+/// to it earlier in the same block. Otherwise promotion replaces the read with
+/// the allocation's default value, which is fir.undefined for fir.alloca but
+/// poison for a memref allocation, so promoting through a cast would change
+/// what reading an uninitialized variable does.
+static bool everyReadFollowsAWrite(mlir::Value pointer, mlir::Type pointee) {
+  // Start at the root of the chain, to see writes through a sibling alias.
+  while (mlir::Value aliased = getAliasedSlotPointer(pointer)) {
+    if (getScalarSlotPointeeType(aliased.getType()) != pointee)
+      break;
+    pointer = aliased;
+  }
+
+  llvm::SmallVector<mlir::Operation *> reads;
+  llvm::DenseMap<mlir::Block *, mlir::Operation *> firstWrite;
+  llvm::SmallVector<mlir::Value> worklist{pointer};
+  llvm::SmallPtrSet<mlir::Value, 8> visited{pointer};
+  while (!worklist.empty()) {
+    for (mlir::OpOperand &use : worklist.pop_back_val().getUses()) {
+      mlir::Operation *user = use.getOwner();
+      if (mlir::isa<fir::ConvertOp, fir::DeclareOp>(user)) {
+        mlir::Value alias = user->getResult(0);
+        // A different pointee is another slot, which blocks promotion itself.
+        if (getScalarSlotPointeeType(alias.getType()) == pointee &&
+            visited.insert(alias).second)
+          worklist.push_back(alias);
+        continue;
+      }
+      // A user that is not a promotable access blocks promotion itself.
+      if (auto memOp = mlir::dyn_cast<mlir::PromotableMemOpInterface>(user)) {
+        mlir::MemorySlot slot{use.get(), pointee};
+        if (memOp.storesTo(slot)) {
+          mlir::Operation *&earliest = firstWrite[user->getBlock()];
+          if (!earliest || user->isBeforeInBlock(earliest))
+            earliest = user;
+        } else if (memOp.loadsFrom(slot)) {
+          reads.push_back(user);
+        }
+      }
+    }
+  }
+
+  return llvm::all_of(reads, [&firstWrite](mlir::Operation *read) {
+    auto write = firstWrite.find(read->getBlock());
+    return write != firstWrite.end() && write->second->isBeforeInBlock(read);
+  });
----------------
cenewcombe wrote:

`everyReadFollowsAWrite` only credits a write that is in the same block as the read, so a variable whose write is in a different block — assigned under an if, updated in a loop — is rejected. Is that intentional? The undef-vs-poison rationale doesn't seem to require it: `@poison_insertion_point` in `mlir/test/Transforms/mem2reg.mlir` already promotes a conditionally-stored `memref.alloca` to a `ub.poison` block argument, which is the shape your new `@partially_written_slot_not_promoted` rejects.

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