[InstCombine] Remove unused entries in gc-liv bundle of statepoint

[skatkov]

Summary

If some of gc live value are not used in gc.relocate we can remove them
from gc-live bundle of statepoint instruction.

Also the CL removes duplicated Values in gc-live bundle.

[dantrushin]

LGTM with couple nits inline

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1549

Change to early exit (break)

1579

Can we avoid creating empty bundles here?

llvm/include/llvm/IR/Instructions.h

public:
   /// insert it before \p InsertPt.
   ///
   /// The returned call instruction is identical \p CI in every way except that
   /// the operand bundles for the new instruction are set to the operand bundles
   /// in \p Bundles.
   static CallInst *Create(CallInst *CI, ArrayRef<OperandBundleDef> Bundles,
                           Instruction *InsertPt = nullptr);
 
+  /// Create a clone of \p CI with a different set of operand bundles and
+  /// insert it before \p InsertPt.
+  ///
+  /// The returned call instruction is identical \p CI in every way except that
+  /// the operand bundle for the new instruction is set to the operand bundle
+  /// in \p Bundle.
+  static CallInst *CreateWithReplacedBundle(CallInst *CI,
+                                            OperandBundleDef Bundle,
+                                            Instruction *InsertPt = nullptr);
+
   /// Generate the IR for a call to malloc:
   /// 1. Compute the malloc call's argument as the specified type's size,
   ///    possibly multiplied by the array size if the array size is not
   ///    constant 1.
   /// 2. Call malloc with that argument.
   /// 3. Bitcast the result of the malloc call to the specified type.
   static Instruction *CreateMalloc(Instruction *InsertBefore, Type *IntPtrTy,
                                    Type *AllocTy, Value *AllocSize,
public:
   /// insert it before \p InsertPt.
   ///
   /// The returned invoke instruction is identical to \p II in every way except
   /// that the operand bundles for the new instruction are set to the operand
   /// bundles in \p Bundles.
   static InvokeInst *Create(InvokeInst *II, ArrayRef<OperandBundleDef> Bundles,
                             Instruction *InsertPt = nullptr);
 
+  /// Create a clone of \p II with a different set of operand bundles and
+  /// insert it before \p InsertPt.
+  ///
+  /// The returned invoke instruction is identical to \p II in every way except
+  /// that the operand bundle for the new instruction is set to the operand
+  /// bundle in \p Bundle.
+  static InvokeInst *CreateWithReplacedBundle(InvokeInst *II,
+                                              OperandBundleDef Bundles,
+                                              Instruction *InsertPt = nullptr);
+
   // get*Dest - Return the destination basic blocks...
   BasicBlock *getNormalDest() const {
     return cast<BasicBlock>(Op<NormalDestOpEndIdx>());
   }
   BasicBlock *getUnwindDest() const {
     return cast<BasicBlock>(Op<UnwindDestOpEndIdx>());
   }
   void setNormalDest(BasicBlock *B) {

llvm/lib/IR/Instructions.cpp

CallInst *CallInst::Create(CallInst *CI, ArrayRef<OperandBundleDef> OpB,
   NewCI->setTailCallKind(CI->getTailCallKind());
   NewCI->setCallingConv(CI->getCallingConv());
   NewCI->SubclassOptionalData = CI->SubclassOptionalData;
   NewCI->setAttributes(CI->getAttributes());
   NewCI->setDebugLoc(CI->getDebugLoc());
   return NewCI;
 }
 
+CallInst *CallInst::CreateWithReplacedBundle(CallInst *CI, OperandBundleDef OpB,
+                                             Instruction *InsertPt) {
+  SmallVector<OperandBundleDef, 2> OpDefs;
+  for (unsigned i = 0, e = CI->getNumOperandBundles(); i < e; ++i) {
+    auto ChildOB = CI->getOperandBundleAt(i);
+    if (ChildOB.getTagName() != OpB.getTag())
+      OpDefs.emplace_back(ChildOB);
+  }
+  OpDefs.emplace_back(OpB);
+  return CallInst::Create(CI, OpDefs, InsertPt);
+}
+
 // Update profile weight for call instruction by scaling it using the ratio
 // of S/T. The meaning of "branch_weights" meta data for call instruction is
 // transfered to represent call count.
 void CallInst::updateProfWeight(uint64_t S, uint64_t T) {
   auto *ProfileData = getMetadata(LLVMContext::MD_prof);
   if (ProfileData == nullptr)
     return;
 
auto *NewII = InvokeInst::Create(
       II->getUnwindDest(), Args, OpB, II->getName(), InsertPt);
   NewII->setCallingConv(II->getCallingConv());
   NewII->SubclassOptionalData = II->SubclassOptionalData;
   NewII->setAttributes(II->getAttributes());
   NewII->setDebugLoc(II->getDebugLoc());
   return NewII;
 }
 
+InvokeInst *InvokeInst::CreateWithReplacedBundle(InvokeInst *II,
+                                                 OperandBundleDef OpB,
+                                                 Instruction *InsertPt) {
+  SmallVector<OperandBundleDef, 2> OpDefs;
+  for (unsigned i = 0, e = II->getNumOperandBundles(); i < e; ++i) {
+    auto ChildOB = II->getOperandBundleAt(i);
+    if (ChildOB.getTagName() != OpB.getTag())
+      OpDefs.emplace_back(ChildOB);
+  }
+  OpDefs.emplace_back(OpB);
+  return InvokeInst::Create(II, OpDefs, InsertPt);
+}
 
 LandingPadInst *InvokeInst::getLandingPadInst() const {
   return cast<LandingPadInst>(getUnwindDest()->getFirstNonPHI());
 }
 
 //===----------------------------------------------------------------------===//
 //                        CallBrInst Implementation
 //===----------------------------------------------------------------------===//

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

if (Known.isAllOnes() && isAssumeWithEmptyBundle(*II))
       return eraseInstFromFunction(*II);
 
     // Update the cache of affected values for this assumption (we might be
     // here because we just simplified the condition).
     AC.updateAffectedValues(II);
     break;
   }
   case Intrinsic::experimental_gc_statepoint: {
-    auto &GCSP = *cast<GCStatepointInst>(II);
+    GCStatepointInst &GCSP = *cast<GCStatepointInst>(II);
+    SmallPtrSet<Value *, 32> LiveGcValues;
     for (const GCRelocateInst *Reloc : GCSP.getGCRelocates()) {
       GCRelocateInst &GCR = *const_cast<GCRelocateInst *>(Reloc);
 
       // Remove the relocation if unused.
       if (GCR.use_empty()) {
         eraseInstFromFunction(GCR);
         continue;
       }
for (const GCRelocateInst *Reloc : GCSP.getGCRelocates()) {
         auto *OpIntTy = GCR.getOperand(2)->getType();
         GCR.setOperand(2, ConstantInt::get(OpIntTy, GCR.getBasePtrIndex()));
       }
 
       // TODO: bitcast(relocate(p)) -> relocate(bitcast(p))
       // Canonicalize on the type from the uses to the defs
 
       // TODO: relocate((gep p, C, C2, ...)) -> gep(relocate(p), C, C2, ...)
+      LiveGcValues.insert(BasePtr);
+      LiveGcValues.insert(DerivedPtr);
     }
+    Optional<OperandBundleUse> Bundle =
+        GCSP.getOperandBundle(LLVMContext::OB_gc_live);
+    unsigned NumOfGCLives = LiveGcValues.size();
+    if (!Bundle.hasValue() || NumOfGCLives == Bundle->Inputs.size())
+      break;
+    // We can reduce the size of gc live bundle.
+    DenseMap<Value *, unsigned> Val2Idx;
+    std::vector<Value *> NewLiveGc;
+    for (unsigned I = 0, E = Bundle->Inputs.size(); I < E; ++I) {
+      Value *V = Bundle->Inputs[I];
+      if (Val2Idx.count(V))
+        continue;
+      if (LiveGcValues.count(V)) {
+        Val2Idx[V] = NewLiveGc.size();
+        NewLiveGc.push_back(V);
+      } else
+        Val2Idx[V] = NumOfGCLives;
+    }
+    // Update all gc.relocates
+    for (const GCRelocateInst *Reloc : GCSP.getGCRelocates()) {
+      GCRelocateInst &GCR = *const_cast<GCRelocateInst *>(Reloc);
+      Value *BasePtr = GCR.getBasePtr();
+      assert(Val2Idx.count(BasePtr) && Val2Idx[BasePtr] != NumOfGCLives &&
+             "Missed live gc for base pointer");
+      auto *OpIntTy1 = GCR.getOperand(1)->getType();
+      GCR.setOperand(1, ConstantInt::get(OpIntTy1, Val2Idx[BasePtr]));
+      Value *DerivedPtr = GCR.getDerivedPtr();
+      assert(Val2Idx.count(DerivedPtr) && Val2Idx[DerivedPtr] != NumOfGCLives &&
+             "Missed live gc for derived pointer");
+      auto *OpIntTy2 = GCR.getOperand(2)->getType();
+      GCR.setOperand(2, ConstantInt::get(OpIntTy2, Val2Idx[DerivedPtr]));
+    }
+    // Create new statepoint instruction.
+    OperandBundleDef NewBundle("gc-live", NewLiveGc);
+    if (isa<CallInst>(II))
+      return CallInst::CreateWithReplacedBundle(cast<CallInst>(II), NewBundle);
+    else
+      return InvokeInst::CreateWithReplacedBundle(cast<InvokeInst>(II),
+                                                  NewBundle);
     break;
   }
   case Intrinsic::experimental_guard: {
     // Is this guard followed by another guard?  We scan forward over a small
     // fixed window of instructions to handle common cases with conditions
     // computed between guards.
     Instruction *NextInst = II->getNextNonDebugInstruction();
     for (unsigned i = 0; i < GuardWideningWindow; i++) {

llvm/test/Transforms/InstCombine/statepoint-iter.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -instcombine-max-iterations=1 -S | FileCheck %s
 ; These tests check the optimizations specific to
 ; pointers being relocated at a statepoint.
 
 
 declare void @func()
 
 define i1 @test_null(i1 %cond) gc "statepoint-example" {
+; CHECK-LABEL: @test_null(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[LEFT:%.*]], label [[RIGHT:%.*]]
+; CHECK:       right:
+; CHECK-NEXT:    br label [[MERGE:%.*]]
+; CHECK:       left:
+; CHECK-NEXT:    [[SAFEPOINT_TOKEN:%.*]] = tail call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 0, i32 0, void ()* nonnull @func, i32 0, i32 0, i32 0, i32 0) [ "gc-live"() ]
+; CHECK-NEXT:    br label [[MERGE]]
+; CHECK:       merge:
+; CHECK-NEXT:    [[SAFEPOINT_TOKEN2:%.*]] = tail call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 0, i32 0, void ()* nonnull @func, i32 0, i32 0, i32 0, i32 0) [ "gc-live"() ]
+; CHECK-NEXT:    ret i1 true
+;
 entry:
   br i1 %cond, label %left, label %right
 
 right:
   br label %merge
 
 left:
   %safepoint_token = tail call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 0, i32 0, void ()* @func, i32 0, i32 0, i32 0, i32 0) ["gc-live" (i32* null)]
   %pnew = call i32* @llvm.experimental.gc.relocate.p1i32(token %safepoint_token,  i32 0, i32 0)
   br label %merge
 
 merge:
   %pnew_phi = phi i32* [null, %right], [%pnew, %left]
   %safepoint_token2 = tail call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 0, i32 0, void ()* @func, i32 0, i32 0, i32 0, i32 0) ["gc-live" (i32* %pnew_phi)]
   %pnew2 = call i32* @llvm.experimental.gc.relocate.p1i32(token %safepoint_token2,  i32 0, i32 0)
   %cmp = icmp eq i32* %pnew2, null
   ret i1 %cmp
