[InstCombine] fold fake floating point vector extract to shift+trunc.

[jacquesguan]

Summary

This patch supports the FP part of D111082.

[RKSimon]

Testing more float types might be a good idea (bfloat / float / double / x86_fp80 / f128 / ppc_f128), even of they're negative tests - not sure how much the isDesirableIntType check is going to interfere?

[jacquesguan]

add more FP type test.

[spatel]

Is there a motivating codegen example for this transform? Or some other IR transform that will fire as a result of this transform?

In the case with a shift, we have an extra IR instruction, so this would be a rare fold that increases instruction count. Maybe that's justifiable just because we want to keep the symmetry with the other patterns, but it would be better to show some kind of win from this patch.

llvm/test/Transforms/InstCombine/extractelement.ll

512

This comment doesn't look accurate. The data layout for both RUN lines says i64 is legal, so that's why we allow converting i64 in the test above. And there's no exception for a 1-element vector. We probably should have a test like this to confirm:

define double @bitcast_fp64vec_index0(i64 %x) {
  %v = bitcast i64 %x to <1 x double>
  %r = extractelement <1 x double> %v, i8 0
  ret double %r
}

Also make sure there's no infinite loop potential with that kind of pattern...I just noticed an odd behavior for it in D125951.

Please pre-commit the baseline tests, so it's easier to see what is or is not changing. Add a RUN with a legal i128 to confirm this is behaving as expected?

[jacquesguan]

add precommit test.

[jacquesguan]

In D125750#3525445, @spatel wrote:

Is there a motivating codegen example for this transform? Or some other IR transform that will fire as a result of this transform?

In the case with a shift, we have an extra IR instruction, so this would be a rare fold that increases instruction count. Maybe that's justifiable just because we want to keep the symmetry with the other patterns, but it would be better to show some kind of win from this patch.

Mostly, it would be much cheaper if we use scalar shift + cast rather than bitcast + vector extractelemt, even the former might cause one more instruction in LLVMIR. For example in RISCV, the former one woule be lower to 3 scalar instructions, but the latter one would firstly move from GPR to vector register and then use 2 vector instruction to extract the element, it is truely much more expensive, even without counting the vector configuration instruction that should be insert for using vector instructions.

llvm/test/Transforms/InstCombine/extractelement.ll

512

Done.

[spatel]

In D125750#3533124, @jacquesguan wrote:

In D125750#3525445, @spatel wrote:

Is there a motivating codegen example for this transform? Or some other IR transform that will fire as a result of this transform?

In the case with a shift, we have an extra IR instruction, so this would be a rare fold that increases instruction count. Maybe that's justifiable just because we want to keep the symmetry with the other patterns, but it would be better to show some kind of win from this patch.

Mostly, it would be much cheaper if we use scalar shift + cast rather than bitcast + vector extractelemt, even the former might cause one more instruction in LLVMIR. For example in RISCV, the former one woule be lower to 3 scalar instructions, but the latter one would firstly move from GPR to vector register and then use 2 vector instruction to extract the element, it is truely much more expensive, even without counting the vector configuration instruction that should be insert for using vector instructions.

Yes, I understand the codegen motivation. I should have been more explicit though - that's generally not enough to justify an IR canonicalization if the backend could just as easily do this transform.

We really want to show that the change in IR leads to an improvement in analysis and/or results in even more optimization. Double-check, but we could probably add a test like this:
https://alive2.llvm.org/ce/z/77k-Zg

[jacquesguan]

Address comment.

[jacquesguan]

In D125750#3534196, @spatel wrote:

In D125750#3533124, @jacquesguan wrote:

In D125750#3525445, @spatel wrote:

Is there a motivating codegen example for this transform? Or some other IR transform that will fire as a result of this transform?

In the case with a shift, we have an extra IR instruction, so this would be a rare fold that increases instruction count. Maybe that's justifiable just because we want to keep the symmetry with the other patterns, but it would be better to show some kind of win from this patch.

