[llvm] 8eaa1ef - [InstCombine] add casts from splat-a-bit pattern if necessary

[Chenbing Zheng via llvm-commits]

Author: Chenbing Zheng
Date: 2022-05-07T15:34:57+08:00
New Revision: 8eaa1ef0d88c36dbf90e191c8852ea46bf4841c9

URL: https://github.com/llvm/llvm-project/commit/8eaa1ef0d88c36dbf90e191c8852ea46bf4841c9
DIFF: https://github.com/llvm/llvm-project/commit/8eaa1ef0d88c36dbf90e191c8852ea46bf4841c9.diff

LOG: [InstCombine] add casts from splat-a-bit pattern if necessary

Splatting a bit of constant-index across a value:
sext (ashr (trunc iN X to iM), M-1) to iN --> ashr (shl X, N-M), N-1
If the dest type is different, use a cast (adjust use check).

https://alive2.llvm.org/ce/z/acAan3

Reviewed By: spatel

Differential Revision: https://reviews.llvm.org/D124590

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
    llvm/test/Transforms/InstCombine/sext.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
index 9a89aa557258a..4fbbdd5d6dbcd 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
@@ -1592,14 +1592,20 @@ Instruction *InstCombinerImpl::visitSExt(SExtInst &CI) {
 
   // Splatting a bit of constant-index across a value:
   // sext (ashr (trunc iN X to iM), M-1) to iN --> ashr (shl X, N-M), N-1
-  // TODO: If the dest type is 
diff erent, use a cast (adjust use check).
+  // If the dest type is 
diff erent, use a cast (adjust use check).
   if (match(Src, m_OneUse(m_AShr(m_Trunc(m_Value(X)),
-                                 m_SpecificInt(SrcBitSize - 1)))) &&
-      X->getType() == DestTy) {
-    Constant *ShlAmtC = ConstantInt::get(DestTy, DestBitSize - SrcBitSize);
-    Constant *AshrAmtC = ConstantInt::get(DestTy, DestBitSize - 1);
-    Value *Shl = Builder.CreateShl(X, ShlAmtC);
-    return BinaryOperator::CreateAShr(Shl, AshrAmtC);
+                                 m_SpecificInt(SrcBitSize - 1))))) {
+    Type *XTy = X->getType();
+    unsigned XBitSize = XTy->getScalarSizeInBits();
+    Constant *ShlAmtC = ConstantInt::get(XTy, XBitSize - SrcBitSize);
+    Constant *AshrAmtC = ConstantInt::get(XTy, XBitSize - 1);
+    if (XTy == DestTy)
+      return BinaryOperator::CreateAShr(Builder.CreateShl(X, ShlAmtC),
+                                        AshrAmtC);
+    if (cast<BinaryOperator>(Src)->getOperand(0)->hasOneUse()) {
+      Value *Ashr = Builder.CreateAShr(Builder.CreateShl(X, ShlAmtC), AshrAmtC);
+      return CastInst::CreateIntegerCast(Ashr, DestTy, /* isSigned */ true);
+    }
   }
 
   if (match(Src, m_VScale(DL))) {

diff  --git a/llvm/test/Transforms/InstCombine/sext.ll b/llvm/test/Transforms/InstCombine/sext.ll
index 2e5735f1c133b..e6d220d8f9b7f 100644
--- a/llvm/test/Transforms/InstCombine/sext.ll
+++ b/llvm/test/Transforms/InstCombine/sext.ll
@@ -381,13 +381,11 @@ define i32 @smear_set_bit_wrong_shift_amount(i32 %x) {
   ret i32 %s
 }
 
-; TODO: this could be mask+compare+sext or shifts+trunc
-
 define i16 @smear_set_bit_
diff erent_dest_type(i32 %x) {
 ; CHECK-LABEL: @smear_set_bit_
diff erent_dest_type(
-; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
-; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[T]], 7
-; CHECK-NEXT:    [[S:%.*]] = sext i8 [[A]] to i16
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X:%.*]], 24
+; CHECK-NEXT:    [[TMP2:%.*]] = ashr i32 [[TMP1]], 31
+; CHECK-NEXT:    [[S:%.*]] = trunc i32 [[TMP2]] to i16
 ; CHECK-NEXT:    ret i16 [[S]]
 ;
   %t = trunc i32 %x to i8
@@ -415,9 +413,9 @@ define i16 @smear_set_bit_
diff erent_dest_type_extra_use(i32 %x) {
 
 define i64 @smear_set_bit_
diff erent_dest_type_wider_dst(i32 %x) {
 ; CHECK-LABEL: @smear_set_bit_
diff erent_dest_type_wider_dst(
-; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
-; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[T]], 7
-; CHECK-NEXT:    [[S:%.*]] = sext i8 [[A]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X:%.*]], 24
+; CHECK-NEXT:    [[TMP2:%.*]] = ashr i32 [[TMP1]], 31
+; CHECK-NEXT:    [[S:%.*]] = sext i32 [[TMP2]] to i64
 ; CHECK-NEXT:    ret i64 [[S]]
 ;
   %t = trunc i32 %x to i8