Preserve more MD_mem_parallel_loop_access and MD_access_group in SROA

[mmendell]

Summary

SROA sometimes preserves MD_mem_parallel_loop_access and MD_access_group metadata on loads/stores, and sometimes fails to do so. This change adds copying of the MD after other CreateAlignedLoad/CreateAlignedStores. Also fix a case where the metadata was being copied from a load, rather than the store.
Added a LIT test to catch one case.

[mmendell]

re-uploading to see if the unrelated Windows build failures go away, as they aren't related at all

[mmendell]

Can someone please review? I haven't heard from anyone in a week

[wsmoses]

Not an expert on llvm.access.group, so someone please chime in with more thoughts. My one comment here is that since this impacts more than a load, I'd like to see test cases added for all the relevant behavior (including store).

[mmendell]

I really don't know enough to recreate the paths through all the places. The LIT test that I submitted is derived from an actual failure to propagate llvm.access.group in our compiler (and the change in visitIntrinsicInst(IntrinsicInst &II) from LI to SI). The others I just added to be consistent with existing calls to propagate the metadata that were already present.

llvm.access.group is used with llvm.loop.parallel_accesses to mark loads/stores that don't have loop carried memory dependencies

[Meinersbur]

llvm/lib/Transforms/Scalar/SROA.cpp

2478

NewAI is the alloca instruction; it does not have parallel loop access or access group metadata which only apply to memory accesses.

2538

Did you mean LI instead of NewAI?

llvm/test/Transforms/SROA/mem-par-metadata-sroa-cast.ll

8–9

This tests the presplitLoadsAndStores case. Shouldn't it also have split the store?

[mmendell]

Removed the copies of metadata from NewAI, as they can't contain any metadata. Only keep the ones from the load/stores

[mmendell]

llvm/test/Transforms/SROA/mem-par-metadata-sroa-cast.ll

8–9

The store is replaced with:

%0 = bitcast i32 %X371 to float
%1 = bitcast i32 %X373 to float

Should I add that to the LIT test? The important part was the loss of the llvm.access.group from the store

[Meinersbur]

llvm/lib/Transforms/Scalar/SROA.cpp

2478

This one might have made sense

llvm/test/Transforms/SROA/mem-par-metadata-sroa-cast.ll

8–9

The bitcasts belong to the load. I guess SROA determined that the store is dead (stack memory before returning) and removed it. Not an indication of a good test case.

I suggest to make %PREV a function pointer so this does not happen. Attach a separate !llvm.access.group !1 to the store and test it as well. Most other instructions in this function seems to be irrelevant.

[mmendell]

llvm/lib/Transforms/Scalar/SROA.cpp

2478

Yes, you are right. I will pass in LI

llvm/test/Transforms/SROA/mem-par-metadata-sroa-cast.ll

8–9

If I make %PREV a function pointer, then the cast isn't done: compiler explorer
I can remove 2 lines (defn and use of %r2) and still get it to fail. Anything else, and it doesn't generate the sroa_cast/sroa_cast2

[mmendell]

Restored the copy of metadata for vectorized loads, passing in LI.
Trimmed the LIT test slightly

[Meinersbur]

LGTM

llvm/test/Transforms/SROA/mem-par-metadata-sroa-cast.ll

8–9

