[InstCombine] canonicalize bitcast after insertelement into undef

[spatel]

Summary

We have a partial transform in the opposite direction, so that needs to be removed while adding a more general transform that moves bitcast after insertelement.

The motivating case from PR45748:
https://bugs.llvm.org/show_bug.cgi?id=45748
...is the last test diff. In that example, we are triggering an existing bitcast transform, so we reduce the number of casts, and that should give us the ideal x86 codegen.

I'm not sure what to do about the mmx diffs. If the x86 backend is expecting something in particular, we need to specify that here (do we need to exclude/add the mmx type to either of these code diffs?).

[craig.topper]

I messed with mmx in this area once and regretted it. See rG5ebbabc1af360756f402203ba7704bb480f279a7

Is canonicalizing towards a vector type that wasn't mentioned in the IR the right way to go? Is that cast free for legal types on all targets? I think its potentially scalarized or becomes a load/store to stack temporary for illegal types in the backend.

[lebedev.ri]

In D79171#2013138, @craig.topper wrote:

I messed with mmx in this area once and regretted it. See rG5ebbabc1af360756f402203ba7704bb480f279a7

Is canonicalizing towards a vector type that wasn't mentioned in the IR the right way to go? Is that cast free for legal types on all targets? I think its potentially scalarized or becomes a load/store to stack temporary for illegal types in the backend.

It kinda sounds to me like such cases should use some custom intrinsics instead of generic bitcasts.

[spatel]

In D79171#2013138, @craig.topper wrote:

I messed with mmx in this area once and regretted it. See rG5ebbabc1af360756f402203ba7704bb480f279a7

Sufficiently scared. Will remove mmx diffs.

Is canonicalizing towards a vector type that wasn't mentioned in the IR the right way to go? Is that cast free for legal types on all targets? I think its potentially scalarized or becomes a load/store to stack temporary for illegal types in the backend.

I think we're fine here. The vector bitcast that we are creating is only changing the element type (FP <-> int), and we assume that both of those are ok since they are already in the code independent of this transform.
The vector cast is likely safer than the scalar cast on most targets because it's probably not causing any underlying register move (assuming vector registers hold both int and FP). So even if the vector type is illegal, the obvious lowering is a no-op cast to the legal vector type.

[spatel]

Patch updated:
Kept existing fold for mmx and added warning comment.

[lebedev.ri]

LG in general, but MMX stuff puzzles me, so would be good for @craig.topper to comment.

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

2487–2488 ↗

