[StackColoring] Update AliasAnalysis information in stack coloring pass

[inouehrs]

Summary

Stack coloring pass need to maintain TBAA metadata when merging stack slots of different types.
Actually, there is a FIXME comment in StackColoring.cpp

// FIXME: In order to enable the use of TBAA when using AA in CodeGen,
// we'll also need to update the TBAA nodes in MMOs with values
// derived from the merged allocas.

But, TBAA has been already enabled in CodeGen without fixing this pass.
The incorrect TBAA metadata results in recent failures in bootstrap test on ppc64le (PR33928) by allowing unsafe instruction scheduling.
Although we observed the problem on ppc64le, this is a platform neutral issue.

This patch makes the stack coloring pass maintains TBAA metadata (somewhat conservatively) when merging multiple stack slots.

This patch fixes PR33928.

[grandinj]

include/llvm/CodeGen/MachineFunction.h

665

AliasAnalisys -> AliasAnalysis

[jtony]

lib/CodeGen/MachineFunction.cpp

338

If we add a temp variable for the first parameter and assign a different value to it according to the guard if(MMO->getValue()), won't the code be simpler? I am talking something like this:

const PseudoSourceValue *Value = MMO->getValue() ?  MMO->getValue() : MMO->getPseudoValue();

return new (Allocator)
             MachineMemOperand(MachinePointerInfo(Value, MMO->getOffset()),
                               MMO->getFlags(), MMO->getSize(),
                               MMO->getBaseAlignment(), AAInfo, 
                               MMO->getRanges(), MMO->getSyncScopeID(),
                               MMO->getOrdering(), MMO->getFailureOrdering());
}

lib/CodeGen/StackColoring.cpp

874

Can we get rid of the static key word for these two functions and and put their declaration into a header and just keep one copy of the source code?

[echristo]

Tagging in Matthias here.

[inouehrs]

• avoid eliminating non-TBAA alias analysis information
• minor touchups

[inouehrs]

include/llvm/CodeGen/MachineFunction.h

665

Fixed. Thanks!

lib/CodeGen/MachineFunction.cpp

338

I hope this looks better than the old code.
Since getValue() and getPseudoValue() return different types, I think this is easier to understand.

lib/CodeGen/StackColoring.cpp

874

Done.

[MatzeB]

I don't know much about our alias analysis infrastructure.

That said the description and changes seem reasonable. Please double check your getUnderlyingObject refactoring before submitting.

include/llvm/CodeGen/MachineFunction.h

664

Don't repeat the method name in the doxygen comment (old code does it, new code shouldn't).

include/llvm/CodeGen/ScheduleDAGInstrs.h

86 ↗

(On Diff #108420)

Explain more what the function does and less what it should be used for.

Also this isn't scheduler related, so at least move it to MachineMemoryOperand or better to ValueTracking.h (which already has a few variations of it, maybe you can even use one of those instead of roling another one).

lib/CodeGen/StackColoring.cpp

40

Instead of documenting why you need to import a scheduling header here, just put the function into an apropriate header (as above) :)

1023

Is MergedAllocas changing here? Looks like you can move the check further up outside this loop.

[hfinkel]

We need to be very careful to document the contracts here and explain what's going on. The implementation of getUnderlyingObjects in ScheduleDAGInstrs.cpp calls its helper function getUnderlyingObjectFromInt which contains this comment:

// If we find an add of a constant, a multiplied value, or a phi, it's
// likely that the other operand will lead us to the base
// object. We don't have to worry about the case where the
// object address is somehow being computed by the multiply,
// because our callers only care when the result is an
// identifiable object.

