[RS4GC] Avoid crashing on gep scalar, vector_idx

[reames]

Summary

This is an alternative to https://reviews.llvm.org/D57103.

The issue at hand is that the base rewriting algorithm doesn't consider the fact that GEPs can turn a scalar input into a vector of outputs. We had handling for scalar GEPs and fully vector GEPs (i.e. all vector operands), but not the scalar-base + vector-index forms. A true fix here requires treating GEP analogously to extractelement or shufflevector.

This patch is merely a workaround. It simply hides the crash at the cost of some ugly code gen for this presumable very rare pattern.

[anna]

LGTM w/ comments.

lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

2628 ↗

(On Diff #183244)

do you think it makes sense to add an assert here as the else part:

else {
   assert(VF == I.getOperand(i)->getType()->getVectorNumElements() && "expected same VF for all operands!");
}

Btw, this property is also verified in the IRVerifier when we come across a vector GEP (vector GEP is one where *any* of the GEP operands is a vector type).

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

if (auto *Cond = getConditionInst(TI))
       // TODO: Handle more than just ICmps here.  We should be able to move
       // most instructions without side effects or memory access.
       if (isa<ICmpInst>(Cond) && Cond->hasOneUse()) {
         MadeChange = true;
         Cond->moveBefore(TI);
       }
   }
 
+  // Nasty workaround - The base computation code in the main algorithm doesn't
+  // consider the fact that a GEP can be used to convert a scalar to a vector.
+  // The right fix for this is to integrate GEPs into the base rewriting
+  // algorithm properly, this is just a short term workaround to prevent
+  // crashes by canonicalizing such GEPs into fully vector GEPs.
+  for (Instruction &I : instructions(F)) {
+    if (!isa<GetElementPtrInst>(I))
+      continue;
+
+    unsigned VF = 0;
+    bool HasScalarOperand = false;
+    for (unsigned i = 0; i < I.getNumOperands(); i++)
+      if (I.getOperand(i)->getType()->isVectorTy())
+        VF = I.getOperand(i)->getType()->getVectorNumElements();
+      else
+        HasScalarOperand = true;
+
+    if (HasScalarOperand && VF != 0) {
+      IRBuilder<> B(&I);
+      for (unsigned i = 0; i < I.getNumOperands(); i++)
+        if (!I.getOperand(i)->getType()->isVectorTy()) {
+          auto *Splat = B.CreateVectorSplat(VF, I.getOperand(i));
+          I.setOperand(i, Splat);
+          MadeChange = true;
+        }
+    }
+  }
+
   MadeChange |= insertParsePoints(F, DT, TTI, ParsePointNeeded);
   return MadeChange;
 }
 
 // liveness computation via standard dataflow
 // -------------------------------------------------------------------
 
 // TODO: Consider using bitvectors for liveness, the set of potentially

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

 }
 declare void @do_safepoint()
 
 define void @test11(<4 x i64 addrspace(1)*> %vec1) gc "statepoint-example" {
 ; CHECK-LABEL: @test11(
 ; CHECK: @llvm.experimental.gc.statepoint.p0f_isVoidf{{.*}}<4 x i64 addrspace(1)*> %vec1)
 ; CHECK: %vec1.relocated = call coldcc <4 x i8 addrspace(1)*> @llvm.experimental.gc.relocate.v4p1i8
 ; CHECK: %vec1.relocated.casted = bitcast <4 x i8 addrspace(1)*> %vec1.relocated to <4 x i64 addrspace(1)*>
-; CHECK: %vec2.remat = getelementptr i64, <4 x i64 addrspace(1)*> %vec1.relocated.casted, i32 1024
+; CHECK: %vec2.remat = getelementptr i64, <4 x i64 addrspace(1)*> %vec1.relocated.casted, <4 x i32> <i32 1024, i32 1024, i32 1024, i32 1024>
 ; CHECK: call void @use_vec(<4 x i64 addrspace(1)*> %vec2.remat)
 entry:
   %vec2 = getelementptr i64, <4 x i64 addrspace(1)*> %vec1, i32 1024
   call void @do_safepoint() [ "deopt"(i32 0, i32 -1, i32 0, i32 0, i32 0) ]
   call void @use_vec(<4 x i64 addrspace(1) *> %vec2)
   ret void
 }
 

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

+; RUN: opt < %s -rewrite-statepoints-for-gc -S | FileCheck  %s
+; RUN: opt < %s -passes=rewrite-statepoints-for-gc -S | FileCheck  %s
+
+declare void @do_safepoint()
+declare i8 addrspace(1)* @def_ptr()
+
+define i32 addrspace(1)* @test1(i8 addrspace(1)* %base1, <2 x i64> %offsets) gc "statepoint-example" {
+entry:
+  br i1 undef, label %first, label %second
+
+first:
+  %base2 = call i8 addrspace(1)* @def_ptr() [ "deopt"(i32 0, i32 -1, i32 0, i32 0, i32 0) ]
+  br label %second
+
+second:
+; CHECK-LABEL: @test1(
+; CHECK: gc.statepoint
+; CHECK-DAG: (%ptr.base, %ptr)
+; CHECK-DAG: (%ptr.base, %ptr.base)
+  %phi = phi i8 addrspace(1)* [ %base1, %entry ], [ %base2, %first ]
+  %base.i32 = bitcast i8 addrspace(1)* %phi to i32 addrspace(1)*
+  %vec = getelementptr i32, i32 addrspace(1)* %base.i32, <2 x i64> %offsets
+  %ptr = extractelement <2 x i32 addrspace(1)*> %vec, i32 1
+  call void @do_safepoint() [ "deopt"(i32 0, i32 -1, i32 0, i32 0, i32 0) ]
+  ret i32 addrspace(1)* %ptr
+}