InstCombine: Scalarize single use icmp/fcmp

[spatel]

I think this is the right change for IR (I'm trying to do something similar with D50992 but that's been held up while we try to get the backend in shape). .

1. I see that the new code is copying the existing handling of binops around line 315, but can we update this to use 'match()' and reduce the duplication?
2. There should be a test that shows the benefit of using 'isCheapToScalarize()'. IIUC, that would involve a sequence with multiple binops and/or compares between an insertelement and the trailing extract.
3. Please commit the tests with baseline checks as a preliminary commit, so we only show the diffs in this patch.

[arsenm]

In D55182#1316746, @spatel wrote:

I think this is the right change for IR (I'm trying to do something similar with D50992 but that's been held up while we try to get the backend in shape). .

1. I see that the new code is copying the existing handling of binops around line 315, but can we update this to use 'match()' and reduce the duplication?
2. There should be a test that shows the benefit of using 'isCheapToScalarize()'. IIUC, that would involve a sequence with multiple binops and/or compares between an insertelement and the trailing extract.
3. Please commit the tests with baseline checks as a preliminary commit, so we only show the diffs in this patch.

I don't think there's much of a useful way to combine this with the BinOp handling. At most it's going to share the pair of extracts, and would require more ugly code when creating the new operation so I think it would be overall worse

[arsenm]

Rearrange code, apply on top of existing tests

[spatel]

In D55182#1316746, @spatel wrote:

2. There should be a test that shows the benefit of using 'isCheapToScalarize()'. IIUC, that would involve a sequence with multiple binops and/or compares between an insertelement and the trailing extract.

I added a binop test like this here:
rL348417
...and it shows a missed opportunity. AFAICT, there was no test coverage for anything beyond the simplest cases for isCheapToScalarize(), so please do add something like that for cmp.

Also, I used 'match' to reduce the binop code:
rL348418

cmp should be similar:

Value *X, *Y;
CmpInst::Predicate Pred;
if (match(SrcVec, m_Cmp(Pred, m_Value(X), m_Value(Y))) &&
    cheapToScalarize(SrcVec, IndexC)) {
  // extelt (cmp X, Y), Index --> cmp (extelt X, Index), (extelt Y, Index)
  Value *E0 = Builder.CreateExtractElement(X, Index);
  Value *E1 = Builder.CreateExtractElement(Y, Index);
  return CmpInst::Create(cast<CmpInst>(SrcVec)->getOpcode(), Pred, E0, E1);
}

[arsenm]

Added more tests in r348456, use match

[spatel]

LGTM - but I think there's a bug in cheapToScalarize() that I didn't notice before. Please put a FIXME note in there.

test/Transforms/InstCombine/scalarization.ll

140–143

This doesn't look right to start with. I don't think we can justify leaving more instructions than we started with as a target-independent combine.

[arsenm]

r348801

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

