[PATCH] D148440: [InstCombine] Fold zext(shl(trunc)) into and(shl)

[Kazu Hirata via Phabricator via llvm-commits]

kazu created this revision.
Herald added a subscriber: hiraditya.
Herald added a project: All.
kazu requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

If we are truncating to a narrower type, shifting left, and extending
back to the original type, we could do a left shift in the original
type followed by masking.

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

This patch fixes:

https://github.com/llvm/llvm-project/issues/61650


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D148440

Files:
  llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
  llvm/test/Transforms/InstCombine/fsh.ll
  llvm/test/Transforms/InstCombine/trunc-shl-zext.ll


Index: llvm/test/Transforms/InstCombine/trunc-shl-zext.ll
===================================================================
--- llvm/test/Transforms/InstCombine/trunc-shl-zext.ll
+++ llvm/test/Transforms/InstCombine/trunc-shl-zext.ll
@@ -4,9 +4,8 @@
 define i32 @trunc_shl_zext_32(i32 %a) {
 ; CHECK-LABEL: define i32 @trunc_shl_zext_32
 ; CHECK-SAME: (i32 [[A:%.*]]) {
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[A]] to i16
-; CHECK-NEXT:    [[SHL:%.*]] = shl i16 [[TRUNC]], 4
-; CHECK-NEXT:    [[EXT:%.*]] = zext i16 [[SHL]] to i32
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[A]], 4
+; CHECK-NEXT:    [[EXT:%.*]] = and i32 [[TMP1]], 65520
 ; CHECK-NEXT:    ret i32 [[EXT]]
 ;
   %trunc = trunc i32 %a to i16
@@ -18,9 +17,8 @@
 define i64 @trunc_shl_zext_64(i64 %a) {
 ; CHECK-LABEL: define i64 @trunc_shl_zext_64
 ; CHECK-SAME: (i64 [[A:%.*]]) {
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i64 [[A]] to i8
-; CHECK-NEXT:    [[SHL:%.*]] = shl i8 [[TRUNC]], 7
-; CHECK-NEXT:    [[EXT:%.*]] = zext i8 [[SHL]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i64 [[A]], 7
+; CHECK-NEXT:    [[EXT:%.*]] = and i64 [[TMP1]], 128
 ; CHECK-NEXT:    ret i64 [[EXT]]
 ;
   %trunc = trunc i64 %a to i8
Index: llvm/test/Transforms/InstCombine/fsh.ll
===================================================================
--- llvm/test/Transforms/InstCombine/fsh.ll
+++ llvm/test/Transforms/InstCombine/fsh.ll
@@ -672,9 +672,8 @@
 
 define i32 @fshl_mask_args_same2(i32 %a) {
 ; CHECK-LABEL: @fshl_mask_args_same2(
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[A:%.*]] to i16
-; CHECK-NEXT:    [[REV:%.*]] = shl i16 [[TRUNC]], 8
-; CHECK-NEXT:    [[T2:%.*]] = zext i16 [[REV]] to i32
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[A:%.*]], 8
+; CHECK-NEXT:    [[T2:%.*]] = and i32 [[TMP1]], 65280
 ; CHECK-NEXT:    ret i32 [[T2]]
 ;
   %t1 = and i32 %a, 255
Index: llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
===================================================================
--- llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
+++ llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
@@ -1207,6 +1207,25 @@
     return BinaryOperator::CreateAnd(X, ZextC);
   }
 
+  // If we are truncating, shifting left, and then zexting back to the original
+  // type, then shift left in the original type and mask.
+  // zext (shl (trunc X), C1) --> and (shl C1), C2
+  const APInt *C1;
+  if (match(Src, m_Shl(m_Trunc(m_Value(X)), m_APInt(C1))) &&
+      X->getType() == DestTy) {
+    // Shift left in the original type.
+    Value *Shl =
+        Builder.CreateShl(X, ConstantInt::get(DestTy, C1->getZExtValue()));
+
+    // Mask the shift result.
+    unsigned SrcSize = Src->getType()->getScalarSizeInBits();
+    unsigned DstSize = DestTy->getScalarSizeInBits();
+    unsigned MaskWidth = SrcSize - C1->getZExtValue();
+    APInt Mask = APInt::getLowBitsSet(DstSize, MaskWidth);
+    Mask <<= C1->getZExtValue();
+    return BinaryOperator::CreateAnd(Shl, ConstantInt::get(DestTy, Mask));
+  }
+
   if (match(Src, m_VScale())) {
     if (Zext.getFunction() &&
         Zext.getFunction()->hasFnAttribute(Attribute::VScaleRange)) {


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D148440.513938.patch
Type: text/x-patch
Size: 3116 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20230415/240a9753/attachment.bin>