[RS4GC] Handle uses of extractelement for conversion from vector to scalar base

[anna]

Summary

As mentioned in the comments, extractelement is special
since we actually want a scalar base for that element we extracted from
the vector (i.e. not a vector base).
This same logic should apply to uses of the extractelement such as phis
and selects which have the same BDV as the extractelement.

Added testcase showcases the problem where the BDV identification chokes
on the incorrect cast from vector to scalar for the phi use of
extractelement.

Tests-Run: make check

[anna]

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

959

Just realized this won't work in all cases since the base ptr identified may not dominate the derived ptr. Needs fixing.

[anna]

fixed dominator error.

[dantrushin]

LGTM, but what about test failure? Please fix it

[anna]

In D75704#1909467, @dantrushin wrote:

LGTM, but what about test failure? Please fix it

will do, thanks! forgot to update the test once I tested the patch internally.

[anna]

changed to early-return, updated comments and test.

[anna]

removed extra line. cleaned up whiitespace. NFC w.r.t. prev diff.

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

if (State.isBase() && isa<ExtractElementInst>(I) &&
       auto *EE = cast<ExtractElementInst>(I);
       // TODO: In many cases, the new instruction is just EE itself.  We should
       // exploit this, but can't do it here since it would break the invariant
       // about the BDV not being known to be a base.
       auto *BaseInst = ExtractElementInst::Create(
           State.getBaseValue(), EE->getIndexOperand(), "base_ee", EE);
       BaseInst->setMetadata("is_base_value", MDNode::get(I->getContext(), {}));
       States[I] = BDVState(BDVState::Base, BaseInst);
+      // Uses of the above EE instruction which have a scalar base and the same
+      // state as the EE instruction's (i.e. vector base) should also be updated
+      // to have this new base.
+      for (User *U : I->users()) {
+        auto *UseI = dyn_cast<Instruction>(U);
+        if (!UseI || !States.count(UseI))
+          continue;
+        if (!isa<VectorType>(UseI->getType()) && States[UseI] == State)
+          States[UseI] = BDVState(BDVState::Base, BaseInst);
+      }
     }
 
     // Since we're joining a vector and scalar base, they can never be the
     // same.  As a result, we should always see insert element having reached
     // the conflict state.
     assert(!isa<InsertElementInst>(I) || State.isConflict());
 
     if (!State.isConflict())