OK, I also tried a bit to get SROA to re-emit a store (and therefore test that the store's metadata is preserved), but without success.

llvm/lib/Transforms/Scalar/SROA.cpp

#endif
 
   Value *rewriteVectorizedLoadInst() {
     unsigned BeginIndex = getIndex(NewBeginOffset);
     unsigned EndIndex = getIndex(NewEndOffset);
     assert(EndIndex > BeginIndex && "Empty vector!");
 
     Value *V = IRB.CreateAlignedLoad(NewAI.getAllocatedType(), &NewAI,
                                      NewAI.getAlign(), "load");
+    cast<Instruction>(V)->copyMetadata(
+        NewAI, {LLVMContext::MD_mem_parallel_loop_access,
+                LLVMContext::MD_access_group});
     return extractVector(IRB, V, BeginIndex, EndIndex, "vec");
   }
 
   Value *rewriteIntegerLoad(LoadInst &LI) {
     assert(IntTy && "We cannot insert an integer to the alloca");
     assert(!LI.isVolatile());
     Value *V = IRB.CreateAlignedLoad(NewAI.getAllocatedType(), &NewAI,
                                      NewAI.getAlign(), "load");
+    cast<Instruction>(V)->copyMetadata(
+        NewAI, {LLVMContext::MD_mem_parallel_loop_access,
+                LLVMContext::MD_access_group});
     V = convertValue(DL, IRB, V, IntTy);
     assert(NewBeginOffset >= NewAllocaBeginOffset && "Out of bounds offset");
     uint64_t Offset = NewBeginOffset - NewAllocaBeginOffset;
     if (Offset > 0 || NewEndOffset < NewAllocaEndOffset) {
       IntegerType *ExtractTy = Type::getIntNTy(LI.getContext(), SliceSize * 8);
       V = extractInteger(DL, IRB, V, ExtractTy, Offset, "extract");
     }
     // It is possible that the extracted type is not the load type. This
bool visitLoadInst(LoadInst &LI) {
     } else if (NewBeginOffset == NewAllocaBeginOffset &&
                NewEndOffset == NewAllocaEndOffset &&
                (canConvertValue(DL, NewAllocaTy, TargetTy) ||
                 (IsLoadPastEnd && NewAllocaTy->isIntegerTy() &&
                  TargetTy->isIntegerTy()))) {
       LoadInst *NewLI = IRB.CreateAlignedLoad(NewAI.getAllocatedType(), &NewAI,
                                               NewAI.getAlign(), LI.isVolatile(),
                                               LI.getName());
+      NewLI->copyMetadata(NewAI, {LLVMContext::MD_mem_parallel_loop_access,
+                                  LLVMContext::MD_access_group});
       if (AATags)
         NewLI->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));
       if (LI.isVolatile())
         NewLI->setAtomic(LI.getOrdering(), LI.getSyncScopeID());
       if (NewLI->isAtomic())
         NewLI->setAlignment(LI.getAlign());
 
       // Any !nonnull metadata or !range metadata on the old load is also valid
if (VecTy) {
       Type *LTy = TargetTy->getPointerTo(AS);
       LoadInst *NewLI =
           IRB.CreateAlignedLoad(TargetTy, getNewAllocaSlicePtr(IRB, LTy),
                                 getSliceAlign(), LI.isVolatile(), LI.getName());
       if (AATags)
         NewLI->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));
       if (LI.isVolatile())
         NewLI->setAtomic(LI.getOrdering(), LI.getSyncScopeID());