and this is indeed true. The only existing caller, getUnderlyingObjectsForInstr, does this:

} else if (const Value *V = MMO->getValue()) {
  SmallVector<Value *, 4> Objs;
  getUnderlyingObjects(V, Objs, DL);

  for (Value *V : Objs) {
    if (!isIdentifiedObject(V))
      return false;

    Objects.push_back(UnderlyingObjectsVector::value_type(V, true));
  }

and this is important because, if GetUnderlyingObjects hits its depth cutoff, it can return arbitrary instructions in the use/def chain of the pointer.

If we're going to make this function a utility that can be used elsewhere, we should make this safer. I'd advocate for moving the isIdentifiedObject check into the implementation of getUnderlyingObjects before:

Objects.push_back(const_cast<Value *>(V));

such that if the check fails we clear the Objects array and return. That should make the interface safe. We should document this part of the functions behavior and remove the corresponding check in getUnderlyingObjectsForInstr.

lib/CodeGen/ScheduleDAGInstrs.cpp

123–124

Please rename this function if you make it non-static. The function we have right now in ValueTracking is called GetUnderlyingObjects, but if we were to rename it follow our coding conventions we'd have a conflict with the name of function function. How about getUnderlyingObjectsForCodeGen?

lib/CodeGen/StackColoring.cpp

917

I'd remove this FIXME. It is not clear to me that there is a well-defined way to do this in general. A single alloca could be accessed at different times by accesses with incompatible TBAA metadata so long as some access compatible with both, or data dependence, keeps the correct ordering. As a result, there's not necessarily a single thing to merge the metadata with. I suppose merging everything (derived from one) with everything (derived from the other) might work, but in any event, it seems non-trivial. I'm happy to have someone revisit this if they find this needs to be more precise.

1003

This is not specific to TBAA. You should replace TBAA with AA metadata in all of the comments.

1012

I don't understand this name. How about MayHaveConflictingAAMD?

1032

This situation is not specific to TBAA. If affects all of the AA metadata. You'll need to unset it all.

[qcolombet]

+Hal

I think Hal would like to double check that.

[qcolombet]

Haha, never mind me, Hal already replied. (Sloppy tab!)

[inouehrs]

• move getUnderlyingObjects from ScheduleDAGInstrs.cpp to ValueTracking.cpp and also renamed to getUnderlyingObjectsForCodeGen.
• check the validity of Values returned from getUnderlyingObjectsForCodeGen.
• remove information on all types of AA as well as information on TBAA.

[inouehrs]

include/llvm/CodeGen/MachineFunction.h

664

Fixed.

lib/CodeGen/ScheduleDAGInstrs.cpp

123–124

I renamed it getUnderlyingObjectsForCodeGen and moved into ValueTracking.cpp.

lib/CodeGen/StackColoring.cpp

40

Moved to ValueTracking.h.

917

I agree. I removed this FIXME.

1003

Done.

1012

Fixed.

1023

We need to check each Value V (and hence AI) with MergedAllocas. So I do this check in the loop.

1032

Done.

[inouehrs]

• minor fix

[inouehrs]

Move check with isIdentifiedObject into getUnderlyingObjectsForCodeGen

[inouehrs]

@hfinkel I am going to commit this soon to fix the buildbot failures. Please let me know if you have further comments. Thanks!

[inouehrs]

committed: https://reviews.llvm.org/rL309651

[MatzeB]

lib/CodeGen/StackColoring.cpp

1033–1034

• Does AAMDNodes() mean it can alias with anything? In that case you could simply use I.dropMemRefs() as we make the same assumption for an instruction without any memory operands.

• Can't you simplify to MemOps[MemOpIdx++] to *MMO = ... and avoid the extra counter?

[MatzeB]

lib/CodeGen/StackColoring.cpp

1033–1034

Ignore my last point about the MemOps. I missed that fact that you are duplicating the information and not just replacing the existing operands (may help understand if you renamed the variable to NewMemOps).

[inouehrs]

updated based on post-commit comments

[inouehrs]

lib/Analysis/ValueTracking.cpp

3332

Comments from Hal

I don't think this check does that you want because, if V is a inttoptr,
this will always be false (i.e. the check for inttoptr will be dead).

To preserve functionality, I believe you need the check here instead.

Yes, exactly. I moved the check with isIdentifiedObject() after check of IntToPtr.

lib/CodeGen/StackColoring.cpp

1033–1034

• AAMDNodes() means no AA information and so optimizers think this instruction can alias with anything. I replace only AAInfo and other information in MMO (e.g. pointer, size etc) are still needed. So I eliminate only AAInfo.
• I renamed MemOps to NewMemOps for readability. Thank you for the suggestion.

[MatzeB]

Thanks, LGTM

[inouehrs]

Additional patch committed in https://reviews.llvm.org/rL309849

include/llvm/Analysis/ValueTracking.h

template <typename T> class ArrayRef;
   ///
   /// Since A[i] and A[i-1] are independent pointers, getUnderlyingObjects
   /// should not assume that Curr and Prev share the same underlying object thus
   /// it shouldn't look through the phi above.
   void GetUnderlyingObjects(Value *V, SmallVectorImpl<Value *> &Objects,
                             const DataLayout &DL, LoopInfo *LI = nullptr,
                             unsigned MaxLookup = 6);
 
+  /// This is a wrapper around GetUnderlyingObjects and adds support for basic
+  /// ptrtoint+arithmetic+inttoptr sequences.
+  /// Caller of this method should check validity of returned Values by
+  /// llvm::isIdentifiedObject(Value*) as necessary.
+  void getUnderlyingObjectsForCodeGen(const Value *V,
+                            SmallVectorImpl<Value *> &Objects,
+                            const DataLayout &DL);
+
   /// Return true if the only users of this pointer are lifetime markers.
   bool onlyUsedByLifetimeMarkers(const Value *V);
 
   /// Return true if the instruction does not have any effects besides
   /// calculating the result and does not have undefined behavior.
   ///
   /// This method never returns true for an instruction that returns true for
   /// mayHaveSideEffects; however, this method also does some other checks in

include/llvm/CodeGen/MachineFunction.h

public:
 
   /// getMachineMemOperand - Allocate a new MachineMemOperand by copying
   /// an existing one, adjusting by an offset and using the given size.
   /// MachineMemOperands are owned by the MachineFunction and need not be
   /// explicitly deallocated.
   MachineMemOperand *getMachineMemOperand(const MachineMemOperand *MMO,
                                           int64_t Offset, uint64_t Size);
 
+  /// Allocate a new MachineMemOperand by copying an existing one,
+  /// replacing only AliasAnalysis information. MachineMemOperands are owned
+  /// by the MachineFunction and need not be explicitly deallocated.
+  MachineMemOperand *getMachineMemOperand(const MachineMemOperand *MMO,
+                                          const AAMDNodes &AAInfo);
+
   using OperandCapacity = ArrayRecycler<MachineOperand>::Capacity;
 
   /// Allocate an array of MachineOperands. This is only intended for use by
   /// internal MachineInstr functions.
   MachineOperand *allocateOperandArray(OperandCapacity Cap) {
     return OperandRecycler.allocate(Cap, Allocator);
   }
 

include/llvm/CodeGen/MachineInstr.h

iterator_range<mmo_iterator> memoperands() const {
     return make_range(memoperands_begin(), memoperands_end());
   }
 
   /// Return true if this instruction has exactly one MachineMemOperand.
   bool hasOneMemOperand() const {
     return NumMemRefs == 1;
   }
 
+  /// Return the number of memory operands.
+  unsigned getNumMemOperands() const { return NumMemRefs; }
+
   /// API for querying MachineInstr properties. They are the same as MCInstrDesc
   /// queries but they are bundle aware.
 
   enum QueryType {
     IgnoreBundle,    // Ignore bundles
     AnyInBundle,     // Return true if any instruction in bundle has property
     AllInBundle      // Return true if all instructions in bundle have property
   };

lib/Analysis/ValueTracking.cpp

if (PHINode *PN = dyn_cast<PHINode>(P)) {
           Worklist.push_back(IncValue);
       continue;
     }
 
     Objects.push_back(P);
   } while (!Worklist.empty());
 }
 
