[llvm] Add the "initializes" attribute inference (PR #97373)

[Jan Voung via llvm-commits]

================
@@ -580,6 +582,205 @@ struct ArgumentUsesTracker : public CaptureTracker {
   const SCCNodeSet &SCCNodes;
 };
 
+struct ArgumentUse {
+  Use *U;
+  std::optional<int64_t> Offset;
+};
+
+// A struct of argument access info. "Unknown" accesses are the cases like
+// unrecognized instructions, instructions that have more than one use of
+// the argument, or volatile memory accesses. "Unknown" implies "IsClobber"
+// and an empty access range.
+// Write or Read accesses can be clobbers as well for example, a Load with
+// scalable type.
+struct ArgumentAccessInfo {
+  enum AccessType { Write, Read, Unknown };
+  AccessType ArgAccessType;
+  ConstantRangeList AccessRanges;
+  bool IsClobber = false;
+};
+
+struct UsesPerBlockInfo {
+  DenseMap<Instruction *, ArgumentAccessInfo> Insts;
+  bool HasWrites;
+  bool HasClobber;
+};
+
+ArgumentAccessInfo GetArgmentAccessInfo(const Instruction *I,
+                                        const ArgumentUse &IU,
+                                        const DataLayout &DL) {
+  auto GetTypeAccessRange =
+      [&DL](Type *Ty,
+            std::optional<int64_t> Offset) -> std::optional<ConstantRange> {
+    auto TypeSize = DL.getTypeStoreSize(Ty);
+    if (!TypeSize.isScalable() && Offset.has_value()) {
+      int64_t Size = TypeSize.getFixedValue();
+      return ConstantRange(APInt(64, Offset.value(), true),
+                           APInt(64, Offset.value() + Size, true));
+    }
+    return std::nullopt;
+  };
+  auto GetConstantIntRange =
+      [](Value *Length,
+         std::optional<int64_t> Offset) -> std::optional<ConstantRange> {
+    auto *ConstantLength = dyn_cast<ConstantInt>(Length);
+    if (ConstantLength && Offset.has_value()) {
+      return ConstantRange(
+          APInt(64, Offset.value(), true),
+          APInt(64, Offset.value() + ConstantLength->getSExtValue(), true));
+    }
+    return std::nullopt;
+  };
+  if (auto *SI = dyn_cast<StoreInst>(I)) {
+    if (&SI->getOperandUse(1) == IU.U) {
+      // Get the fixed type size of "SI". Since the access range of a write
+      // will be unioned, if "SI" doesn't have a fixed type size, we just set
+      // the access range to empty.
+      ConstantRangeList AccessRanges;
+      auto TypeAccessRange = GetTypeAccessRange(SI->getAccessType(), IU.Offset);
+      if (TypeAccessRange.has_value())
+        AccessRanges.insert(TypeAccessRange.value());
+      return {ArgumentAccessInfo::AccessType::Write, AccessRanges,
+              /*IsClobber=*/false};
+    }
+  } else if (auto *LI = dyn_cast<LoadInst>(I)) {
+    if (&LI->getOperandUse(0) == IU.U) {
+      // Get the fixed type size of "LI". Different from Write, if "LI"
+      // doesn't have a fixed type size, we conservatively set as a clobber
+      // with an empty access range.
+      auto TypeAccessRange = GetTypeAccessRange(LI->getAccessType(), IU.Offset);
+      if (TypeAccessRange.has_value())
+        return {ArgumentAccessInfo::AccessType::Read,
+                {TypeAccessRange.value()},
+                /*IsClobber=*/false};
+      else
+        return {ArgumentAccessInfo::AccessType::Read, {}, /*IsClobber=*/true};
+    }
+  } else if (auto *MemSet = dyn_cast<MemSetInst>(I)) {
+    if (!MemSet->isVolatile()) {
+      ConstantRangeList AccessRanges;
+      auto AccessRange = GetConstantIntRange(MemSet->getLength(), IU.Offset);
+      if (AccessRange.has_value())
+        AccessRanges.insert(AccessRange.value());
+      return {ArgumentAccessInfo::AccessType::Write, AccessRanges,
+              /*IsClobber=*/false};
+    }
+  } else if (auto *MemCpy = dyn_cast<MemCpyInst>(I)) {
+    if (!MemCpy->isVolatile()) {
+      if (&MemCpy->getOperandUse(0) == IU.U) {
+        ConstantRangeList AccessRanges;
+        auto AccessRange = GetConstantIntRange(MemCpy->getLength(), IU.Offset);
+        if (AccessRange.has_value())
+          AccessRanges.insert(AccessRange.value());
+        return {ArgumentAccessInfo::AccessType::Write, AccessRanges,
+                /*IsClobber=*/false};
+      } else if (&MemCpy->getOperandUse(1) == IU.U) {
+        auto AccessRange = GetConstantIntRange(MemCpy->getLength(), IU.Offset);
+        if (AccessRange.has_value())
+          return {ArgumentAccessInfo::AccessType::Read,
+                  {AccessRange.value()},
+                  /*IsClobber=*/false};
+        else
+          return {ArgumentAccessInfo::AccessType::Read, {}, /*IsClobber=*/true};
+      }
+    }
+  } else if (auto *CB = dyn_cast<CallBase>(I)) {
+    if (CB->isArgOperand(IU.U)) {
+      unsigned ArgNo = CB->getArgOperandNo(IU.U);
+      bool IsInitialize = CB->paramHasAttr(ArgNo, Attribute::Initializes);
+      // Argument is only not clobbered when parameter is writeonly/readnone
+      // and nocapture.
+      bool IsClobber = !(CB->onlyWritesMemory(ArgNo) &&
+                         CB->paramHasAttr(ArgNo, Attribute::NoCapture));
+      ConstantRangeList AccessRanges;
+      if (IsInitialize && IU.Offset.has_value()) {
+        Attribute Attr = CB->getParamAttr(ArgNo, Attribute::Initializes);
+        if (!Attr.isValid()) {
+          Attr = CB->getCalledFunction()->getParamAttribute(
+              ArgNo, Attribute::Initializes);
+        }
+        ConstantRangeList CBCRL = Attr.getValueAsConstantRangeList();
+        for (ConstantRange &CR : CBCRL) {
+          AccessRanges.insert(ConstantRange(CR.getLower() + IU.Offset.value(),
+                                            CR.getUpper() + IU.Offset.value()));
+        }
+        return {ArgumentAccessInfo::AccessType::Write, AccessRanges, IsClobber};
+      }
+    }
+  }
+  // Unrecognized instructions are considered clobbers.
+  return {ArgumentAccessInfo::AccessType::Unknown, {}, /*IsClobber=*/true};
+}
+
+std::pair<bool, bool> CollectArgumentUsesPerBlock(
+    Argument &A, Function &F,
+    DenseMap<const BasicBlock *, UsesPerBlockInfo> &UsesPerBlock) {
+  auto &DL = F.getParent()->getDataLayout();
+  auto PointerSize =
+      DL.getIndexSizeInBits(A.getType()->getPointerAddressSpace());
+
+  bool HasAnyWrite = false;
+  bool HasWriteOutsideEntryBB = false;
+
+  BasicBlock &EntryBB = F.getEntryBlock();
+  SmallVector<ArgumentUse, 4> Worklist;
+  for (Use &U : A.uses())
+    Worklist.push_back({&U, 0});
+
+  auto UpdateUseInfo = [&UsesPerBlock](Instruction *I,
----------------
jvoung wrote:

brief comment for what the lambda's return value means?

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