+      NewLI->copyMetadata(LI, {LLVMContext::MD_mem_parallel_loop_access,
+                               LLVMContext::MD_access_group});
 
       V = NewLI;
       IsPtrAdjusted = true;
     }
     V = convertValue(DL, IRB, V, TargetTy);
 
     if (IsSplit) {
       assert(!LI.isVolatile());
if (V->getType() != VecTy) {
         V = convertValue(DL, IRB, V, SliceTy);
 
       // Mix in the existing elements.
       Value *Old = IRB.CreateAlignedLoad(NewAI.getAllocatedType(), &NewAI,
                                          NewAI.getAlign(), "load");
       V = insertVector(IRB, Old, V, BeginIndex, "vec");
     }
     StoreInst *Store = IRB.CreateAlignedStore(V, &NewAI, NewAI.getAlign());
+    Store->copyMetadata(SI, {LLVMContext::MD_mem_parallel_loop_access,
+                             LLVMContext::MD_access_group});
     if (AATags)
       Store->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));
     Pass.DeadInsts.push_back(&SI);
 
     LLVM_DEBUG(dbgs() << "          to: " << *Store << "\n");
     return true;
   }
 
   bool rewriteIntegerStore(Value *V, StoreInst &SI, AAMDNodes AATags) {
     assert(IntTy && "We cannot extract an integer from the alloca");
     assert(!SI.isVolatile());
     if (DL.getTypeSizeInBits(V->getType()).getFixedSize() !=
         IntTy->getBitWidth()) {
       Value *Old = IRB.CreateAlignedLoad(NewAI.getAllocatedType(), &NewAI,
                                          NewAI.getAlign(), "oldload");
+      cast<Instruction>(Old)->copyMetadata(
+          NewAI, {LLVMContext::MD_mem_parallel_loop_access,
+                  LLVMContext::MD_access_group});
       Old = convertValue(DL, IRB, Old, IntTy);
       assert(BeginOffset >= NewAllocaBeginOffset && "Out of bounds offset");
       uint64_t Offset = BeginOffset - NewAllocaBeginOffset;
       V = insertInteger(DL, IRB, Old, SI.getValueOperand(), Offset, "insert");
     }
     V = convertValue(DL, IRB, V, NewAllocaTy);
     StoreInst *Store = IRB.CreateAlignedStore(V, &NewAI, NewAI.getAlign());
     Store->copyMetadata(SI, {LLVMContext::MD_mem_parallel_loop_access,
if (VecTy) {
         V = getVectorSplat(
             V, cast<FixedVectorType>(AllocaVecTy)->getNumElements());
 
       V = convertValue(DL, IRB, V, AllocaTy);
     }
 
     StoreInst *New =
         IRB.CreateAlignedStore(V, &NewAI, NewAI.getAlign(), II.isVolatile());
+    New->copyMetadata(II, {LLVMContext::MD_mem_parallel_loop_access,
+                           LLVMContext::MD_access_group});
     if (AATags)
       New->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));
     LLVM_DEBUG(dbgs() << "          to: " << *New << "\n");
     return !II.isVolatile();
   }
 
   bool visitMemTransferInst(MemTransferInst &II) {
     // Rewriting of memory transfer instructions can be a bit tricky. We break
if (VecTy && !IsWholeAlloca && !IsDest) {
       Src = IRB.CreateAlignedLoad(NewAI.getAllocatedType(), &NewAI,
                                   NewAI.getAlign(), "load");
       Src = convertValue(DL, IRB, Src, IntTy);
       uint64_t Offset = NewBeginOffset - NewAllocaBeginOffset;
       Src = extractInteger(DL, IRB, Src, SubIntTy, Offset, "extract");
     } else {
       LoadInst *Load = IRB.CreateAlignedLoad(OtherTy, SrcPtr, SrcAlign,
                                              II.isVolatile(), "copyload");
+      Load->copyMetadata(II, {LLVMContext::MD_mem_parallel_loop_access,
+                              LLVMContext::MD_access_group});
       if (AATags)
         Load->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));
       Src = Load;
     }
 
     if (VecTy && !IsWholeAlloca && IsDest) {
       Value *Old = IRB.CreateAlignedLoad(NewAI.getAllocatedType(), &NewAI,
                                          NewAI.getAlign(), "oldload");
       Src = insertVector(IRB, Old, Src, BeginIndex, "vec");
     } else if (IntTy && !IsWholeAlloca && IsDest) {
       Value *Old = IRB.CreateAlignedLoad(NewAI.getAllocatedType(), &NewAI,
                                          NewAI.getAlign(), "oldload");
       Old = convertValue(DL, IRB, Old, IntTy);
       uint64_t Offset = NewBeginOffset - NewAllocaBeginOffset;
       Src = insertInteger(DL, IRB, Old, Src, Offset, "insert");
       Src = convertValue(DL, IRB, Src, NewAllocaTy);
     }
 
     StoreInst *Store = cast<StoreInst>(
         IRB.CreateAlignedStore(Src, DstPtr, DstAlign, II.isVolatile()));
