[llvm] [DSE] Apply initializes attribute to DSE (PR #107282)

[Nikita Popov via llvm-commits]

================
@@ -820,20 +828,126 @@ struct MemoryLocationWrapper {
   const Value *UnderlyingObject;
   MemoryDef *MemDef;
   Instruction *DefInst;
+  bool DefByInitializesAttr = false;
 };
 
 // A memory def wrapper that represents a MemoryDef and the MemoryLocation(s)
 // defined by this MemoryDef.
 struct MemoryDefWrapper {
-  MemoryDefWrapper(MemoryDef *MemDef, std::optional<MemoryLocation> MemLoc) {
+  MemoryDefWrapper(
+      MemoryDef *MemDef,
+      const SmallVectorImpl<std::pair<MemoryLocation, bool>> &MemLocations) {
     DefInst = MemDef->getMemoryInst();
-    if (MemLoc.has_value())
-      DefinedLocation = MemoryLocationWrapper(*MemLoc, MemDef);
+    for (auto &[MemLoc, DefByInitializesAttr] : MemLocations)
+      DefinedLocations.push_back(
+          MemoryLocationWrapper(MemLoc, MemDef, DefByInitializesAttr));
   }
   Instruction *DefInst;
-  std::optional<MemoryLocationWrapper> DefinedLocation = std::nullopt;
+  SmallVector<MemoryLocationWrapper, 1> DefinedLocations;
+};
+
+bool hasInitializesAttr(Instruction *I) {
+  CallBase *CB = dyn_cast<CallBase>(I);
+  if (!CB)
+    return false;
+
+  for (unsigned Idx = 0, Count = CB->arg_size(); Idx < Count; ++Idx)
+    if (CB->paramHasAttr(Idx, Attribute::Initializes))
+      return true;
+  return false;
+}
+
+struct ArgumentInitInfo {
+  unsigned Idx;
+  bool HasDeadOnUnwindAttr;
+  ConstantRangeList Inits;
 };
 
+// Return the intersected range list of the initializes attributes of "Args".
+// "Args" are call arguments that alias to each other.
+// If any argument in "Args" doesn't have dead_on_unwind attr and
+// "FuncHasNoUnwindAttr" is false, return empty.
+ConstantRangeList
+getIntersectedInitRangeList(const SmallVectorImpl<ArgumentInitInfo> &Args,
+                            bool FuncHasNoUnwindAttr) {
+  if (Args.empty())
+    return {};
+
+  // To address unwind, the function should have nounwind attribute or the
+  // arguments have dead_on_unwind attribute. Otherwise, return empty.
+  for (const auto &Arg : Args) {
+    if (!FuncHasNoUnwindAttr && !Arg.HasDeadOnUnwindAttr)
+      return {};
+    if (Arg.Inits.empty())
+      return {};
+  }
+
+  if (Args.size() == 1)
+    return Args[0].Inits;
+
+  ConstantRangeList IntersectedIntervals = Args[0].Inits;
+  for (size_t I = 1, Count = Args.size(); I < Count; ++I)
+    IntersectedIntervals = IntersectedIntervals.intersectWith(Args[I].Inits);
+
+  return IntersectedIntervals;
+}
+
+// Return the locations wrote by the initializes attribute.
+// Note that this function considers:
+// 1. Unwind edge: apply "initializes" attribute only if the callee has
+//    "nounwind" attribute or the argument has "dead_on_unwind" attribute.
+// 2. Argument alias: for aliasing arguments, the "initializes" attribute is
+//    the intersected range list of their "initializes" attributes.
+SmallVector<MemoryLocation, 1>
+getInitializesArgMemLoc(const Instruction *I, BatchAAResults &BatchAA) {
+  const CallBase *CB = dyn_cast<CallBase>(I);
+  if (!CB)
+    return {};
+
+  // Collect aliasing arguments and their initializes ranges.
+  SmallMapVector<Value *, SmallVector<ArgumentInitInfo, 2>, 2> Arguments;
+  for (unsigned Idx = 0, Count = CB->arg_size(); Idx < Count; ++Idx) {
+    ConstantRangeList Inits;
+    if (CB->paramHasAttr(Idx, Attribute::Initializes))
+      Inits = CB->getParamAttr(Idx, Attribute::Initializes)
+                  .getValueAsConstantRangeList();
+
+    bool HasDeadOnUnwindAttr = CB->paramHasAttr(Idx, Attribute::DeadOnUnwind);
+    ArgumentInitInfo InitInfo{Idx, HasDeadOnUnwindAttr, Inits};
+    Value *CurArg = CB->getArgOperand(Idx);
+    bool FoundAliasing = false;
+    for (auto &[Arg, AliasList] : Arguments) {
+      if (BatchAA.isMustAlias(Arg, CurArg)) {
----------------
nikic wrote:

Just handling them the same as MustAlias isn't correct. If it's PartialAlias then there is an offset between the arguments, so initializes refers to different offsets. And if it's MayAlias, then there may be an unknown offset, or the arguments may be unrelated entirely.

For the PartialAlias/MayAlias cases we should discard the initializes information entirely.

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