-; CHECK-LABEL: test_null
-; CHECK-NOT: %pnew
-; CHECK-NOT: %pnew2
-; CHECK: ret i1 true
 }
 
 define i32* @test_undef(i1 %cond) gc "statepoint-example" {
+; CHECK-LABEL: @test_undef(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[LEFT:%.*]], label [[RIGHT:%.*]]
+; CHECK:       right:
+; CHECK-NEXT:    br label [[MERGE:%.*]]
+; CHECK:       left:
+; CHECK-NEXT:    [[SAFEPOINT_TOKEN:%.*]] = tail call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 0, i32 0, void ()* nonnull @func, i32 0, i32 0, i32 0, i32 0) [ "gc-live"() ]
+; CHECK-NEXT:    br label [[MERGE]]
+; CHECK:       merge:
+; CHECK-NEXT:    [[SAFEPOINT_TOKEN2:%.*]] = tail call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 0, i32 0, void ()* nonnull @func, i32 0, i32 0, i32 0, i32 0) [ "gc-live"() ]
+; CHECK-NEXT:    ret i32* undef
+;
 entry:
   br i1 %cond, label %left, label %right
 
 right:
   br label %merge
 
 left:
   %safepoint_token = tail call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 0, i32 0, void ()* @func, i32 0, i32 0, i32 0, i32 0) ["gc-live" (i32* undef)]
   %pnew = call i32* @llvm.experimental.gc.relocate.p1i32(token %safepoint_token,  i32 0, i32 0)
   br label %merge
 
 merge:
   %pnew_phi = phi i32* [undef, %right], [%pnew, %left]
   %safepoint_token2 = tail call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 0, i32 0, void ()* @func, i32 0, i32 0, i32 0, i32 0) ["gc-live" (i32* %pnew_phi)]
   %pnew2 = call i32* @llvm.experimental.gc.relocate.p1i32(token %safepoint_token2,  i32 0, i32 0)
   ret i32* %pnew2
-; CHECK-LABEL: test_undef
-; CHECK-NOT: %pnew
-; CHECK-NOT: %pnew2
-; CHECK: ret i32* undef
 }
 
 declare token @llvm.experimental.gc.statepoint.p0f_isVoidf(i64, i32, void ()*, i32, i32, ...)
 declare i32* @llvm.experimental.gc.relocate.p1i32(token, i32, i32)