[InstCombine] matchFunnelShift - support non-uniform constant vector shift amounts (PR46895)

Complete basic PR46895 fixes by refactoring D87452/D88402 to allow us to match non-uniform constant values.

We still don't handle non-uniform vectors that contain undef elements, but that can wait until we have a decent generic mechanism for this.

Differential Revision: https://reviews.llvm.org/D88420

Author: RKSimon

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

@@ -2079,11 +2079,20 @@ static Instruction *matchFunnelShift(Instruction &Or) {
   // matches a subtraction on the R operand.
   auto matchShiftAmount = [&](Value *L, Value *R, unsigned Width) -> Value * {
     // Check for constant shift amounts that sum to the bitwidth.
-    // TODO: Support non-uniform shift amounts.
-    const APInt *LC, *RC;
-    if (match(L, m_APIntAllowUndef(LC)) && match(R, m_APIntAllowUndef(RC)))
-      if (LC->ult(Width) && RC->ult(Width) && (*LC + *RC) == Width)
-        return ConstantInt::get(L->getType(), *LC);
+    const APInt *LI, *RI;
+    if (match(L, m_APIntAllowUndef(LI)) && match(R, m_APIntAllowUndef(RI)))
+      if (LI->ult(Width) && RI->ult(Width) && (*LI + *RI) == Width)
+        return ConstantInt::get(L->getType(), *LI);
+
+    // TODO: Support undefs in non-uniform shift amounts.
+    Constant *LC, *RC;
+    if (match(L, m_Constant(LC)) && !LC->containsUndefElement() &&
+        match(R, m_Constant(RC)) && !RC->containsUndefElement() &&
+        match(L, m_SpecificInt_ICMP(ICmpInst::ICMP_ULT, APInt(Width, Width))) &&
+        match(R, m_SpecificInt_ICMP(ICmpInst::ICMP_ULT, APInt(Width, Width)))) {
+      if (match(ConstantExpr::getAdd(LC, RC), m_SpecificInt(Width)))
+        return L;
+    }
 
     // For non-constant cases, the following patterns currently only work for
     // rotation patterns.

llvm/test/Transforms/InstCombine/funnel.ll

@@ -100,13 +100,12 @@ define <2 x i17> @fshr_v2i17_constant_splat_undef1(<2 x i17> %x, <2 x i17> %y) {
   ret <2 x i17> %r
 }
 
-; TODO: Allow arbitrary shift constants.
+; Allow arbitrary shift constants.
+; TODO: Support undef elements.
 
 define <2 x i32> @fshr_v2i32_constant_nonsplat(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @fshr_v2i32_constant_nonsplat(
-; CHECK-NEXT:    [[SHR:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 17, i32 19>
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i32> [[Y:%.*]], <i32 15, i32 13>
-; CHECK-NEXT:    [[R:%.*]] = or <2 x i32> [[SHL]], [[SHR]]
+; CHECK-NEXT:    [[R:%.*]] = call <2 x i32> @llvm.fshl.v2i32(<2 x i32> [[Y:%.*]], <2 x i32> [[X:%.*]], <2 x i32> <i32 15, i32 13>)
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %shr = lshr <2 x i32> %x, <i32 17, i32 19>
@@ -143,9 +142,7 @@ define <2 x i32> @fshr_v2i32_constant_nonsplat_undef1(<2 x i32> %x, <2 x i32> %y
 
 define <2 x i36> @fshl_v2i36_constant_nonsplat(<2 x i36> %x, <2 x i36> %y) {
 ; CHECK-LABEL: @fshl_v2i36_constant_nonsplat(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i36> [[X:%.*]], <i36 21, i36 11>
-; CHECK-NEXT:    [[SHR:%.*]] = lshr <2 x i36> [[Y:%.*]], <i36 15, i36 25>
-; CHECK-NEXT:    [[R:%.*]] = or <2 x i36> [[SHL]], [[SHR]]
+; CHECK-NEXT:    [[R:%.*]] = call <2 x i36> @llvm.fshl.v2i36(<2 x i36> [[X:%.*]], <2 x i36> [[Y:%.*]], <2 x i36> <i36 21, i36 11>)
 ; CHECK-NEXT:    ret <2 x i36> [[R]]
 ;
   %shl = shl <2 x i36> %x, <i36 21, i36 11>

llvm/test/Transforms/InstCombine/rotate.ll

@@ -122,13 +122,12 @@ define <2 x i17> @rotr_v2i17_constant_splat_undef1(<2 x i17> %x) {
   ret <2 x i17> %r
 }
 
-; TODO: Allow arbitrary shift constants.
+; Allow arbitrary shift constants.
+; TODO: Support undef elements.
 
 define <2 x i32> @rotr_v2i32_constant_nonsplat(<2 x i32> %x) {
 ; CHECK-LABEL: @rotr_v2i32_constant_nonsplat(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i32> [[X:%.*]], <i32 17, i32 19>
-; CHECK-NEXT:    [[SHR:%.*]] = lshr <2 x i32> [[X]], <i32 15, i32 13>
-; CHECK-NEXT:    [[R:%.*]] = or <2 x i32> [[SHL]], [[SHR]]
+; CHECK-NEXT:    [[R:%.*]] = call <2 x i32> @llvm.fshl.v2i32(<2 x i32> [[X:%.*]], <2 x i32> [[X]], <2 x i32> <i32 17, i32 19>)
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %shl = shl <2 x i32> %x, <i32 17, i32 19>
@@ -165,9 +164,7 @@ define <2 x i32> @rotr_v2i32_constant_nonsplat_undef1(<2 x i32> %x) {
 
 define <2 x i36> @rotl_v2i36_constant_nonsplat(<2 x i36> %x) {
 ; CHECK-LABEL: @rotl_v2i36_constant_nonsplat(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i36> [[X:%.*]], <i36 21, i36 11>
-; CHECK-NEXT:    [[SHR:%.*]] = lshr <2 x i36> [[X]], <i36 15, i36 25>
-; CHECK-NEXT:    [[R:%.*]] = or <2 x i36> [[SHL]], [[SHR]]
+; CHECK-NEXT:    [[R:%.*]] = call <2 x i36> @llvm.fshl.v2i36(<2 x i36> [[X:%.*]], <2 x i36> [[X]], <2 x i36> <i36 21, i36 11>)
 ; CHECK-NEXT:    ret <2 x i36> [[R]]
 ;
   %shl = shl <2 x i36> %x, <i36 21, i36 11>