Mostly, it would be much cheaper if we use scalar shift + cast rather than bitcast + vector extractelemt, even the former might cause one more instruction in LLVMIR. For example in RISCV, the former one woule be lower to 3 scalar instructions, but the latter one would firstly move from GPR to vector register and then use 2 vector instruction to extract the element, it is truely much more expensive, even without counting the vector configuration instruction that should be insert for using vector instructions.

Yes, I understand the codegen motivation. I should have been more explicit though - that's generally not enough to justify an IR canonicalization if the backend could just as easily do this transform.

We really want to show that the change in IR leads to an improvement in analysis and/or results in even more optimization. Double-check, but we could probably add a test like this:
https://alive2.llvm.org/ce/z/77k-Zg

Thanks, I added this test at the last.

[spatel]

llvm/test/Transforms/InstCombine/extractelement.ll

648

We were missing a transform that would reduce this, so I added it here:
482777123427

Please rebase/update to generate new output.

[jacquesguan]

Address comment.

llvm/test/Transforms/InstCombine/extractelement.ll

648

Done, thanks.

[spatel]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

Instruction *InstCombinerImpl::foldBitcastExtElt(ExtractElementInst &Ext) {
 
   ElementCount NumElts =
       cast<VectorType>(Ext.getVectorOperandType())->getElementCount();
   Type *DestTy = Ext.getType();
   bool IsBigEndian = DL.isBigEndian();
 
   // If we are casting an integer to vector and extracting a portion, that is
   // a shift-right and truncate.
-  // TODO: Allow FP dest type by casting the trunc to FP?
-  if (X->getType()->isIntegerTy() && DestTy->isIntegerTy() &&
+  if (X->getType()->isIntegerTy() &&
       isDesirableIntType(X->getType()->getPrimitiveSizeInBits())) {
     assert(isa<FixedVectorType>(Ext.getVectorOperand()->getType()) &&
            "Expected fixed vector type for bitcast from scalar integer");
 
     // Big endian requires adjusting the extract index since MSB is at index 0.
     // LittleEndian: extelt (bitcast i32 X to v4i8), 0 -> trunc i32 X to i8
     // BigEndian: extelt (bitcast i32 X to v4i8), 0 -> trunc i32 (X >> 24) to i8
     if (IsBigEndian)
       ExtIndexC = NumElts.getKnownMinValue() - 1 - ExtIndexC;
     unsigned ShiftAmountC = ExtIndexC * DestTy->getPrimitiveSizeInBits();
     if (!ShiftAmountC || Ext.getVectorOperand()->hasOneUse()) {
       Value *Lshr = Builder.CreateLShr(X, ShiftAmountC, "extelt.offset");
+      if (DestTy->isFloatingPointTy()) {
+        Type *DstIntTy = IntegerType::getIntNTy(
+            Lshr->getContext(), DestTy->getPrimitiveSizeInBits());
+        Value *Trunc = Builder.CreateTrunc(Lshr, DstIntTy);
+        return new BitCastInst(Trunc, DestTy);
+      }
       return new TruncInst(Lshr, DestTy);
     }
   }
 
   if (!X->getType()->isVectorTy())
     return nullptr;
 
   // If this extractelement is using a bitcast from a vector of the same number

llvm/test/Transforms/InstCombine/extractelement.ll

 ; ANY-NEXT:    [[R:%.*]] = extractelement <4 x i8> [[V]], i64 0
 ; ANY-NEXT:    ret i8 [[R]]
 ;
   %v = bitcast float %x to <4 x i8>
   %r = extractelement <4 x i8> %v, i8 0
   ret i8 %r
 }
 
-; negative test - can't have FP dest type without a cast
-
-define half @bitcast_fpvec_index0(i32 %x) {
-; ANY-LABEL: @bitcast_fpvec_index0(
-; ANY-NEXT:    [[V:%.*]] = bitcast i32 [[X:%.*]] to <2 x half>
-; ANY-NEXT:    [[R:%.*]] = extractelement <2 x half> [[V]], i64 0
-; ANY-NEXT:    ret half [[R]]
+define half @bitcast_fp16vec_index0(i32 %x) {
+; LE-LABEL: @bitcast_fp16vec_index0(
+; LE-NEXT:    [[TMP1:%.*]] = trunc i32 [[X:%.*]] to i16
+; LE-NEXT:    [[R:%.*]] = bitcast i16 [[TMP1]] to half
+; LE-NEXT:    ret half [[R]]
+;
+; BE-LABEL: @bitcast_fp16vec_index0(
+; BE-NEXT:    [[EXTELT_OFFSET:%.*]] = lshr i32 [[X:%.*]], 16
+; BE-NEXT:    [[TMP1:%.*]] = trunc i32 [[EXTELT_OFFSET]] to i16
+; BE-NEXT:    [[R:%.*]] = bitcast i16 [[TMP1]] to half
+; BE-NEXT:    ret half [[R]]
 ;
   %v = bitcast i32 %x to <2 x half>
   %r = extractelement <2 x half> %v, i8 0
   ret half %r
 }
 
+define half @bitcast_fp16vec_index1(i32 %x) {
+; LE-LABEL: @bitcast_fp16vec_index1(
+; LE-NEXT:    [[EXTELT_OFFSET:%.*]] = lshr i32 [[X:%.*]], 16
+; LE-NEXT:    [[TMP1:%.*]] = trunc i32 [[EXTELT_OFFSET]] to i16
+; LE-NEXT:    [[R:%.*]] = bitcast i16 [[TMP1]] to half
+; LE-NEXT:    ret half [[R]]
+;
+; BE-LABEL: @bitcast_fp16vec_index1(
+; BE-NEXT:    [[TMP1:%.*]] = trunc i32 [[X:%.*]] to i16
+; BE-NEXT:    [[R:%.*]] = bitcast i16 [[TMP1]] to half
+; BE-NEXT:    ret half [[R]]
+;
+  %v = bitcast i32 %x to <2 x half>
+  %r = extractelement <2 x half> %v, i8 1
+  ret half %r
+}
+
+define bfloat @bitcast_bfp16vec_index0(i32 %x) {
+; LE-LABEL: @bitcast_bfp16vec_index0(
+; LE-NEXT:    [[TMP1:%.*]] = trunc i32 [[X:%.*]] to i16
+; LE-NEXT:    [[R:%.*]] = bitcast i16 [[TMP1]] to bfloat
+; LE-NEXT:    ret bfloat [[R]]
+;
+; BE-LABEL: @bitcast_bfp16vec_index0(
+; BE-NEXT:    [[EXTELT_OFFSET:%.*]] = lshr i32 [[X:%.*]], 16
+; BE-NEXT:    [[TMP1:%.*]] = trunc i32 [[EXTELT_OFFSET]] to i16
+; BE-NEXT:    [[R:%.*]] = bitcast i16 [[TMP1]] to bfloat
+; BE-NEXT:    ret bfloat [[R]]
+;
+  %v = bitcast i32 %x to <2 x bfloat>
+  %r = extractelement <2 x bfloat> %v, i8 0
+  ret bfloat %r
+}
+
+define bfloat @bitcast_bfp16vec_index1(i32 %x) {
+; LE-LABEL: @bitcast_bfp16vec_index1(
+; LE-NEXT:    [[EXTELT_OFFSET:%.*]] = lshr i32 [[X:%.*]], 16
+; LE-NEXT:    [[TMP1:%.*]] = trunc i32 [[EXTELT_OFFSET]] to i16
+; LE-NEXT:    [[R:%.*]] = bitcast i16 [[TMP1]] to bfloat
+; LE-NEXT:    ret bfloat [[R]]
+;
+; BE-LABEL: @bitcast_bfp16vec_index1(
+; BE-NEXT:    [[TMP1:%.*]] = trunc i32 [[X:%.*]] to i16
+; BE-NEXT:    [[R:%.*]] = bitcast i16 [[TMP1]] to bfloat
+; BE-NEXT:    ret bfloat [[R]]
+;
+  %v = bitcast i32 %x to <2 x bfloat>
+  %r = extractelement <2 x bfloat> %v, i8 1
+  ret bfloat %r
+}
+
+define float @bitcast_fp32vec_index0(i64 %x) {
+; LE-LABEL: @bitcast_fp32vec_index0(
+; LE-NEXT:    [[TMP1:%.*]] = trunc i64 [[X:%.*]] to i32
+; LE-NEXT:    [[R:%.*]] = bitcast i32 [[TMP1]] to float
+; LE-NEXT:    ret float [[R]]
+;
+; BE-LABEL: @bitcast_fp32vec_index0(
+; BE-NEXT:    [[EXTELT_OFFSET:%.*]] = lshr i64 [[X:%.*]], 32
+; BE-NEXT:    [[TMP1:%.*]] = trunc i64 [[EXTELT_OFFSET]] to i32
+; BE-NEXT:    [[R:%.*]] = bitcast i32 [[TMP1]] to float
+; BE-NEXT:    ret float [[R]]
+;
+  %v = bitcast i64 %x to <2 x float>
+  %r = extractelement <2 x float> %v, i8 0
+  ret float %r
+}
+
+define float @bitcast_fp32vec_index1(i64 %x) {
+; LE-LABEL: @bitcast_fp32vec_index1(
+; LE-NEXT:    [[EXTELT_OFFSET:%.*]] = lshr i64 [[X:%.*]], 32
+; LE-NEXT:    [[TMP1:%.*]] = trunc i64 [[EXTELT_OFFSET]] to i32
+; LE-NEXT:    [[R:%.*]] = bitcast i32 [[TMP1]] to float
+; LE-NEXT:    ret float [[R]]
+;
+; BE-LABEL: @bitcast_fp32vec_index1(
+; BE-NEXT:    [[TMP1:%.*]] = trunc i64 [[X:%.*]] to i32
+; BE-NEXT:    [[R:%.*]] = bitcast i32 [[TMP1]] to float
+; BE-NEXT:    ret float [[R]]
+;
+  %v = bitcast i64 %x to <2 x float>
+  %r = extractelement <2 x float> %v, i8 1
+  ret float %r
+}
+
+; negative test - only support the width <= 32
+
+define double @bitcast_fp64vec_index0(i128 %x) {
+; ANY-LABEL: @bitcast_fp64vec_index0(
+; ANY-NEXT:    [[V:%.*]] = bitcast i128 [[X:%.*]] to <2 x double>
+; ANY-NEXT:    [[R:%.*]] = extractelement <2 x double> [[V]], i64 0
+; ANY-NEXT:    ret double [[R]]
+;
+  %v = bitcast i128 %x to <2 x double>
+  %r = extractelement <2 x double> %v, i8 0
+  ret double %r
+}
+
+; negative test - only support the width <= 32
+
+define double @bitcast_fp64vec_index1(i128 %x) {
+; ANY-LABEL: @bitcast_fp64vec_index1(
+; ANY-NEXT:    [[V:%.*]] = bitcast i128 [[X:%.*]] to <2 x double>
+; ANY-NEXT:    [[R:%.*]] = extractelement <2 x double> [[V]], i64 1
+; ANY-NEXT:    ret double [[R]]
+;
+  %v = bitcast i128 %x to <2 x double>
+  %r = extractelement <2 x double> %v, i8 1
+  ret double %r
+}
+
+; negative test - only support the width <= 32
+
+define x86_fp80 @bitcast_x86fp80vec_index0(i160 %x) {
+; ANY-LABEL: @bitcast_x86fp80vec_index0(
+; ANY-NEXT:    [[V:%.*]] = bitcast i160 [[X:%.*]] to <2 x x86_fp80>
+; ANY-NEXT:    [[R:%.*]] = extractelement <2 x x86_fp80> [[V]], i64 0
+; ANY-NEXT:    ret x86_fp80 [[R]]
+;
+  %v = bitcast i160 %x to <2 x x86_fp80>
+  %r = extractelement <2 x x86_fp80> %v, i8 0
+  ret x86_fp80 %r
+}
+
+; negative test - only support the width <= 32
+
+define x86_fp80 @bitcast_x86fp80vec_index1(i160 %x) {
+; ANY-LABEL: @bitcast_x86fp80vec_index1(
+; ANY-NEXT:    [[V:%.*]] = bitcast i160 [[X:%.*]] to <2 x x86_fp80>
+; ANY-NEXT:    [[R:%.*]] = extractelement <2 x x86_fp80> [[V]], i64 1
+; ANY-NEXT:    ret x86_fp80 [[R]]
+;
+  %v = bitcast i160 %x to <2 x x86_fp80>
+  %r = extractelement <2 x x86_fp80> %v, i8 1
+  ret x86_fp80 %r
+}
+
+; negative test - only support the width <= 32
+
+define fp128 @bitcast_fp128vec_index0(i256 %x) {
+; ANY-LABEL: @bitcast_fp128vec_index0(
+; ANY-NEXT:    [[V:%.*]] = bitcast i256 [[X:%.*]] to <2 x fp128>
+; ANY-NEXT:    [[R:%.*]] = extractelement <2 x fp128> [[V]], i64 0
+; ANY-NEXT:    ret fp128 [[R]]
+;
+  %v = bitcast i256 %x to <2 x fp128>
+  %r = extractelement <2 x fp128> %v, i8 0
+  ret fp128 %r
+}
+
+; negative test - only support the width <= 32
+
+define fp128 @bitcast_fp128vec_index1(i256 %x) {
+; ANY-LABEL: @bitcast_fp128vec_index1(
+; ANY-NEXT:    [[V:%.*]] = bitcast i256 [[X:%.*]] to <2 x fp128>
+; ANY-NEXT:    [[R:%.*]] = extractelement <2 x fp128> [[V]], i64 1
+; ANY-NEXT:    ret fp128 [[R]]
+;
+  %v = bitcast i256 %x to <2 x fp128>
+  %r = extractelement <2 x fp128> %v, i8 1
+  ret fp128 %r
+}
+
+; negative test - only support the width <= 32
+
+define ppc_fp128 @bitcast_ppcfp128vec_index0(i256 %x) {
+; ANY-LABEL: @bitcast_ppcfp128vec_index0(
+; ANY-NEXT:    [[V:%.*]] = bitcast i256 [[X:%.*]] to <2 x ppc_fp128>
+; ANY-NEXT:    [[R:%.*]] = extractelement <2 x ppc_fp128> [[V]], i64 0
+; ANY-NEXT:    ret ppc_fp128 [[R]]
+;
+  %v = bitcast i256 %x to <2 x ppc_fp128>
+  %r = extractelement <2 x ppc_fp128> %v, i8 0
+  ret ppc_fp128 %r
+}
+
+; negative test - only support the width <= 32
+
+define ppc_fp128 @bitcast_ppcfp128vec_index1(i256 %x) {
+; ANY-LABEL: @bitcast_ppcfp128vec_index1(
+; ANY-NEXT:    [[V:%.*]] = bitcast i256 [[X:%.*]] to <2 x ppc_fp128>
+; ANY-NEXT:    [[R:%.*]] = extractelement <2 x ppc_fp128> [[V]], i64 1
+; ANY-NEXT:    ret ppc_fp128 [[R]]
+;
+  %v = bitcast i256 %x to <2 x ppc_fp128>
+  %r = extractelement <2 x ppc_fp128> %v, i8 1
+  ret ppc_fp128 %r
+}
+
 ; negative test - need constant index
 
 define i8 @bitcast_scalar_index_variable(i32 %x, i64 %y) {
 ; ANY-LABEL: @bitcast_scalar_index_variable(
 ; ANY-NEXT:    [[V:%.*]] = bitcast i32 [[X:%.*]] to <4 x i8>
 ; ANY-NEXT:    [[R:%.*]] = extractelement <4 x i8> [[V]], i64 [[Y:%.*]]
 ; ANY-NEXT:    ret i8 [[R]]
 ;

 ; BE-NEXT:    call void @use(<8 x i8> [[V]])
 ; BE-NEXT:    [[R:%.*]] = extractelement <8 x i8> [[V]], i64 0
 ; BE-NEXT:    ret i8 [[R]]
 ;
   %v = bitcast i64 %x to <8 x i8>
   call void @use(<8 x i8> %v)
   %r = extractelement <8 x i8> %v, i64 0
   ret i8 %r
 }