+/// This is the function that does the work of looking through basic
+/// ptrtoint+arithmetic+inttoptr sequences.
+static const Value *getUnderlyingObjectFromInt(const Value *V) {
+  do {
+    if (const Operator *U = dyn_cast<Operator>(V)) {
+      // If we find a ptrtoint, we can transfer control back to the
+      // regular getUnderlyingObjectFromInt.
+      if (U->getOpcode() == Instruction::PtrToInt)
+        return U->getOperand(0);
+      // If we find an add of a constant, a multiplied value, or a phi, it's
+      // likely that the other operand will lead us to the base
+      // object. We don't have to worry about the case where the
+      // object address is somehow being computed by the multiply,
+      // because our callers only care when the result is an
+      // identifiable object.
+      if (U->getOpcode() != Instruction::Add ||
+          (!isa<ConstantInt>(U->getOperand(1)) &&
+           Operator::getOpcode(U->getOperand(1)) != Instruction::Mul &&
+           !isa<PHINode>(U->getOperand(1))))
+        return V;
+      V = U->getOperand(0);
+    } else {
+      return V;
+    }
+    assert(V->getType()->isIntegerTy() && "Unexpected operand type!");
+  } while (true);
+}
+
+/// This is a wrapper around GetUnderlyingObjects and adds support for basic
+/// ptrtoint+arithmetic+inttoptr sequences.
+/// Caller of this method should check validity of returned Values by
+/// llvm::isIdentifiedObject(Value*) as necessary.
+void llvm::getUnderlyingObjectsForCodeGen(const Value *V,
+                          SmallVectorImpl<Value *> &Objects,
+                          const DataLayout &DL) {
+  SmallPtrSet<const Value *, 16> Visited;
+  SmallVector<const Value *, 4> Working(1, V);
+  do {
+    V = Working.pop_back_val();
+
+    SmallVector<Value *, 4> Objs;
+    GetUnderlyingObjects(const_cast<Value *>(V), Objs, DL);
+
+    for (Value *V : Objs) {
+      if (!Visited.insert(V).second)
+        continue;
+      if (Operator::getOpcode(V) == Instruction::IntToPtr) {
+        const Value *O =
+          getUnderlyingObjectFromInt(cast<User>(V)->getOperand(0));
+        if (O->getType()->isPointerTy()) {
+          Working.push_back(O);
+          continue;
+        }
+      }
+      Objects.push_back(const_cast<Value *>(V));
+    }
+  } while (!Working.empty());
+}
+
 /// Return true if the only users of this pointer are lifetime markers.
 bool llvm::onlyUsedByLifetimeMarkers(const Value *V) {
   for (const User *U : V->users()) {
     const IntrinsicInst *II = dyn_cast<IntrinsicInst>(U);
     if (!II) return false;
 
     if (II->getIntrinsicID() != Intrinsic::lifetime_start &&
         II->getIntrinsicID() != Intrinsic::lifetime_end)

lib/CodeGen/MachineFunction.cpp

MachineFunction::getMachineMemOperand(const MachineMemOperand *MMO,
   return new (Allocator)
              MachineMemOperand(MachinePointerInfo(MMO->getPseudoValue(),
                                                   MMO->getOffset()+Offset),
                                MMO->getFlags(), Size, MMO->getBaseAlignment(),
                                AAMDNodes(), nullptr, MMO->getSyncScopeID(),
                                MMO->getOrdering(), MMO->getFailureOrdering());
 }
 
+MachineMemOperand *
+MachineFunction::getMachineMemOperand(const MachineMemOperand *MMO,
+                                      const AAMDNodes &AAInfo) {
+  MachinePointerInfo MPI = MMO->getValue() ?
+             MachinePointerInfo(MMO->getValue(), MMO->getOffset()) :
+             MachinePointerInfo(MMO->getPseudoValue(), MMO->getOffset());
+
+  return new (Allocator)
+             MachineMemOperand(MPI, MMO->getFlags(), MMO->getSize(),
+                               MMO->getBaseAlignment(), AAInfo,
+                               MMO->getRanges(), MMO->getSyncScopeID(),
+                               MMO->getOrdering(), MMO->getFailureOrdering());
+}
+
 MachineInstr::mmo_iterator
 MachineFunction::allocateMemRefsArray(unsigned long Num) {
   return Allocator.Allocate<MachineMemOperand *>(Num);
 }
 
 std::pair<MachineInstr::mmo_iterator, MachineInstr::mmo_iterator>
 MachineFunction::extractLoadMemRefs(MachineInstr::mmo_iterator Begin,
                                     MachineInstr::mmo_iterator End) {

lib/CodeGen/ScheduleDAGInstrs.cpp

: ScheduleDAG(mf), MLI(mli), MFI(mf.getFrameInfo()),
       RemoveKillFlags(RemoveKillFlags),
       UnknownValue(UndefValue::get(
                              Type::getVoidTy(mf.getFunction()->getContext()))) {
   DbgValues.clear();
 
   const TargetSubtargetInfo &ST = mf.getSubtarget();
   SchedModel.init(ST.getSchedModel(), &ST, TII);
 }
 
-/// This is the function that does the work of looking through basic
-/// ptrtoint+arithmetic+inttoptr sequences.
-static const Value *getUnderlyingObjectFromInt(const Value *V) {
-  do {
-    if (const Operator *U = dyn_cast<Operator>(V)) {
-      // If we find a ptrtoint, we can transfer control back to the
-      // regular getUnderlyingObjectFromInt.
-      if (U->getOpcode() == Instruction::PtrToInt)
-        return U->getOperand(0);
-      // If we find an add of a constant, a multiplied value, or a phi, it's
-      // likely that the other operand will lead us to the base
-      // object. We don't have to worry about the case where the
-      // object address is somehow being computed by the multiply,
-      // because our callers only care when the result is an
-      // identifiable object.
-      if (U->getOpcode() != Instruction::Add ||
-          (!isa<ConstantInt>(U->getOperand(1)) &&
-           Operator::getOpcode(U->getOperand(1)) != Instruction::Mul &&
-           !isa<PHINode>(U->getOperand(1))))
-        return V;
-      V = U->getOperand(0);
-    } else {
-      return V;
-    }
-    assert(V->getType()->isIntegerTy() && "Unexpected operand type!");
-  } while (true);
-}
-
-/// This is a wrapper around GetUnderlyingObjects and adds support for basic
-/// ptrtoint+arithmetic+inttoptr sequences.
-static void getUnderlyingObjects(const Value *V,
-                                 SmallVectorImpl<Value *> &Objects,
-                                 const DataLayout &DL) {
-  SmallPtrSet<const Value *, 16> Visited;
-  SmallVector<const Value *, 4> Working(1, V);
-  do {
-    V = Working.pop_back_val();
-
-    SmallVector<Value *, 4> Objs;
-    GetUnderlyingObjects(const_cast<Value *>(V), Objs, DL);
-
-    for (Value *V : Objs) {
-      if (!Visited.insert(V).second)
-        continue;
-      if (Operator::getOpcode(V) == Instruction::IntToPtr) {
-        const Value *O =
-          getUnderlyingObjectFromInt(cast<User>(V)->getOperand(0));
-        if (O->getType()->isPointerTy()) {
-          Working.push_back(O);
-          continue;
-        }
-      }
-      Objects.push_back(const_cast<Value *>(V));
-    }
-  } while (!Working.empty());
-}
-
 /// If this machine instr has memory reference information and it can be tracked
 /// to a normal reference to a known object, return the Value for that object.
 static void getUnderlyingObjectsForInstr(const MachineInstr *MI,
                                          const MachineFrameInfo &MFI,
                                          UnderlyingObjectsVector &Objects,
                                          const DataLayout &DL) {
   auto allMMOsOkay = [&]() {
     for (const MachineMemOperand *MMO : MI->memoperands()) {
       if (MMO->isVolatile())
for (const MachineMemOperand *MMO : MI->memoperands()) {
         // such aliases.
         if (PSV->isAliased(&MFI))
           return false;
 
         bool MayAlias = PSV->mayAlias(&MFI);
         Objects.push_back(UnderlyingObjectsVector::value_type(PSV, MayAlias));
       } else if (const Value *V = MMO->getValue()) {
         SmallVector<Value *, 4> Objs;
-        getUnderlyingObjects(V, Objs, DL);
+        getUnderlyingObjectsForCodeGen(V, Objs, DL);
 
         for (Value *V : Objs) {
           if (!isIdentifiedObject(V))
             return false;
 
           Objects.push_back(UnderlyingObjectsVector::value_type(V, true));
         }
       } else

lib/CodeGen/StackColoring.cpp

 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/PseudoSourceValue.h"
+#include "llvm/CodeGen/SelectionDAGNodes.h"
 #include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/CodeGen/StackProtector.h"
 #include "llvm/CodeGen/WinEHFuncInfo.h"
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
for (MachineInstr *MI : Markers) {
     Count++;
   }
   Markers.clear();
 
   DEBUG(dbgs()<<"Removed "<<Count<<" markers.\n");
   return Count;
 }
 
 void StackColoring::remapInstructions(DenseMap<int, int> &SlotRemap) {
   unsigned FixedInstr = 0;
   unsigned FixedMemOp = 0;
   unsigned FixedDbg = 0;
 
   // Remap debug information that refers to stack slots.
   for (auto &VI : MF->getVariableDbgInfo()) {
     if (!VI.Var)
       continue;
     if (SlotRemap.count(VI.Slot)) {
       DEBUG(dbgs() << "Remapping debug info for ["
                    << cast<DILocalVariable>(VI.Var)->getName() << "].\n");
       VI.Slot = SlotRemap[VI.Slot];
       FixedDbg++;
     }
   }
 
   // Keep a list of *allocas* which need to be remapped.
   DenseMap<const AllocaInst*, const AllocaInst*> Allocas;
+
+  // Keep a list of allocas which has been affected by the remap.
+  SmallPtrSet<const AllocaInst*, 32> MergedAllocas;
+
   for (const std::pair<int, int> &SI : SlotRemap) {
     const AllocaInst *From = MFI->getObjectAllocation(SI.first);
     const AllocaInst *To = MFI->getObjectAllocation(SI.second);
     assert(To && From && "Invalid allocation object");
     Allocas[From] = To;
 
     // AA might be used later for instruction scheduling, and we need it to be
     // able to deduce the correct aliasing releationships between pointers
     // derived from the alloca being remapped and the target of that remapping.
     // The only safe way, without directly informing AA about the remapping
     // somehow, is to directly update the IR to reflect the change being made
     // here.
     Instruction *Inst = const_cast<AllocaInst *>(To);
     if (From->getType() != To->getType()) {
       BitCastInst *Cast = new BitCastInst(Inst, From->getType());
       Cast->insertAfter(Inst);
       Inst = Cast;
     }
 
+    // We keep both slots to maintain AliasAnalysis metadata later.
+    MergedAllocas.insert(From);
+    MergedAllocas.insert(To);
+
     // Allow the stack protector to adjust its value map to account for the
     // upcoming replacement.
     SP->adjustForColoring(From, To);
 
     // The new alloca might not be valid in a llvm.dbg.declare for this
     // variable, so undef out the use to make the verifier happy.
     AllocaInst *FromAI = const_cast<AllocaInst *>(From);
     if (FromAI->isUsedByMetadata())
for (MachineBasicBlock &BB : *MF)
     for (MachineInstr &I : BB) {
       // Skip lifetime markers. We'll remove them soon.
       if (I.getOpcode() == TargetOpcode::LIFETIME_START ||
           I.getOpcode() == TargetOpcode::LIFETIME_END)
         continue;
 
       // Update the MachineMemOperand to use the new alloca.
       for (MachineMemOperand *MMO : I.memoperands()) {
-        // FIXME: In order to enable the use of TBAA when using AA in CodeGen,
-        // we'll also need to update the TBAA nodes in MMOs with values
-        // derived from the merged allocas. When doing this, we'll need to use
-        // the same variant of GetUnderlyingObjects that is used by the
-        // instruction scheduler (that can look through ptrtoint/inttoptr
-        // pairs).
-
         // We've replaced IR-level uses of the remapped allocas, so we only
         // need to replace direct uses here.
         const AllocaInst *AI = dyn_cast_or_null<AllocaInst>(MMO->getValue());
         if (!AI)
           continue;
 
         if (!Allocas.count(AI))
           continue;
#ifndef NDEBUG
         }
 #endif
 
         // Fix the machine instructions.
         int ToSlot = SlotRemap[FromSlot];
         MO.setIndex(ToSlot);
         FixedInstr++;
       }
+
+      // We adjust AliasAnalysis information for merged stack slots.
+      MachineSDNode::mmo_iterator MemOps =
+          MF->allocateMemRefsArray(I.getNumMemOperands());
+      unsigned MemOpIdx = 0;
+      bool ReplaceMemOps = false;
+      for (MachineMemOperand *MMO : I.memoperands()) {
+        // If this memory location can be a slot remapped here,
+        // we remove AA information.
+        bool MayHaveConflictingAAMD = false;
+        if (MMO->getAAInfo()) {
+          if (const Value *MMOV = MMO->getValue()) {
+            SmallVector<Value *, 4> Objs;
+            getUnderlyingObjectsForCodeGen(MMOV, Objs, MF->getDataLayout());
+
+            if (Objs.empty())
+              MayHaveConflictingAAMD = true;
+            else
+              for (Value *V : Objs) {
+                // If this memory location comes from a known stack slot
+                // that is not remapped, we continue checking.
+                // Otherwise, we need to invalidate AA infomation.
+                if (!isIdentifiedObject(V))
+                  MayHaveConflictingAAMD = true;
+
+                const AllocaInst *AI = dyn_cast_or_null<AllocaInst>(V);
+                if (AI && MergedAllocas.count(AI))
+                  MayHaveConflictingAAMD = true;
+
+                if (MayHaveConflictingAAMD)
+                  break;
+              }
+          }
+        }
+        if (MayHaveConflictingAAMD) {
+          MemOps[MemOpIdx++] = MF->getMachineMemOperand(MMO, AAMDNodes());
+          ReplaceMemOps = true;
+        }
+        else
+          MemOps[MemOpIdx++] = MMO;
+      }
+
+      // If any memory operand is updated, set memory references of
+      // this instruction.
+      if (ReplaceMemOps)
+        I.setMemRefs(std::make_pair(MemOps, I.getNumMemOperands()));
     }
 
   // Update the location of C++ catch objects for the MSVC personality routine.
   if (WinEHFuncInfo *EHInfo = MF->getWinEHFuncInfo())
     for (WinEHTryBlockMapEntry &TBME : EHInfo->TryBlockMap)
       for (WinEHHandlerType &H : TBME.HandlerArray)
         if (H.CatchObj.FrameIndex != INT_MAX &&
             SlotRemap.count(H.CatchObj.FrameIndex))