(On Diff #261290)

Am i reading this correctly that this is NFC in context of rG5ebbabc1af360756f402203ba7704bb480f279a7?
Can this be split off?

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1035

I guess this won't work for scalable vectors?
Can we somehow just replace the elt type in VecOp->getType() instead?

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

2487–2488 ↗

(On Diff #261290)

It's not quite NFC - this code diff results in the test diff in llvm/test/Transforms/InstCombine/bitcast-vec-canon.ll test "@c". That does seem like an independent improvement though, so I can break that into a preliminary commit.

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1035

I think this is safe for scalable vectors, so I better add a test. See similar construct annotated below - around line 1279.
cc @ctetreau

1279

Safe for scalable vectors?

[spatel]

Patch updated:

1. Pulled the MMX diff into a preliminary commit - rG2058c98715f6
2. Added a scalable vector test.

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

Instruction *InstCombiner::visitInsertElementInst(InsertElementInst &IE) {
   Value *VecOp    = IE.getOperand(0);
   Value *ScalarOp = IE.getOperand(1);
   Value *IdxOp    = IE.getOperand(2);
 
   if (auto *V = SimplifyInsertElementInst(
           VecOp, ScalarOp, IdxOp, SQ.getWithInstruction(&IE)))
     return replaceInstUsesWith(IE, V);
 
+  // If the scalar is bitcast and inserted into undef, do the insert in the
+  // source type followed by bitcast.
+  // TODO: Generalize for insert into any constant, not just undef?
+  Value *ScalarSrc;
+  if (match(VecOp, m_Undef()) &&
+      match(ScalarOp, m_OneUse(m_BitCast(m_Value(ScalarSrc)))) &&
+      (ScalarSrc->getType()->isIntegerTy() ||
+       ScalarSrc->getType()->isFloatingPointTy())) {
+    // inselt undef, (bitcast ScalarSrc), IdxOp -->
+    //   bitcast (inselt undef, ScalarSrc, IdxOp)
+    Type *ScalarTy = ScalarSrc->getType();
+    Type *VecTy = VectorType::get(ScalarTy, IE.getType()->getElementCount());
+    UndefValue *NewUndef = UndefValue::get(VecTy);
+    Value *NewInsElt = Builder.CreateInsertElement(NewUndef, ScalarSrc, IdxOp);
+    return new BitCastInst(NewInsElt, IE.getType());
+  }
+
   // If the vector and scalar are both bitcast from the same element type, do
   // the insert in that source type followed by bitcast.
-  Value *VecSrc, *ScalarSrc;
+  Value *VecSrc;
   if (match(VecOp, m_BitCast(m_Value(VecSrc))) &&
       match(ScalarOp, m_BitCast(m_Value(ScalarSrc))) &&
       (VecOp->hasOneUse() || ScalarOp->hasOneUse()) &&
       VecSrc->getType()->isVectorTy() && !ScalarSrc->getType()->isVectorTy() &&
       cast<VectorType>(VecSrc->getType())->getElementType() ==
           ScalarSrc->getType()) {
     // inselt (bitcast VecSrc), (bitcast ScalarSrc), IdxOp -->
     //   bitcast (inselt VecSrc, ScalarSrc, IdxOp)
switch (I->getOpcode()) {
     case Instruction::UIToFP:
     case Instruction::SIToFP:
     case Instruction::FPTrunc:
     case Instruction::FPExt: {
       // It's possible that the mask has a different number of elements from
       // the original cast. We recompute the destination type to match the mask.
       Type *DestTy = VectorType::get(
           I->getType()->getScalarType(),
           cast<VectorType>(NewOps[0]->getType())->getElementCount());
       assert(NewOps.size() == 1 && "cast with #ops != 1");
       return CastInst::Create(cast<CastInst>(I)->getOpcode(), NewOps[0], DestTy,
                               "", I);
     }
     case Instruction::GetElementPtr: {
       Value *Ptr = NewOps[0];
       ArrayRef<Value*> Idx = NewOps.slice(1);
       GetElementPtrInst *GEP = GetElementPtrInst::Create(

llvm/test/Transforms/InstCombine/bitcast-vec-canon.ll

 ; CHECK-NEXT:    ret double [[TMP0]]
 ;
 entry:
   %0 = bitcast x86_mmx %x to <1 x i64>
   %1 = bitcast <1 x i64> %0 to double
   ret double %1
 }
 
+; FP source is ok.
+
 define <3 x i64> @bitcast_inselt_undef(double %x, i32 %idx) {
 ; CHECK-LABEL: @bitcast_inselt_undef(
-; CHECK-NEXT:    [[XB:%.*]] = bitcast double [[X:%.*]] to i64
-; CHECK-NEXT:    [[I:%.*]] = insertelement <3 x i64> undef, i64 [[XB]], i32 [[IDX:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <3 x double> undef, double [[X:%.*]], i32 [[IDX:%.*]]
+; CHECK-NEXT:    [[I:%.*]] = bitcast <3 x double> [[TMP1]] to <3 x i64>
 ; CHECK-NEXT:    ret <3 x i64> [[I]]
 ;
   %xb = bitcast double %x to i64
   %i = insertelement <3 x i64> undef, i64 %xb, i32 %idx
   ret <3 x i64> %i
 }
 
+; Integer source is ok; index is anything.
+
 define <3 x float> @bitcast_inselt_undef_fp(i32 %x, i567 %idx) {
 ; CHECK-LABEL: @bitcast_inselt_undef_fp(
-; CHECK-NEXT:    [[XB:%.*]] = bitcast i32 [[X:%.*]] to float
-; CHECK-NEXT:    [[I:%.*]] = insertelement <3 x float> undef, float [[XB]], i567 [[IDX:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <3 x i32> undef, i32 [[X:%.*]], i567 [[IDX:%.*]]
+; CHECK-NEXT:    [[I:%.*]] = bitcast <3 x i32> [[TMP1]] to <3 x float>
 ; CHECK-NEXT:    ret <3 x float> [[I]]
 ;
   %xb = bitcast i32 %x to float
   %i = insertelement <3 x float> undef, float %xb, i567 %idx
   ret <3 x float> %i
 }
 
+define <vscale x 3 x float> @bitcast_inselt_undef_vscale(i32 %x, i567 %idx) {
+; CHECK-LABEL: @bitcast_inselt_undef_vscale(
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <vscale x 3 x i32> undef, i32 [[X:%.*]], i567 [[IDX:%.*]]
+; CHECK-NEXT:    [[I:%.*]] = bitcast <vscale x 3 x i32> [[TMP1]] to <vscale x 3 x float>
+; CHECK-NEXT:    ret <vscale x 3 x float> [[I]]
+;
+  %xb = bitcast i32 %x to float
+  %i = insertelement <vscale x 3 x float> undef, float %xb, i567 %idx
+  ret <vscale x 3 x float> %i
+}
+
 declare void @use(i64)
 
+; Negative test - extra use prevents canonicalization
+
 define <3 x i64> @bitcast_inselt_undef_extra_use(double %x, i32 %idx) {
 ; CHECK-LABEL: @bitcast_inselt_undef_extra_use(
 ; CHECK-NEXT:    [[XB:%.*]] = bitcast double [[X:%.*]] to i64
 ; CHECK-NEXT:    call void @use(i64 [[XB]])
 ; CHECK-NEXT:    [[I:%.*]] = insertelement <3 x i64> undef, i64 [[XB]], i32 [[IDX:%.*]]
 ; CHECK-NEXT:    ret <3 x i64> [[I]]
 ;
   %xb = bitcast double %x to i64
   call void @use(i64 %xb)
   %i = insertelement <3 x i64> undef, i64 %xb, i32 %idx
   ret <3 x i64> %i
 }
 
+; Negative test - source type must be scalar
+
 define <3 x i64> @bitcast_inselt_undef_vec_src(<2 x i32> %x, i32 %idx) {
 ; CHECK-LABEL: @bitcast_inselt_undef_vec_src(
 ; CHECK-NEXT:    [[XB:%.*]] = bitcast <2 x i32> [[X:%.*]] to i64
 ; CHECK-NEXT:    [[I:%.*]] = insertelement <3 x i64> undef, i64 [[XB]], i32 [[IDX:%.*]]
 ; CHECK-NEXT:    ret <3 x i64> [[I]]
 ;
   %xb = bitcast <2 x i32> %x to i64
   %i = insertelement <3 x i64> undef, i64 %xb, i32 %idx
   ret <3 x i64> %i
 }
 
+; Negative test - source type must be scalar
+
 define <3 x i64> @bitcast_inselt_undef_from_mmx(x86_mmx %x, i32 %idx) {
 ; CHECK-LABEL: @bitcast_inselt_undef_from_mmx(
 ; CHECK-NEXT:    [[XB:%.*]] = bitcast x86_mmx [[X:%.*]] to i64
 ; CHECK-NEXT:    [[I:%.*]] = insertelement <3 x i64> undef, i64 [[XB]], i32 [[IDX:%.*]]
 ; CHECK-NEXT:    ret <3 x i64> [[I]]
 ;
   %xb = bitcast x86_mmx %x to i64
   %i = insertelement <3 x i64> undef, i64 %xb, i32 %idx
   ret <3 x i64> %i
 }
 
+; Reduce number of casts
+
 define <2 x i64> @PR45748(double %x, double %y) {
 ; CHECK-LABEL: @PR45748(
-; CHECK-NEXT:    [[XB:%.*]] = bitcast double [[X:%.*]] to i64
-; CHECK-NEXT:    [[I0:%.*]] = insertelement <2 x i64> undef, i64 [[XB]], i32 0
-; CHECK-NEXT:    [[YB:%.*]] = bitcast double [[Y:%.*]] to i64
-; CHECK-NEXT:    [[I1:%.*]] = insertelement <2 x i64> [[I0]], i64 [[YB]], i32 1
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x double> undef, double [[X:%.*]], i32 0
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <2 x double> [[TMP1]], double [[Y:%.*]], i32 1
+; CHECK-NEXT:    [[I1:%.*]] = bitcast <2 x double> [[TMP2]] to <2 x i64>
 ; CHECK-NEXT:    ret <2 x i64> [[I1]]
 ;
   %xb = bitcast double %x to i64
   %i0 = insertelement <2 x i64> undef, i64 %xb, i32 0
   %yb = bitcast double %y to i64
   %i1 = insertelement <2 x i64> %i0, i64 %yb, i32 1
   ret <2 x i64> %i1
 }