Instruction *InstCombiner::visitExtractElementInst(ExtractElementInst &EI) {
   if (match(SrcVec, m_BinOp(BO)) && cheapToScalarize(SrcVec, IndexC)) {
     // extelt (binop X, Y), Index --> binop (extelt X, Index), (extelt Y, Index)
     Value *X = BO->getOperand(0), *Y = BO->getOperand(1);
     Value *E0 = Builder.CreateExtractElement(X, Index);
     Value *E1 = Builder.CreateExtractElement(Y, Index);
     return BinaryOperator::CreateWithCopiedFlags(BO->getOpcode(), E0, E1, BO);
   }
 
+  Value *X, *Y;
+  CmpInst::Predicate Pred;
+  if (match(SrcVec, m_Cmp(Pred, m_Value(X), m_Value(Y))) &&
+      cheapToScalarize(SrcVec, IndexC)) {
+    // extelt (cmp X, Y), Index --> cmp (extelt X, Index), (extelt Y, Index)
+    Value *E0 = Builder.CreateExtractElement(X, Index);
+    Value *E1 = Builder.CreateExtractElement(Y, Index);
+    return CmpInst::Create(cast<CmpInst>(SrcVec)->getOpcode(), Pred, E0, E1);
+  }
+
   if (auto *I = dyn_cast<Instruction>(SrcVec)) {
     if (auto *IE = dyn_cast<InsertElementInst>(I)) {
       // Extracting the inserted element?
       if (IE->getOperand(2) == Index)
         return replaceInstUsesWith(EI, IE->getOperand(1));
       // If the inserted and extracted elements are constants, they must not
       // be the same value, extract from the pre-inserted value instead.
       if (isa<Constant>(IE->getOperand(2)) && IndexC) {

test/Transforms/InstCombine/scalarization.ll

 ; CHECK-NEXT:    ret float [[R]]
 ;
   %r = extractelement <4 x float> <float 1.0, float 2.0, float 3.0, float 4.0>, i32 %y
   ret float %r
 }
 
 define i1 @cheap_to_extract_icmp(<4 x i32> %x, <4 x i1> %y) {
 ; CHECK-LABEL: @cheap_to_extract_icmp(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <4 x i32> [[X:%.*]], zeroinitializer
-; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <4 x i1> [[CMP]], i32 2
-; CHECK-NEXT:    [[TMP2:%.*]] = extractelement <4 x i1> [[Y:%.*]], i32 2
-; CHECK-NEXT:    [[R:%.*]] = and i1 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <4 x i32> [[X:%.*]], i32 2
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i32 [[TMP1]], 0
+; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <4 x i1> [[Y:%.*]], i32 2
+; CHECK-NEXT:    [[R:%.*]] = and i1 [[TMP2]], [[TMP3]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %cmp = icmp eq <4 x i32> %x, zeroinitializer
   %and = and <4 x i1> %cmp, %y
   %r = extractelement <4 x i1> %and, i32 2
   ret i1 %r
 }
 
 define i1 @cheap_to_extract_fcmp(<4 x float> %x, <4 x i1> %y) {
 ; CHECK-LABEL: @cheap_to_extract_fcmp(
-; CHECK-NEXT:    [[CMP:%.*]] = fcmp oeq <4 x float> [[X:%.*]], zeroinitializer
-; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <4 x i1> [[CMP]], i32 2
-; CHECK-NEXT:    [[TMP2:%.*]] = extractelement <4 x i1> [[Y:%.*]], i32 2
-; CHECK-NEXT:    [[R:%.*]] = and i1 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <4 x float> [[X:%.*]], i32 2
+; CHECK-NEXT:    [[TMP2:%.*]] = fcmp oeq float [[TMP1]], 0.000000e+00
+; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <4 x i1> [[Y:%.*]], i32 2
+; CHECK-NEXT:    [[R:%.*]] = and i1 [[TMP2]], [[TMP3]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %cmp = fcmp oeq <4 x float> %x, zeroinitializer
   %and = and <4 x i1> %cmp, %y
   %r = extractelement <4 x i1> %and, i32 2
   ret i1 %r
 }
 
 define i1 @extractelt_vector_icmp_constrhs(<2 x i32> %arg) {
 ; CHECK-LABEL: @extractelt_vector_icmp_constrhs(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <2 x i32> [[ARG:%.*]], zeroinitializer
-; CHECK-NEXT:    [[EXT:%.*]] = extractelement <2 x i1> [[CMP]], i32 0
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x i32> [[ARG:%.*]], i32 0
+; CHECK-NEXT:    [[EXT:%.*]] = icmp eq i32 [[TMP1]], 0
 ; CHECK-NEXT:    ret i1 [[EXT]]
 ;
   %cmp = icmp eq <2 x i32> %arg, zeroinitializer
   %ext = extractelement <2 x i1> %cmp, i32 0
   ret i1 %ext
 }
 
 define i1 @extractelt_vector_fcmp_constrhs(<2 x float> %arg) {
 ; CHECK-LABEL: @extractelt_vector_fcmp_constrhs(
-; CHECK-NEXT:    [[CMP:%.*]] = fcmp oeq <2 x float> [[ARG:%.*]], zeroinitializer
-; CHECK-NEXT:    [[EXT:%.*]] = extractelement <2 x i1> [[CMP]], i32 0
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x float> [[ARG:%.*]], i32 0
+; CHECK-NEXT:    [[EXT:%.*]] = fcmp oeq float [[TMP1]], 0.000000e+00
 ; CHECK-NEXT:    ret i1 [[EXT]]
 ;
   %cmp = fcmp oeq <2 x float> %arg, zeroinitializer
   %ext = extractelement <2 x i1> %cmp, i32 0
   ret i1 %ext
 }
 
 define i1 @extractelt_vector_icmp_constrhs_dynidx(<2 x i32> %arg, i32 %idx) {
 ; CHECK-LABEL: @extractelt_vector_icmp_constrhs_dynidx(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <2 x i32> [[ARG:%.*]], zeroinitializer
-; CHECK-NEXT:    [[EXT:%.*]] = extractelement <2 x i1> [[CMP]], i32 [[IDX:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x i32> [[ARG:%.*]], i32 [[IDX:%.*]]
+; CHECK-NEXT:    [[EXT:%.*]] = icmp eq i32 [[TMP1]], 0
 ; CHECK-NEXT:    ret i1 [[EXT]]
 ;
   %cmp = icmp eq <2 x i32> %arg, zeroinitializer
   %ext = extractelement <2 x i1> %cmp, i32 %idx
   ret i1 %ext
 }
 
 define i1 @extractelt_vector_fcmp_constrhs_dynidx(<2 x float> %arg, i32 %idx) {
 ; CHECK-LABEL: @extractelt_vector_fcmp_constrhs_dynidx(
-; CHECK-NEXT:    [[CMP:%.*]] = fcmp oeq <2 x float> [[ARG:%.*]], zeroinitializer
-; CHECK-NEXT:    [[EXT:%.*]] = extractelement <2 x i1> [[CMP]], i32 [[IDX:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x float> [[ARG:%.*]], i32 [[IDX:%.*]]
+; CHECK-NEXT:    [[EXT:%.*]] = fcmp oeq float [[TMP1]], 0.000000e+00
 ; CHECK-NEXT:    ret i1 [[EXT]]
 ;
   %cmp = fcmp oeq <2 x float> %arg, zeroinitializer
   %ext = extractelement <2 x i1> %cmp, i32 %idx
   ret i1 %ext
 }
 
 define i1 @extractelt_vector_fcmp_not_cheap_to_scalarize_multi_use(<2 x float> %arg0, <2 x float> %arg1, <2 x float> %arg2, i32 %idx) {