diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
index 59dae932ae495..6e94e58234337 100644
--- 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 319c07d9f8a96..ebab23a195254 100644
--- 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 d9b02eaa16981..fc13d86bcfc39 100644
--- 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]]
 ;