[InstCombine] Retain TBAA when narrowing loads

[loladiro]

Summary

As discussed on the mailing list it is legal to propagate TBAA to loads from
smaller regions of a larger load tagged with TBAA. Do so for
(load->extractvalue)=>(gep->load) folding.

[loladiro]

There were a couple more of instances of the same pattern in this code.
I'll write some tests for all these cases next week (famous last words).

[loladiro]

Now with testcases for the extra cases I discovered.

[loladiro]

Run clang-format over the diff

[loladiro]

Bump.

[sanjoy]

lgtm

lib/Transforms/InstCombine/InstructionCombining.cpp

2433 ↗

(On Diff #95497)

Call this NL like you've used NS before?

llvm/trunk/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

static Instruction *unpackLoadToAggregate(InstCombiner &IC, LoadInst &LI) {
   assert(LI.getAlignment() && "Alignment must be set at this point");
 
   if (auto *ST = dyn_cast<StructType>(T)) {
     // If the struct only have one element, we unpack.
     auto NumElements = ST->getNumElements();
     if (NumElements == 1) {
       LoadInst *NewLoad = combineLoadToNewType(IC, LI, ST->getTypeAtIndex(0U),
                                                ".unpack");
+      AAMDNodes AAMD;
+      LI.getAAMetadata(AAMD);
+      NewLoad->setAAMetadata(AAMD);
       return IC.replaceInstUsesWith(LI, IC.Builder->CreateInsertValue(
         UndefValue::get(T), NewLoad, 0, Name));
     }
 
     // We don't want to break loads with padding here as we'd loose
     // the knowledge that padding exists for the rest of the pipeline.
     const DataLayout &DL = IC.getDataLayout();
     auto *SL = DL.getStructLayout(ST);
for (unsigned i = 0; i < NumElements; i++) {
       Value *Indices[2] = {
         Zero,
         ConstantInt::get(IdxType, i),
       };
       auto *Ptr = IC.Builder->CreateInBoundsGEP(ST, Addr, makeArrayRef(Indices),
                                                 Name + ".elt");
       auto EltAlign = MinAlign(Align, SL->getElementOffset(i));
       auto *L = IC.Builder->CreateAlignedLoad(Ptr, EltAlign, Name + ".unpack");
+      // Propagate AA metadata. It'll still be valid on the narrowed load.
+      AAMDNodes AAMD;
+      LI.getAAMetadata(AAMD);
+      L->setAAMetadata(AAMD);
       V = IC.Builder->CreateInsertValue(V, L, i);
     }
 
     V->setName(Name);
     return IC.replaceInstUsesWith(LI, V);
   }
 
   if (auto *AT = dyn_cast<ArrayType>(T)) {
     auto *ET = AT->getElementType();
     auto NumElements = AT->getNumElements();
     if (NumElements == 1) {
       LoadInst *NewLoad = combineLoadToNewType(IC, LI, ET, ".unpack");
+      AAMDNodes AAMD;
+      LI.getAAMetadata(AAMD);
+      NewLoad->setAAMetadata(AAMD);
       return IC.replaceInstUsesWith(LI, IC.Builder->CreateInsertValue(
         UndefValue::get(T), NewLoad, 0, Name));
     }
 
     // Bail out if the array is too large. Ideally we would like to optimize
     // arrays of arbitrary size but this has a terrible impact on compile time.
     // The threshold here is chosen arbitrarily, maybe needs a little bit of
     // tuning.
for (uint64_t i = 0; i < NumElements; i++) {
       Value *Indices[2] = {
         Zero,
         ConstantInt::get(IdxType, i),
       };
       auto *Ptr = IC.Builder->CreateInBoundsGEP(AT, Addr, makeArrayRef(Indices),
                                                 Name + ".elt");
       auto *L = IC.Builder->CreateAlignedLoad(Ptr, MinAlign(Align, Offset),
                                               Name + ".unpack");
+      AAMDNodes AAMD;
+      LI.getAAMetadata(AAMD);
+      L->setAAMetadata(AAMD);
       V = IC.Builder->CreateInsertValue(V, L, i);
       Offset += EltSize;
     }
 
     V->setName(Name);
     return IC.replaceInstUsesWith(LI, V);
   }
 
for (unsigned i = 0; i < Count; i++) {
       Value *Indices[2] = {
         Zero,
         ConstantInt::get(IdxType, i),
       };
       auto *Ptr = IC.Builder->CreateInBoundsGEP(ST, Addr, makeArrayRef(Indices),
                                                 AddrName);
       auto *Val = IC.Builder->CreateExtractValue(V, i, EltName);
       auto EltAlign = MinAlign(Align, SL->getElementOffset(i));
+      llvm::Instruction *NS =
           IC.Builder->CreateAlignedStore(Val, Ptr, EltAlign);
+      AAMDNodes AAMD;
+      SI.getAAMetadata(AAMD);
+      NS->setAAMetadata(AAMD);
     }
 
     return true;
   }
 
   if (auto *AT = dyn_cast<ArrayType>(T)) {
     // If the array only have one element, we unpack.
     auto NumElements = AT->getNumElements();
for (uint64_t i = 0; i < NumElements; i++) {
       Value *Indices[2] = {
         Zero,
         ConstantInt::get(IdxType, i),
       };
       auto *Ptr = IC.Builder->CreateInBoundsGEP(AT, Addr, makeArrayRef(Indices),
                                                 AddrName);
       auto *Val = IC.Builder->CreateExtractValue(V, i, EltName);
       auto EltAlign = MinAlign(Align, Offset);
-      IC.Builder->CreateAlignedStore(Val, Ptr, EltAlign);
+      Instruction *NS = IC.Builder->CreateAlignedStore(Val, Ptr, EltAlign);
+      AAMDNodes AAMD;
+      SI.getAAMetadata(AAMD);
+      NS->setAAMetadata(AAMD);
       Offset += EltSize;
     }
 
     return true;
   }
 
   return false;
 }

