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

[Chenbing.Zheng]

Summary

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

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

1595

The TODO says we need to adjust (add) a use check, but this patch did not do that.

Please add a test like this to check that we have the correct behavior:

declare void @use(i8)
define i16 @smear_set_bit_different_dest_type_extra_use(i32 %x) {
  %t = trunc i32 %x to i8
  call void @use(i8 %t)
  %a = ashr i8 %t, 7
  %s = sext i8 %a to i16
  ret i16 %s
}

1612–1614

These 3 values are common in both cases - I think it's better to not repeat the code on both sides of the if/else.

1616

We need another test to verify what happens when DestTy is wider than XTy.

[Chenbing.Zheng]

llvm/test/Transforms/InstCombine/sext.ll

397

For this case, we get a correct result, but it is longer. Should I limit trunc one-use to break it?

[spatel]

llvm/test/Transforms/InstCombine/sext.ll

397

Yes - we do not want to create more instructions than we started with. That is the purpose for the code comment ("adjust use check)".

[Chenbing.Zheng]

adjust use check

[RKSimon]

Please can you pre-commit the new tests and rebase the patch to show the codegen diff

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

1611–1612

There's a potential subtle difference with constant expressions, but you could reduce this to:

if (cast<BinaryOperator>(Src)->getOperand(0)->hasOneUse()) {

The m_Ashr guarantees that we matched a binary operator, so it's ok to plain "cast" it to that type. You could also capture the trunc in the first match with something like this:

Instruction *Trunc;
if (match(Src, m_OneUse(m_AShr(
                   m_CombineAnd(m_Trunc(m_Value(X)), m_Instruction(Trunc)),
                   m_SpecificInt(SrcBitSize - 1))))) {

Then the use check is the straightforward:

if (Trunc->hasOneUse()) {

llvm/test/Transforms/InstCombine/sext.ll

413

Please pre-commit the baseline tests and update here, so we just show the diffs.

[Chenbing.Zheng]

address comments

[Chenbing.Zheng]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

1611–1612

Done and thank you for your detailed explanation.

[spatel]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

if (match(Src, m_AShr(m_Shl(m_Trunc(m_Value(A)), m_Constant(BA)),
     NewShAmt =
         Constant::mergeUndefsWith(Constant::mergeUndefsWith(NewShAmt, BA), CA);
     A = Builder.CreateShl(A, NewShAmt, CI.getName());
     return BinaryOperator::CreateAShr(A, NewShAmt);
   }
 
   // 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 different, use a cast (adjust use check).
+  // If the dest type is different, 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);
+                                 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);
     Value *Shl = Builder.CreateShl(X, ShlAmtC);
+    if (XTy == DestTy) {
       return BinaryOperator::CreateAShr(Shl, AshrAmtC);
+    } else {
+      Value *Ashr = Builder.CreateAShr(Shl, AshrAmtC);
+      return CastInst::CreateIntegerCast(Ashr, DestTy, /* isSigned */ true);
+    }
   }
 
   if (match(Src, m_VScale(DL))) {
     if (CI.getFunction() &&
         CI.getFunction()->hasFnAttribute(Attribute::VScaleRange)) {
       Attribute Attr = CI.getFunction()->getFnAttribute(Attribute::VScaleRange);
       if (Optional<unsigned> MaxVScale = Attr.getVScaleRangeMax()) {
         if (Log2_32(MaxVScale.getValue()) < (SrcBitSize - 1)) {
           Value *VScale = Builder.CreateVScale(ConstantInt::get(DestTy, 1));
           return replaceInstUsesWith(CI, VScale);
         }
       }
     }
   }
 
   return nullptr;

llvm/test/Transforms/InstCombine/sext.ll

 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %t = trunc i32 %x to i8
   %a = ashr i8 %t, 6
   %s = sext i8 %a to i32
   ret i32 %s
 }
 
-; TODO: this could be mask+compare+sext or shifts+trunc
-
 define i16 @smear_set_bit_different_dest_type(i32 %x) {
 ; CHECK-LABEL: @smear_set_bit_different_dest_type(
+; 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
+  %a = ashr i8 %t, 7
+  %s = sext i8 %a to i16
+  ret i16 %s
+}
+
+define i16 @smear_set_bit_different_dest_type_extra_use(i32 %x) {
+; CHECK-LABEL: @smear_set_bit_different_dest_type_extra_use(
 ; 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:    call void @use(i8 [[T]])
+; 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
+  call void @use(i8 %t)
   %a = ashr i8 %t, 7
   %s = sext i8 %a to i16
   ret i16 %s
 }
+
+define i64 @smear_set_bit_different_dest_type_wider_dst(i32 %x) {
+; CHECK-LABEL: @smear_set_bit_different_dest_type_wider_dst(
+; 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
+  %a = ashr i8 %t, 7
+  %s = sext i8 %a to i64
+  ret i64 %s
+}