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

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 319c07d..ebab23a 100644 (file)
--- a/llvm/test/Transforms/InstCombine/sext.ll
+++ b/llvm/test/Transforms/InstCombine/sext.ll
@@ -167,10 +167,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>
@@ -181,10 +179,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>
@@ -195,8 +191,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>

diff --git a/llvm/test/Transforms/InstCombine/trunc.ll b/llvm/test/Transforms/InstCombine/trunc.ll
index d9b02ea..fc13d86 100644 (file)
--- a/llvm/test/Transforms/InstCombine/trunc.ll
+++ b/llvm/test/Transforms/InstCombine/trunc.ll
@@ -114,8 +114,8 @@ define <2 x i64> @test2_vec_nonuniform(<2 x i64> %a) {
 define <2 x i64> @test2_vec_undef(<2 x i64> %a) {
 ; CHECK-LABEL: @test2_vec_undef(
 ; CHECK-NEXT:    [[B:%.*]] = trunc <2 x i64> [[A:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[D1:%.*]] = shl <2 x i64> [[A]], <i64 36, i64 0>
-; CHECK-NEXT:    [[D:%.*]] = ashr <2 x i64> [[D1]], <i64 36, i64 0>
+; CHECK-NEXT:    [[D1:%.*]] = shl <2 x i64> [[A]], <i64 36, i64 undef>
+; CHECK-NEXT:    [[D:%.*]] = ashr <2 x i64> [[D1]], <i64 36, i64 undef>
 ; CHECK-NEXT:    call void @use_vec(<2 x i32> [[B]])
 ; CHECK-NEXT:    ret <2 x i64> [[D]]
 ;