[llvm] [InstCombine] Add handling for (or (zext x), (shl (zext (ashr x, bw-1))), bw) -> (sext x) fold (PR #130316)

[Simon Pilgrim via llvm-commits]

https://github.com/RKSimon created https://github.com/llvm/llvm-project/pull/130316

Minor tweak to #129363 which handled all the cases where there was a sext for the original source value, but not for cases where the source is already half the size of the destination type

Another regression noticed in #76524

>From 635e8bd961b1e49c4c24d68f6bfb97417dcd18d2 Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev at redking.me.uk>
Date: Fri, 7 Mar 2025 17:54:08 +0000
Subject: [PATCH] [InstCombine] Add handling for (or (zext x), (shl (zext (ashr
 x, bw-1))), bw) -> (sext x) fold

Minor tweak #129363 which handled all the cases where there was a sext for the original source value, but not for cases where the source is already half the size of the destination type

Another regression noticed in #76524
---
 llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp | 7 +++----
 llvm/test/Transforms/InstCombine/iX-ext-split.ll        | 6 +-----
 2 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
index 330358ee749e5..f81e2de7b365c 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
@@ -3122,13 +3122,12 @@ static Value *matchOrConcat(Instruction &Or, InstCombiner::BuilderTy &Builder) {
   // iX ext split: extending or(zext(x),shl(zext(y),bw/2) pattern
   // to consume sext/ashr: or(zext(sext(x)),shl(zext(sext(ashr(x))),bw/2)
   Value *X;
-  if (match(LowerSrc, m_SExt(m_Value(X))) &&
+  if (match(LowerSrc, m_SExtOrSelf(m_Value(X))) &&
       match(UpperSrc,
-            m_SExt(m_AShr(
+            m_SExtOrSelf(m_AShr(
                 m_Specific(X),
-                m_SpecificInt(X->getType()->getScalarSizeInBits() - 1))))) {
+                m_SpecificInt(X->getType()->getScalarSizeInBits() - 1)))))
     return Builder.CreateSExt(X, Ty);
-  }
 
   return nullptr;
 }
diff --git a/llvm/test/Transforms/InstCombine/iX-ext-split.ll b/llvm/test/Transforms/InstCombine/iX-ext-split.ll
index 5f8276934c2ef..fc804df0e4bec 100644
--- a/llvm/test/Transforms/InstCombine/iX-ext-split.ll
+++ b/llvm/test/Transforms/InstCombine/iX-ext-split.ll
@@ -94,11 +94,7 @@ define void @i128_ext_split_store_i64(i64 %x, ptr %out)  {
 ; CHECK-LABEL: define void @i128_ext_split_store_i64(
 ; CHECK-SAME: i64 [[X:%.*]], ptr [[OUT:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[LO:%.*]] = zext i64 [[X]] to i128
-; CHECK-NEXT:    [[SIGN:%.*]] = ashr i64 [[X]], 63
-; CHECK-NEXT:    [[WIDEN:%.*]] = zext i64 [[SIGN]] to i128
-; CHECK-NEXT:    [[HI:%.*]] = shl nuw i128 [[WIDEN]], 64
-; CHECK-NEXT:    [[RES:%.*]] = or disjoint i128 [[HI]], [[LO]]
+; CHECK-NEXT:    [[RES:%.*]] = sext i64 [[X]] to i128
 ; CHECK-NEXT:    store i128 [[RES]], ptr [[OUT]], align 16
 ; CHECK-NEXT:    ret void
 ;