Fix a use-after-RAUW bug in large GEP splitting

[amharc]

Summary

Large GEP splitting, introduced in rL332015, uses a DenseMap<AssertingVH<Value>, ...>. This causes an assertion to fail (in debug builds) or undefined behaviour to occur (in release builds) when a value is RAUWed.

This manifested itself in the 7zip benchmark from the llvm test suite built on ARM with -fstrict-vtable-pointers enabled while RAUWing invariant group launders and splits in CodeGenPrepare.

This patch merges the large offsets of the argument and the result of an invariant.group strip/launder intrinsic before RAUWing.

[amharc]

The other AssertingVHs held by large GEP splitting code are either AssertingVH<GetElementPtrInst> or refer to possibly bitcasted GetElementPtrInsts. As far as I can see, there is no RAUWing of those being done in CGP at the moment.

This patch should probably include a custom ValueMapConfig which merges the vectors of GEPs while RAUWing.

[Prazek]

LGTM, but please wait for Javed to review

[efriedma]

llvm/lib/CodeGen/CodeGenPrepare.cpp

281

Please add a comment explaining why the keys can be RAUW'ed, but the values can't.

Please add a comment explaining what happens if the result of the RAUW already exists in the map... or if that isn't possible, why it isn't possible.

[efriedma]

What happened to this patch?

[amharc]

Updated the patch to manually merge vectors of large GEP offsets when doing a RAUW of an invariant.group strip/launder intrinsic. This requires only handling keys and not values of the map, as values are always GetElementPtr instructions and not intrinsic calls.

I do not have the necessary level of understanding of CGP to determine if there are any other RAUWs there that need attention too, but so far I wasn't able to reproduce in cases that do not relate to invariant.group strips/launders.

[efriedma]

LGTM

llvm/lib/CodeGen/CodeGenPrepare.cpp

class CodeGenPrepare : public FunctionPass {
     SetOfInstrs RemovedInsts;
 
     /// Keep track of sext chains based on their initial value.
     DenseMap<Value *, Instruction *> SeenChainsForSExt;
 
     /// Keep track of GEPs accessing the same data structures such as structs or
     /// arrays that are candidates to be split later because of their large
     /// size.
-    DenseMap<
-        AssertingVH<Value>,
+    ValueMap<
+        Value*,
         SmallVector<std::pair<AssertingVH<GetElementPtrInst>, int64_t>, 32>>
         LargeOffsetGEPMap;
 
     /// Keep track of new GEP base after splitting the GEPs having large offset.
     SmallSet<AssertingVH<Value>, 2> NewGEPBases;
 
     /// Map serial numbers to Large offset GEPs.
     DenseMap<AssertingVH<GetElementPtrInst>, int> LargeOffsetGEPID;

 //   %load1    = load i32, i32* %new_gep0
 //   %load2    = load i32, i32* %new_gep1
 //   %load3    = load i32, i32* %new_gep2
 //
 // %new_gep1 and %new_gep2 can be sunk to BB2 now after the splitting because
 // their offsets are smaller enough to fit into the addressing mode.
 bool CodeGenPrepare::splitLargeGEPOffsets() {
   bool Changed = false;
-  for (auto &Entry : LargeOffsetGEPMap) {
+  for (const auto &Entry : LargeOffsetGEPMap) {
     Value *OldBase = Entry.first;
     SmallVectorImpl<std::pair<AssertingVH<GetElementPtrInst>, int64_t>>
         &LargeOffsetGEPs = Entry.second;
     auto compareGEPOffset =
         [&](const std::pair<GetElementPtrInst *, int64_t> &LHS,
             const std::pair<GetElementPtrInst *, int64_t> &RHS) {
           if (LHS.first == RHS.first)
             return false;

llvm/test/Transforms/CodeGenPrepare/AArch64/large-offset-gep.ll

; CHECK:     str      w{{[0-9]+}}, [x{{[0-9]+}}, #4]
   %i = add i32 %phi, 1
   store i32 %i, i32* %gep0
   store i32 %phi, i32* %gep1
   br label %while_cond
 
 while_end:
   ret void
 }
+
+declare i8* @llvm.strip.invariant.group.p0i8(i8*)
+
+define void @test_invariant_group(i32) {
+; CHECK-LABEL: test_invariant_group
+  br i1 undef, label %8, label %7
+
+; <label>:2:                                      ; preds = %8, %2
+  br i1 undef, label %2, label %7
+
+; <label>:3:                                      ; preds = %8
+  %4 = getelementptr inbounds i8, i8* %9, i32 40000
+  %5 = bitcast i8* %4 to i64*
+  br i1 undef, label %7, label %6
+
+; <label>:6:                                      ; preds = %3
+  store i64 1, i64* %5, align 8
+  br label %7
+
+; <label>:7:                                      ; preds = %6, %3, %2, %1
+  ret void
+
+; <label>:8:                                      ; preds = %1
+  %9 = call i8* @llvm.strip.invariant.group.p0i8(i8* nonnull undef)
+  %10 = icmp eq i32 %0, 0
+  br i1 %10, label %3, label %2
+}
+