llvm/trunk/lib/Transforms/InstCombine/InstructionCombining.cpp

if (L->isSimple() && L->hasOneUse()) {
             I != E; ++I)
         Indices.push_back(Builder->getInt32(*I));
 
       // We need to insert these at the location of the old load, not at that of
       // the extractvalue.
       Builder->SetInsertPoint(L);
       Value *GEP = Builder->CreateInBoundsGEP(L->getType(),
                                               L->getPointerOperand(), Indices);
+      Instruction *NL = Builder->CreateLoad(GEP);
+      // Whatever aliasing information we had for the orignal load must also
+      // hold for the smaller load, so propagate the annotations.
+      AAMDNodes Nodes;
+      L->getAAMetadata(Nodes);
+      NL->setAAMetadata(Nodes);
       // Returning the load directly will cause the main loop to insert it in
       // the wrong spot, so use replaceInstUsesWith().
-      return replaceInstUsesWith(EV, Builder->CreateLoad(GEP));
+      return replaceInstUsesWith(EV, NL);
     }
   // We could simplify extracts from other values. Note that nested extracts may
   // already be simplified implicitly by the above: extract (extract (insert) )
   // will be translated into extract ( insert ( extract ) ) first and then just
   // the value inserted, if appropriate. Similarly for extracts from single-use
   // loads: extract (extract (load)) will be translated to extract (load (gep))
   // and if again single-use then via load (gep (gep)) to load (gep).
   // However, double extracts from e.g. function arguments or return values

llvm/trunk/test/Transforms/InstCombine/extractinsert-tbaa.ll

+; RUN: opt -S -instcombine %s -o - | FileCheck %s
+
+%Complex = type { double, double }
+
+; Check that instcombine preserves TBAA when narrowing loads
+define double @teststructextract(%Complex *%val) {
+; CHECK: load double, {{.*}}, !tbaa
+; CHECK-NOT: load %Complex
+    %loaded = load %Complex, %Complex *%val, !tbaa !1
+    %real = extractvalue %Complex %loaded, 0
+    ret double %real
+}
+
+define double @testarrayextract([2 x double] *%val) {
+; CHECK: load double, {{.*}}, !tbaa
+; CHECK-NOT: load [2 x double]
+    %loaded = load [2 x double], [2 x double] *%val, !tbaa !1
+    %real = extractvalue [2 x double] %loaded, 0
+    ret double %real
+}
+
+; Check that inscombine preserves TBAA when breaking up stores
+define void @teststructinsert(%Complex *%loc, double %a, double %b) {
+; CHECK: store double %a, {{.*}}, !tbaa
+; CHECK: store double %b, {{.*}}, !tbaa
+; CHECK-NOT: store %Complex
+    %inserted  = insertvalue %Complex undef,      double %a, 0
+    %inserted2 = insertvalue %Complex %inserted,  double %b, 1
+    store %Complex %inserted2, %Complex *%loc, !tbaa !1
+    ret void
+}
+
+define void @testarrayinsert([2 x double] *%loc, double %a, double %b) {
+; CHECK: store double %a, {{.*}}, !tbaa
+; CHECK: store double %b, {{.*}}, !tbaa
+; CHECK-NOT: store [2 x double]
+    %inserted  = insertvalue [2 x double] undef,      double %a, 0
+    %inserted2 = insertvalue [2 x double] %inserted,  double %b, 1
+    store [2 x double] %inserted2, [2 x double] *%loc, !tbaa !1
+    ret void
+}
+
+!0 = !{!"tbaa_root"}
+!1 = !{!2, !2, i64 0}
+!2 = !{!"Complex", !0, i64 0}