#endif
   // know it's base defining value is a base, or have inserted a new
   // instruction to propagate the base of it's BDV and have entered that newly
   // introduced instruction into the state table.  In either case, we are
   // assured to be able to determine an instruction which produces it's base
   // pointer.
   auto getBaseForInput = [&](Value *Input, Instruction *InsertPt) {
     Value *BDV = findBaseOrBDV(Input, Cache);
     Value *Base = nullptr;
-    if (isKnownBaseResult(BDV)) {
+    // Avoid the optimization of choosing a KnownBase BDV when the BDV is a
+    // vector type, while the input is scalar.
+    const bool VectorTypeBDV =
+        isa<VectorType>(BDV->getType()) && !isa<VectorType>(Input->getType());
+    if (isKnownBaseResult(BDV) && !VectorTypeBDV) {
       Base = BDV;
     } else {
       // Either conflict or base.
       assert(States.count(BDV));
       Base = States[BDV].getBaseValue();
     }
     assert(Base && "Can't be null");
     // The cast is needed since base traversal may strip away bitcasts

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

 ;
 entry:
   %base.i32 = bitcast <2 x i8 addrspace(1)*> %base to <2 x i32 addrspace(1)*>
   %vec = getelementptr i32, <2 x i32 addrspace(1)*> %base.i32, <2 x i64> %offsets
   call void @do_safepoint()
   %ptr = extractelement <2 x i32 addrspace(1)*> %vec, i32 0
   ret i32 addrspace(1)* %ptr
 }
+
+define void @test6() gc "statepoint-example" {
+; CHECK-LABEL: @test6(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    br label [[HEADER:%.*]]
+; CHECK:       header:
+; CHECK-NEXT:    [[TMP_BASE:%.*]] = phi i8 addrspace(1)* [ [[BASE_EE:%.*]], [[LATCH:%.*]] ], [ null, [[BB:%.*]] ], !is_base_value !0
+; CHECK-NEXT:    [[TMP:%.*]] = phi i8 addrspace(1)* [ [[TMP6:%.*]], [[LATCH]] ], [ undef, [[BB]] ]
+; CHECK-NEXT:    br label [[BB10:%.*]]
+; CHECK:       bb10:
+; CHECK-NEXT:    [[STATEPOINT_TOKEN:%.*]] = call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 2882400000, i32 0, void ()* @spam, i32 0, i32 0, i32 0, i32 1, i8 addrspace(1)* [[TMP]], i8 addrspace(1)* [[TMP]], i8 addrspace(1)* [[TMP_BASE]])
+; CHECK-NEXT:    [[TMP_RELOCATED:%.*]] = call coldcc i8 addrspace(1)* @llvm.experimental.gc.relocate.p1i8(token [[STATEPOINT_TOKEN]], i32 9, i32 8)
+; CHECK-NEXT:    [[TMP_BASE_RELOCATED:%.*]] = call coldcc i8 addrspace(1)* @llvm.experimental.gc.relocate.p1i8(token [[STATEPOINT_TOKEN]], i32 9, i32 9)
+; CHECK-NEXT:    br label [[BB25:%.*]]
+; CHECK:       bb25:
+; CHECK-NEXT:    [[STATEPOINT_TOKEN1:%.*]] = call token (i64, i32, <2 x i8 addrspace(1)*> ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_v2p1i8f(i64 2882400000, i32 0, <2 x i8 addrspace(1)*> ()* @baz, i32 0, i32 0, i32 0, i32 0)
+; CHECK-NEXT:    [[TMP262:%.*]] = call <2 x i8 addrspace(1)*> @llvm.experimental.gc.result.v2p1i8(token [[STATEPOINT_TOKEN1]])
+; CHECK-NEXT:    [[BASE_EE]] = extractelement <2 x i8 addrspace(1)*> [[TMP262]], i32 0, !is_base_value !0
+; CHECK-NEXT:    [[TMP27:%.*]] = extractelement <2 x i8 addrspace(1)*> [[TMP262]], i32 0
+; CHECK-NEXT:    br i1 undef, label [[BB7:%.*]], label [[LATCH]]
+; CHECK:       bb7:
+; CHECK-NEXT:    br label [[LATCH]]
+; CHECK:       latch:
+; CHECK-NEXT:    [[TMP6]] = phi i8 addrspace(1)* [ [[TMP27]], [[BB25]] ], [ [[TMP27]], [[BB7]] ]
+; CHECK-NEXT:    br label [[HEADER]]
+;
+bb:
+  br label %header
+
+header:                                              ; preds = %latch, %bb
+  %tmp = phi i8 addrspace(1)* [ %tmp6, %latch ], [ undef, %bb ]
+  br label %bb10
+
+bb10:                                             ; preds = %bb2
+  call void @spam() [ "deopt"(i8 addrspace(1)* %tmp) ]
+  br label %bb25
+
+bb25:                                             ; preds = %bb24
+  %tmp26 = call <2 x i8 addrspace(1)*> @baz()
+  %tmp27 = extractelement <2 x i8 addrspace(1)*> %tmp26, i32 0
+  br i1 undef, label %bb7, label %latch
+
+bb7:                                              ; preds = %bb25
+  br label %latch
+
+latch:                                              ; preds = %bb25, %bb7
+  %tmp6 = phi i8 addrspace(1)* [ %tmp27, %bb25 ], [ %tmp27, %bb7 ]
+  br label %header
+}
+
+declare void @spam()
+declare <2 x i8 addrspace(1)*> @baz()