[InstCombine] Improve vector undef handling for sext(ashr(shl(trunc()))) fold

If the shift amount was undef for some lane, the shift amount in opposite
shift is irrelevant for that lane, and the new shift amount for that lane
can be undef.

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
index 59dae93..6e94e58 100644 (file)
--- a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
@@ -1520,13 +1520,15 @@ Instruction *InstCombinerImpl::visitSExt(SExtInst &CI) {
   Constant *BA = nullptr, *CA = nullptr;
   if (match(Src, m_AShr(m_Shl(m_Trunc(m_Value(A)), m_Constant(BA)),
                         m_Constant(CA))) &&
-      BA == CA && A->getType() == DestTy) {
+      BA->isElementWiseEqual(CA) && A->getType() == DestTy) {
     Constant *WideCurrShAmt = ConstantExpr::getSExt(CA, DestTy);
     Constant *NumLowbitsLeft = ConstantExpr::getSub(
         ConstantInt::get(DestTy, SrcTy->getScalarSizeInBits()), WideCurrShAmt);
     Constant *NewShAmt = ConstantExpr::getSub(
         ConstantInt::get(DestTy, DestTy->getScalarSizeInBits()),
         NumLowbitsLeft);
+    NewShAmt =
+        Constant::mergeUndefsWith(Constant::mergeUndefsWith(NewShAmt, BA), CA);
     A = Builder.CreateShl(A, NewShAmt, CI.getName());
     return BinaryOperator::CreateAShr(A, NewShAmt);
   }

diff --git a/llvm/test/Transforms/InstCombine/sext.ll b/llvm/test/Transforms/InstCombine/sext.ll
index 23639e9..31923fd 100644 (file)
--- a/llvm/test/Transforms/InstCombine/sext.ll
+++ b/llvm/test/Transforms/InstCombine/sext.ll
@@ -166,10 +166,8 @@ define <2 x i32> @test10_vec_nonuniform(<2 x i32> %i) {
 
 define <2 x i32> @test10_vec_undef0(<2 x i32> %i) {
 ; CHECK-LABEL: @test10_vec_undef0(
-; CHECK-NEXT:    [[A:%.*]] = trunc <2 x i32> [[I:%.*]] to <2 x i8>
-; CHECK-NEXT:    [[B:%.*]] = shl <2 x i8> [[A]], <i8 6, i8 0>
-; CHECK-NEXT:    [[C:%.*]] = ashr <2 x i8> [[B]], <i8 6, i8 undef>
-; CHECK-NEXT:    [[D:%.*]] = sext <2 x i8> [[C]] to <2 x i32>
+; CHECK-NEXT:    [[D1:%.*]] = shl <2 x i32> [[I:%.*]], <i32 30, i32 undef>
+; CHECK-NEXT:    [[D:%.*]] = ashr <2 x i32> [[D1]], <i32 30, i32 undef>
 ; CHECK-NEXT:    ret <2 x i32> [[D]]
 ;
   %A = trunc <2 x i32> %i to <2 x i8>
@@ -180,10 +178,8 @@ define <2 x i32> @test10_vec_undef0(<2 x i32> %i) {
 }
 define <2 x i32> @test10_vec_undef1(<2 x i32> %i) {
 ; CHECK-LABEL: @test10_vec_undef1(
-; CHECK-NEXT:    [[A:%.*]] = trunc <2 x i32> [[I:%.*]] to <2 x i8>
-; CHECK-NEXT:    [[B:%.*]] = shl <2 x i8> [[A]], <i8 6, i8 undef>
-; CHECK-NEXT:    [[C:%.*]] = ashr <2 x i8> [[B]], <i8 6, i8 0>
-; CHECK-NEXT:    [[D:%.*]] = sext <2 x i8> [[C]] to <2 x i32>
+; CHECK-NEXT:    [[D1:%.*]] = shl <2 x i32> [[I:%.*]], <i32 30, i32 undef>
+; CHECK-NEXT:    [[D:%.*]] = ashr <2 x i32> [[D1]], <i32 30, i32 undef>
 ; CHECK-NEXT:    ret <2 x i32> [[D]]
 ;
   %A = trunc <2 x i32> %i to <2 x i8>
@@ -194,8 +190,8 @@ define <2 x i32> @test10_vec_undef1(<2 x i32> %i) {
 }
 define <2 x i32> @test10_vec_undef2(<2 x i32> %i) {
 ; CHECK-LABEL: @test10_vec_undef2(
-; CHECK-NEXT:    [[D1:%.*]] = shl <2 x i32> [[I:%.*]], <i32 30, i32 24>
-; CHECK-NEXT:    [[D:%.*]] = ashr <2 x i32> [[D1]], <i32 30, i32 24>
+; CHECK-NEXT:    [[D1:%.*]] = shl <2 x i32> [[I:%.*]], <i32 30, i32 undef>
+; CHECK-NEXT:    [[D:%.*]] = ashr <2 x i32> [[D1]], <i32 30, i32 undef>
 ; CHECK-NEXT:    ret <2 x i32> [[D]]
 ;
   %A = trunc <2 x i32> %i to <2 x i8>