+    Store->copyMetadata(II, {LLVMContext::MD_mem_parallel_loop_access,
+                             LLVMContext::MD_access_group});
     if (AATags)
       Store->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));
     LLVM_DEBUG(dbgs() << "          to: " << *Store << "\n");
     return !II.isVolatile();
   }
 
   bool visitIntrinsicInst(IntrinsicInst &II) {
     assert((II.isLifetimeStartOrEnd() || II.isDroppable()) &&
for (User *LU : LI->users()) {
         auto AS = SI->getPointerAddressSpace();
         StoreInst *PStore = IRB.CreateAlignedStore(
             PLoad,
             getAdjustedPtr(IRB, DL, StoreBasePtr,
                            APInt(DL.getIndexSizeInBits(AS), PartOffset),
                            PartPtrTy, StoreBasePtr->getName() + "."),
             getAdjustedAlignment(SI, PartOffset),
             /*IsVolatile*/ false);
-        PStore->copyMetadata(*LI, {LLVMContext::MD_mem_parallel_loop_access,
+        PStore->copyMetadata(*SI, {LLVMContext::MD_mem_parallel_loop_access,
                                    LLVMContext::MD_access_group});
         LLVM_DEBUG(dbgs() << "      +" << PartOffset << ":" << *PStore << "\n");
       }
 
       // We want to immediately iterate on any allocas impacted by splitting
       // this store, and we have to track any promotable alloca (indicated by
       // a direct store) as needing to be resplit because it is no longer
       // promotable.
for (;;) {
         auto AS = LI->getPointerAddressSpace();
         PLoad = IRB.CreateAlignedLoad(
             PartTy,
             getAdjustedPtr(IRB, DL, LoadBasePtr,
                            APInt(DL.getIndexSizeInBits(AS), PartOffset),
                            LoadPartPtrTy, LoadBasePtr->getName() + "."),
             getAdjustedAlignment(LI, PartOffset),
             /*IsVolatile*/ false, LI->getName());
+        PLoad->copyMetadata(*LI, {LLVMContext::MD_mem_parallel_loop_access,
+                                  LLVMContext::MD_access_group});
       }
 
       // And store this partition.
       IRB.SetInsertPoint(SI);
       auto AS = SI->getPointerAddressSpace();
       StoreInst *PStore = IRB.CreateAlignedStore(
           PLoad,
           getAdjustedPtr(IRB, DL, StoreBasePtr,
                          APInt(DL.getIndexSizeInBits(AS), PartOffset),
                          StorePartPtrTy, StoreBasePtr->getName() + "."),
           getAdjustedAlignment(SI, PartOffset),
           /*IsVolatile*/ false);
+      PStore->copyMetadata(*SI, {LLVMContext::MD_mem_parallel_loop_access,
+                                 LLVMContext::MD_access_group});
 
       // Now build a new slice for the alloca.
       NewSlices.push_back(
           Slice(BaseOffset + PartOffset, BaseOffset + PartOffset + PartSize,
                 &PStore->getOperandUse(PStore->getPointerOperandIndex()),
                 /*IsSplittable*/ false));
       LLVM_DEBUG(dbgs() << "    new slice [" << NewSlices.back().beginOffset()
                         << ", " << NewSlices.back().endOffset()

llvm/test/Transforms/SROA/mem-par-metadata-sroa-cast.ll

+; RUN: opt < %s -sroa -S | FileCheck %s
+;
+; Make sure the llvm.access.group meta-data is preserved
+; when a load/store is replaced with another load/store by sroa
+; Ensure this is done for casting too.
+;
+; CHECK: entry:
+; CHECK: load i32, i32* {{.*}}, !llvm.access.group [[DISTINCT:![0-9]*]]
+; CHECK: load i32, i32* {{.*}}, !llvm.access.group [[DISTINCT]]
+; CHECK: ret void
+; CHECK: [[DISTINCT]] = distinct !{}
+
+%CMPLX = type { float, float }
+
+define dso_local void @test() {
+entry:
+  %PART = alloca %CMPLX, align 8
+  %PREV = alloca %CMPLX, align 8
+  %r2 = getelementptr %CMPLX, %CMPLX* %PREV, i32 0, i32 0
+  store float 0.000000e+00, float* %r2, align 4
+  %i2 = getelementptr %CMPLX, %CMPLX* %PREV, i32 0, i32 1
+  store float 0.000000e+00, float* %i2, align 4
+  %dummy = sext i16 0 to i64
+  %T = getelementptr %CMPLX, %CMPLX* %PART, i64 %dummy
+  %X35 = bitcast %CMPLX* %T to i64*
+  %X36 = bitcast %CMPLX* %PREV to i64*
+  %X37 = load i64, i64* %X35, align 8, !llvm.access.group !0
+  store i64 %X37, i64* %X36, align 8
+  ret void
+}
+
+!0 = distinct !{}