[CodeGen] [ExpandReduction] Fix the bug for ExpandReduction() when vector size isn't power of 2

[ZhangKang]

Summary

For below test case, we will get assert error except for AArch64 and ARM:

declare i8 @llvm.experimental.vector.reduce.and.i8.v3i8(<3 x i8> %a)
define i8 @test_v3i8(<3 x i8> %a) nounwind {
  %b = call i8 @llvm.experimental.vector.reduce.and.i8.v3i8(<3 x i8> %a)
  ret i8 %b
}

In the function getShuffleReduction (), we can see it needs the vector size must be power of 2.

This patch is fix below error when the number of element is not power of 2 for those llvm.experimental.vector.reduce.* function.

[jsji]

LGTM.

llvm/lib/CodeGen/ExpandReductions.cpp

for (auto *II : Worklist) {
     case Intrinsic::experimental_vector_reduce_v2_fmul: {
       // FMFs must be attached to the call, otherwise it's an ordered reduction
       // and it can't be handled by generating a shuffle sequence.
       Value *Acc = II->getArgOperand(0);
       Value *Vec = II->getArgOperand(1);
       if (!FMF.allowReassoc())
         Rdx = getOrderedReduction(Builder, Acc, Vec, getOpcode(ID), MRK);
       else {
+        if (!isPowerOf2_32(Vec->getType()->getVectorNumElements()))
+          continue;
+
         Rdx = getShuffleReduction(Builder, Vec, getOpcode(ID), MRK);
         Rdx = Builder.CreateBinOp((Instruction::BinaryOps)getOpcode(ID),
                                   Acc, Rdx, "bin.rdx");
       }
     } break;
     case Intrinsic::experimental_vector_reduce_add:
     case Intrinsic::experimental_vector_reduce_mul:
     case Intrinsic::experimental_vector_reduce_and:
     case Intrinsic::experimental_vector_reduce_or:
     case Intrinsic::experimental_vector_reduce_xor:
     case Intrinsic::experimental_vector_reduce_smax:
     case Intrinsic::experimental_vector_reduce_smin:
     case Intrinsic::experimental_vector_reduce_umax:
     case Intrinsic::experimental_vector_reduce_umin:
     case Intrinsic::experimental_vector_reduce_fmax:
     case Intrinsic::experimental_vector_reduce_fmin: {
       Value *Vec = II->getArgOperand(0);
+      if (!isPowerOf2_32(Vec->getType()->getVectorNumElements()))
+        continue;
+
       Rdx = getShuffleReduction(Builder, Vec, getOpcode(ID), MRK);
     } break;
     default:
       continue;
     }
     II->replaceAllUsesWith(Rdx);
     II->eraseFromParent();
     Changed = true;

llvm/test/CodeGen/Generic/expand-experimental-reductions.ll

 declare i64 @llvm.experimental.vector.reduce.smax.v2i64(<2 x i64>)
 declare i64 @llvm.experimental.vector.reduce.smin.v2i64(<2 x i64>)
 declare i64 @llvm.experimental.vector.reduce.umax.v2i64(<2 x i64>)
 declare i64 @llvm.experimental.vector.reduce.umin.v2i64(<2 x i64>)
 
 declare double @llvm.experimental.vector.reduce.fmax.v2f64(<2 x double>)
 declare double @llvm.experimental.vector.reduce.fmin.v2f64(<2 x double>)
 
+declare i8 @llvm.experimental.vector.reduce.and.i8.v3i8(<3 x i8>)
 
 define i64 @add_i64(<2 x i64> %vec) {
 ; CHECK-LABEL: @add_i64(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[RDX_SHUF:%.*]] = shufflevector <2 x i64> [[VEC:%.*]], <2 x i64> undef, <2 x i32> <i32 1, i32 undef>
 ; CHECK-NEXT:    [[BIN_RDX:%.*]] = add <2 x i64> [[VEC]], [[RDX_SHUF]]
 ; CHECK-NEXT:    [[TMP0:%.*]] = extractelement <2 x i64> [[BIN_RDX]], i32 0
 ; CHECK-NEXT:    ret i64 [[TMP0]]

 ; CHECK-NEXT:    [[RDX_MINMAX_SELECT:%.*]] = select fast <2 x i1> [[RDX_MINMAX_CMP]], <2 x double> [[VEC]], <2 x double> [[RDX_SHUF]]
 ; CHECK-NEXT:    [[TMP0:%.*]] = extractelement <2 x double> [[RDX_MINMAX_SELECT]], i32 0
 ; CHECK-NEXT:    ret double [[TMP0]]
 ;
 entry:
   %r = call double @llvm.experimental.vector.reduce.fmin.v2f64(<2 x double> %vec)
   ret double %r
 }
+
+; Test when the vector size is not power of two.
+define i8 @test_v3i8(<3 x i8> %a) nounwind {
+; CHECK-LABEL: @test_v3i8(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    %b = call i8 @llvm.experimental.vector.reduce.and.v3i8(<3 x i8> %a)
+; CHECK-NEXT:    ret i8 %b
+;
+entry:
+  %b = call i8 @llvm.experimental.vector.reduce.and.i8.v3i8(<3 x i8> %a)
+  ret i8 %b
+}