[RewriteStatepointsForGC] Reduce the number of new instructions for base pointers

[reames]

Summary

When computing base pointers, we introduce new instructions to propagate the base of existing instructions which might not be bases. However, the algorithm doesn't make any effort to recognize when the new instruction to be inserted is the same as an existing one already in the IR. Since this is happening immediately before rewriting, we don't really have a chance to fix it after the pass runs without teaching loop passes about statepoints.

I'm really not thrilled with this patch. I've rewritten it 4 different ways now, but this is the best I've come up with. The case where the new instruction is just the original base defining value could be merged into the existing algorithm with some complexity. The problem is that we might have something like an extractelement from a phi of two vectors. It may be trivially obvious that the base of the 0th element is an existing instruction, but I can't see how to make the algorithm itself figure that out. Thus, I resort to the call to SimplifyInstruction instead.

Note that we can only adjust the instructions we've inserted ourselves. The live sets are still being tracked in side structures at this point in the code. We can't easily muck with instructions which might be in them. Long term, I'm really thinking we need to materialize the live pointer sets explicitly in the IR somehow rather than using side structures to track them.

[sanjoy]

I don't see any particular problem with the code, but I have a few questions (inline) to ensure that I've understood the code correctly.

lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

1018 ↗

(On Diff #32014)

There are some whitespace issues here -- do you mind running this patch through clang-format before checkin?

1020 ↗

(On Diff #32014)

Should be Item.

1032 ↗

(On Diff #32014)

Can't this set be populated as we create the instructions above?

1048 ↗

(On Diff #32014)

Can PushNewUsers(BaseI) be hoisted out of this if block with a if (Visited.insert(BaseI).second) check? Or does this need to remain conditional for correctness?

1055 ↗

(On Diff #32014)

I'd hoist this bit outside the loop.

1057 ↗

(On Diff #32014)

Will re-pushing V onto Worklist help? I'm thinking of cases where SimplifyInstruction(BaseI)->isIdenticalTo(BdvI).

[reames]

Anything I didn't explicitly comment on, I accept and will address in a revised patch once the comments are settled.

lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

1018 ↗

(On Diff #32014)

Sure.

1032 ↗

(On Diff #32014)

It absolutely could. I was trying to keep the code separate to make it standalone, but I can merge if you'd prefer.

1048 ↗

(On Diff #32014)

It's required for correctness. Consider a self referential phi node. You only want to add it to the worklist if you modify it, otherwise the loop would run forever.

Oh, wait, you were proposing a separate Visited structure. Yes, that would work. Do you have a preference?

1057 ↗

(On Diff #32014)

I don't think this would ever happen in practice, but it's theoretically sound to do.

[sanjoy]

lgtm with comments addressed.

lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

1032 ↗

(On Diff #32014)

I'd prefer that -- it makes it more obvious what happened.

1048 ↗

(On Diff #32014)

Do you have a preference?

No, I was just trying to make sure I understood what was going on. :)

[reames]

lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

1032 ↗

(On Diff #32014)

I spoke with Sanjoy offline. We convinced ourselves that tracking all the new instructions was actually more complex than the current patch, so I'm going ahead with the current scheme.

Two options which might be worth exploring in the future:

• Using a custom IR builder with an insertion callback as InstCombine does. We'd have to be a bit careful here since we purposely have partially constructed phis/selects.
• Switch the states map to being entirely bi-directional and maintain that all the way through.

llvm/trunk/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

 //===- RewriteStatepointsForGC.cpp - Make GC relocations explicit ---------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // Rewrite an existing set of gc.statepoints such that they make potential
 // relocations performed by the garbage collector explicit in the IR.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Pass.h"
 #include "llvm/Analysis/CFG.h"
+#include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CallSite.h"
#endif
         Base = states[Base].getBase();
         assert(Base != nullptr && "unknown BDVState!");
       }
       assert(Base && "can't be null");
       BaseEE->setOperand(0, Base);
     }
   }
 
+  // Now that we're done with the algorithm, see if we can optimize the 
+  // results slightly by reducing the number of new instructions needed. 
+  // Arguably, this should be integrated into the algorithm above, but 
+  // doing as a post process step is easier to reason about for the moment.
+  DenseMap<Value *, Value *> ReverseMap;
+  SmallPtrSet<Instruction *, 16> NewInsts;
+  SmallSetVector<Instruction *, 16> Worklist;
+  for (auto Item : states) {
+    Value *V = Item.first;
+    Value *Base = Item.second.getBase();
+    assert(V && Base);
+    assert(!isKnownBaseResult(V) && "why did it get added?");
+    assert(isKnownBaseResult(Base) &&
+           "must be something we 'know' is a base pointer");
+    if (!Item.second.isConflict())
+      continue;
+
+    ReverseMap[Base] = V;
+    if (auto *BaseI = dyn_cast<Instruction>(Base)) {
+      NewInsts.insert(BaseI);
+      Worklist.insert(BaseI);
+    }
+  }
+  auto PushNewUsers = [&](Instruction *I) {
+    for (User *U : I->users())
+      if (auto *UI = dyn_cast<Instruction>(U))
+        if (NewInsts.count(UI))
+          Worklist.insert(UI);
+  };
+  const DataLayout &DL = cast<Instruction>(def)->getModule()->getDataLayout();
+  while (!Worklist.empty()) {
+    Instruction *BaseI = Worklist.pop_back_val();
+    Value *Bdv = ReverseMap[BaseI];
+    if (auto *BdvI = dyn_cast<Instruction>(Bdv))
+      if (BaseI->isIdenticalTo(BdvI)) {
+        DEBUG(dbgs() << "Identical Base: " << *BaseI << "\n");
+        PushNewUsers(BaseI);
+        BaseI->replaceAllUsesWith(Bdv);
+        BaseI->eraseFromParent();
+        states[Bdv] = BDVState(BDVState::Conflict, Bdv);
+        NewInsts.erase(BaseI);
+        ReverseMap.erase(BaseI);
+        continue;
+      }
+    if (Value *V = SimplifyInstruction(BaseI, DL)) {
+      DEBUG(dbgs() << "Base " << *BaseI << " simplified to " << *V << "\n");
+      PushNewUsers(BaseI);
+      BaseI->replaceAllUsesWith(V);
+      BaseI->eraseFromParent();
+      states[Bdv] = BDVState(BDVState::Conflict, V);
+      NewInsts.erase(BaseI);
+      ReverseMap.erase(BaseI);
+      continue;
+    }
+  }
+
   // Cache all of our results so we can cheaply reuse them
   // NOTE: This is actually two caches: one of the base defining value
   // relation and one of the base pointer relation!  FIXME
   for (auto item : states) {
     Value *v = item.first;
     Value *base = item.second.getBase();
     assert(v && base);
-    assert(!isKnownBaseResult(v) && "why did it get added?");
 
     if (TraceLSP) {
       std::string fromstr =
           cache.count(v) ? (cache[v]->hasName() ? cache[v]->getName() : "")
                          : "none";
       errs() << "Updating base value cache"
              << " for: " << (v->hasName() ? v->getName() : "")
              << " from: " << fromstr
              << " to: " << (base->hasName() ? base->getName() : "") << "\n";
     }
 
-    assert(isKnownBaseResult(base) &&
-           "must be something we 'know' is a base pointer");
     if (cache.count(v)) {
       // Once we transition from the BDV relation being store in the cache to
       // the base relation being stored, it must be stable
       assert((!isKnownBaseResult(cache[v]) || cache[v] == base) &&
              "base relation should be stable");
     }
     cache[v] = base;
   }

llvm/trunk/test/Transforms/RewriteStatepointsForGC/base-pointers-4.ll

 ; RUN: opt %s -rewrite-statepoints-for-gc -spp-print-base-pointers -S 2>&1 | FileCheck %s
 
 
-; CHECK: derived %obj_to_consume base %obj_to_consume.base
+; CHECK: derived %obj_to_consume base %obj_to_consume
 
 declare void @foo()
 declare i64 addrspace(1)* @generate_obj()
 declare void @consume_obj(i64 addrspace(1)*)
 
 define void @test(i32 %condition) gc "statepoint-example" {
 entry:
   br label %loop

 dest_b:
   br label %merge
 
 dest_c:
   br label %merge
 
 merge:
 ; CHECK: merge:
-; CHECK:  %obj_to_consume.base = phi i64 addrspace(1)* [ %obj2, %dest_a ], [ null, %dest_b ], [ null, %dest_c ]
 ; CHECK:  %obj_to_consume = phi i64 addrspace(1)* [ %obj2, %dest_a ], [ null, %dest_b ], [ null, %dest_c ]
 
   %obj_to_consume = phi i64 addrspace(1)* [ %obj2, %dest_a ], [ null, %dest_b ], [ null, %dest_c ]
   %safepoint_token3 = call i32 (i64, i32, void (i64 addrspace(1)*)*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidp1i64f(i64 0, i32 0, void (i64 addrspace(1)*)* @consume_obj, i32 1, i32 0, i64 addrspace(1)* %obj_to_consume, i32 0, i32 5, i32 0, i32 -1, i32 0, i32 0, i32 0)
   br label %merge.split
 
 merge.split:                                      ; preds = %merge
   %safepoint_token = call i32 (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 0, i32 0, void ()* @foo, i32 0, i32 0, i32 0, i32 5, i32 0, i32 -1, i32 0, i32 0, i32 0)

llvm/trunk/test/Transforms/RewriteStatepointsForGC/base-pointers.ll

 loop:                                             ; preds = %loop, %entry
 ; CHECK-LABEL: loop
 ; CHECK:   %current.base = phi i64 addrspace(1)*
 ; CHECK-DAG: [ %base_obj, %entry ]
 ; Given the two selects are equivelent, so are their base phis - ideally,
 ; we'd have commoned these, but that's a missed optimization, not correctness.
 ; CHECK-DAG: [ [[DISCARD:%.*.base.relocated.casted]], %loop ]
 ; CHECK-NOT: extra.base
-; CHECK: next.base = select
 ; CHECK: next = select
 ; CHECK: extra2.base = select
 ; CHECK: extra2 = select
 ; CHECK: statepoint
 ;; Both 'next' and 'extra2' are live across the backedge safepoint...
   %current = phi i64 addrspace(1)* [ %obj, %entry ], [ %next, %loop ]
   %extra = phi i64 addrspace(1)* [ %obj, %entry ], [ %extra2, %loop ]
   %nexta = getelementptr i64, i64 addrspace(1)* %current, i32 1
   %next = select i1 %cnd, i64 addrspace(1)* %nexta, i64 addrspace(1)* %base_arg2
   %extra2 = select i1 %cnd, i64 addrspace(1)* %nexta, i64 addrspace(1)* %base_arg2
   %safepoint_token = call i32 (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 0, i32 0, void ()* @foo, i32 0, i32 0, i32 0, i32 5, i32 0, i32 -1, i32 0, i32 0, i32 0)
   br label %loop
 }
 
+define i64 addrspace(1)* @test3(i1 %cnd, i64 addrspace(1)* %obj, 
+                                i64 addrspace(1)* %obj2)
+    gc "statepoint-example" {
+; CHECK-LABEL: @test3
+entry:
+  br i1 %cnd, label %merge, label %taken
+taken:
+  br label %merge
+merge:
+; CHECK-LABEL: merge:
+; CHECK-NEXT: %bdv = phi
+; CHECK-NEXT: gc.statepoint
+  %bdv = phi i64 addrspace(1)* [ %obj, %entry ], [ %obj2, %taken ]
+  %safepoint_token = call i32 (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 0, i32 0, void ()* @foo, i32 0, i32 0, i32 0, i32 5, i32 0, i32 -1, i32 0, i32 0, i32 0)
+  ret i64 addrspace(1)* %bdv
+}
+
+
 declare void @foo()
 declare i32 @llvm.experimental.gc.statepoint.p0f_isVoidf(i64, i32, void ()*, i32, i32, ...)
 declare i32 @llvm.experimental.gc.statepoint.p0f_isVoidp1i64f(i64, i32, void (i64 addrspace(1)*)*, i32, i32, ...)

llvm/trunk/test/Transforms/RewriteStatepointsForGC/base-vector.ll

 taken2:
   %obj0 = extractelement <2 x i64 addrspace(1)*> %vec, i32 %idx1
   br label %merge2
 untaken2:
   %obj1 = extractelement <2 x i64 addrspace(1)*> %vec, i32 %idx2
   br label %merge2
 merge2:
 ; CHECK-LABEL: merge2:
-; CHECK: %obj.base = phi i64 addrspace(1)*
-; CHECK: %obj = phi i64 addrspace(1)*
-; CHECK: statepoint
-; CHECK: gc.relocate
-; CHECK-DAG: ; (%obj.base, %obj)
+; CHECK-NEXT: %obj = phi i64 addrspace(1)*
+; CHECK-NEXT: statepoint
 ; CHECK: gc.relocate
-; CHECK-DAG: ; (%obj.base, %obj.base)
+; CHECK-DAG: ; (%obj, %obj)
   %obj = phi i64 addrspace(1)* [%obj0, %taken2], [%obj1, %untaken2]
   %safepoint_token = call i32 (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 0, i32 0, void ()* @do_safepoint, i32 0, i32 0, i32 0, i32 0)
   ret i64 addrspace(1)* %obj
 }
 
 define i64 addrspace(1)* @test3(i64 addrspace(1)* %ptr) 
     gc "statepoint-example" {
 ; CHECK-LABEL: test3

llvm/trunk/test/Transforms/RewriteStatepointsForGC/live-vector.ll

 
 
 ; A base vector from a load
 define <2 x i64 addrspace(1)*> @test6(i1 %cnd, <2 x i64 addrspace(1)*>* %ptr) 
     gc "statepoint-example" {
 ; CHECK-LABEL: test6
 ; CHECK-LABEL: merge:
 ; CHECK-NEXT: = phi
-; CHECK-NEXT: = phi
-; CHECK-NEXT: extractelement
-; CHECK-NEXT: extractelement
 ; CHECK-NEXT: extractelement
 ; CHECK-NEXT: extractelement
 ; CHECK-NEXT: gc.statepoint
 ; CHECK-NEXT: gc.relocate
 ; CHECK-NEXT: bitcast
 ; CHECK-NEXT: gc.relocate
 ; CHECK-NEXT: bitcast
-; CHECK-NEXT: gc.relocate
-; CHECK-NEXT: bitcast
-; CHECK-NEXT: gc.relocate
-; CHECK-NEXT: bitcast
-; CHECK-NEXT: insertelement
-; CHECK-NEXT: insertelement
 ; CHECK-NEXT: insertelement
 ; CHECK-NEXT: insertelement
 ; CHECK-NEXT: ret <2 x i64 addrspace(1)*>
 entry:
   br i1 %cnd, label %taken, label %untaken
 taken:
   %obja = load <2 x i64 addrspace(1)*>, <2 x i64 addrspace(1)*>* %ptr
   br label %merge