[InstCombine] Teach canEvaluateTruncated to handle arithmetic shift (including those with vector splat shift amount)

[aaboud]

Summary

There are cases where AShr can not be converted to LShr, thus, we need to handle AShr in canEvaluateTruncated.

[spatel]

LGTM.

Just to make sure I'm understanding: the multiply patch (D36679) is independent of this?

[aaboud]

In D36784#843566, @spatel wrote:

LGTM.

Just to make sure I'm understanding: the multiply patch (D36679) is independent of this?

Yes, I managed to find an example where I am using OR instruction and that AShr instruction cannot be replaced by LShr instruction.
So, we can commit D36679 separately after this one.

lib/Transforms/InstCombine/InstCombineCasts.cpp

if (match(I->getOperand(1), m_APInt(Amt))) {
       if (IC.MaskedValueIsZero(I->getOperand(0),
             APInt::getHighBitsSet(OrigBitWidth, OrigBitWidth-BitWidth), 0, CxtI) &&
           Amt->getLimitedValue(BitWidth) < BitWidth) {
         return canEvaluateTruncated(I->getOperand(0), Ty, IC, CxtI);
       }
     }
     break;
   }
+  case Instruction::AShr: {
+    // If this is a truncate of an arithmetic shr, we can truncate it to a
+    // smaller ashr iff we know that all the bits from the sign bit of the
+    // original type and the sign bit of the truncate type are similar.
+    // TODO: It is enough to check that the bits we would be shifting in are
+    //       similar to sign bit of the truncate type.
+    const APInt *Amt;
+    if (match(I->getOperand(1), m_APInt(Amt))) {
+      uint32_t OrigBitWidth = OrigTy->getScalarSizeInBits();
+      uint32_t BitWidth = Ty->getScalarSizeInBits();
+      if (Amt->getLimitedValue(BitWidth) < BitWidth &&
+          OrigBitWidth - BitWidth <
+              IC.ComputeNumSignBits(I->getOperand(0), 0, CxtI))
+        return canEvaluateTruncated(I->getOperand(0), Ty, IC, CxtI);
+    }
+    break;
+  }
   case Instruction::Trunc:
     // trunc(trunc(x)) -> trunc(x)
     return true;
   case Instruction::ZExt:
   case Instruction::SExt:
     // trunc(ext(x)) -> ext(x) if the source type is smaller than the new dest
     // trunc(ext(x)) -> trunc(x) if the source type is larger than the new dest
     return true;

test/Transforms/InstCombine/trunc.ll

 ; CHECK-NEXT:    ret i32 [[D]]
 ;
   %B = zext i64 %A to i128
   %C = lshr i128 %B, 32
   %D = trunc i128 %C to i32
   ret i32 %D
 }
 
+define i32 @trunc_ashr(i32 %X) {
+; CHECK-LABEL: @trunc_ashr(
+; CHECK-NEXT:    [[B:%.*]] = or i32 [[X:%.*]], -2147483648
+; CHECK-NEXT:    [[C:%.*]] = ashr i32 [[B]], 8
+; CHECK-NEXT:    ret i32 [[C]]
+;
+  %A = zext i32 %X to i36
+  %B = or i36 %A, -2147483648 ; 0xF80000000
+  %C = ashr i36 %B, 8
+  %T = trunc i36 %C to i32
+  ret i32  %T
+}
+
+define <2 x i32> @trunc_ashr_vec(<2 x i32> %X) {
+; CHECK-LABEL: @trunc_ashr_vec(
+; CHECK-NEXT:    [[B:%.*]] = or <2 x i32> [[X:%.*]], <i32 -2147483648, i32 -2147483648>
+; CHECK-NEXT:    [[C:%.*]] = ashr <2 x i32> [[B]], <i32 8, i32 8>
+; CHECK-NEXT:    ret <2 x i32> [[C]]
+;
+  %A = zext <2 x i32> %X to <2 x i36>
+  %B = or <2 x i36> %A, <i36 -2147483648, i36 -2147483648> ; 0xF80000000
+  %C = ashr <2 x i36> %B, <i36 8, i36 8>
+  %T = trunc <2 x i36> %C to <2 x i32>
+  ret <2 x i32>  %T
+}
+
 define i92 @test7(i64 %A) {
 ; CHECK-LABEL: @test7(
 ; CHECK-NEXT:    [[TMP1:%.*]] = lshr i64 %A, 32
 ; CHECK-NEXT:    [[D:%.*]] = zext i64 [[TMP1]] to i92
 ; CHECK-NEXT:    ret i92 [[D]]
 ;
   %B = zext i64 %A to i128
   %C = lshr i128 %B, 32