[InstCombine] Fold zext(shl(trunc)) into and(shl)

[kazu]

Summary

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

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

1214

Why only shl, this seems to be pretty generically true:
https://alive2.llvm.org/ce/z/kTBKC7

[kazu]

I missed the entire block of code in InstCombinerImpl::visitZExt to widen an expression tree:

// Try to extend the entire expression tree to the wide destination type.                                            
unsigned BitsToClear;
if (shouldChangeType(SrcTy, DestTy) &&
    canEvaluateZExtd(Src, DestTy, BitsToClear, *this, &Zext)) {

Let me see if I could modify this block code to accommodate our trunc-shl-zext sequence.

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

1214

Also, why just APInt? We should be able to perform this for any Constant (and just create the Mask below generically with Builder).

[kazu]

In D148440#4271660, @kazu wrote:

I missed the entire block of code in InstCombinerImpl::visitZExt to widen an expression tree:

// Try to extend the entire expression tree to the wide destination type.                                            
unsigned BitsToClear;
if (shouldChangeType(SrcTy, DestTy) &&
    canEvaluateZExtd(Src, DestTy, BitsToClear, *this, &Zext)) {

Let me see if I could modify this block code to accommodate our trunc-shl-zext sequence.

On X86 at least, shouldChangeType returns false on going from 16-bit integers to 32-bit integers because DataLayout::isLegalInteger always returns false. Note that LegalIntWidths is empty on X86.

Widening narrow computation to well supported types like i32 or i64 seems to be a worthwhile transformation.

Do we happen to know why LegalIntWidths is empty on X86? Anyone?

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

Instruction *InstCombinerImpl::visitZExt(ZExtInst &Zext) {
   // a non-constant mask operand.
   // zext (and (trunc X), C) --> and X, (zext C)
   if (match(Src, m_And(m_Trunc(m_Value(X)), m_Constant(C))) &&
       X->getType() == DestTy) {
     Constant *ZextC = ConstantExpr::getZExt(C, DestTy);
     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)) {
       Attribute Attr =
           Zext.getFunction()->getFnAttribute(Attribute::VScaleRange);
       if (std::optional<unsigned> MaxVScale = Attr.getVScaleRangeMax()) {
         unsigned TypeWidth = Src->getType()->getScalarSizeInBits();
         if (Log2_32(*MaxVScale) < TypeWidth) {

llvm/test/Transforms/InstCombine/fsh.ll

 ;
   %t1 = and i32 %a, 4294901760 ; 0xffff0000
   %t2 = call i32 @llvm.fshl.i32(i32 %t1, i32 %t1, i32 16)
   ret i32 %t2
 }
 
 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
   %t2 = call i32 @llvm.fshl.i32(i32 %t1, i32 %t1, i32 8)
   ret i32 %t2
 }
 
 define i32 @fshl_mask_args_same3(i32 %a) {

llvm/test/Transforms/InstCombine/trunc-shl-zext.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
 ; RUN: opt < %s -passes=instcombine -S | FileCheck %s
 
 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
   %shl = shl i16 %trunc, 4
   %ext = zext i16 %shl to i32
   ret i32 %ext
 }
 
 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
   %shl = shl i8 %trunc, 7
   %ext = zext i8 %shl to i64
   ret i64 %ext
 }