[ARM] IselLowering unsigned overflow to crash using APInt in PerformSHLSimplify

[Peter]

Summary

This diff fixes issue https://github.com/llvm/llvm-project/issues/59317

We should check if bitwidth is lower than the shift amount before we subtract them to avoid unsigned overflow.

[Peter]

format

[RKSimon]

llvm/lib/Target/ARM/ARMISelLowering.cpp

13778

I'm not sure if you can encounter types > 64 bits here - but you might be better doing the test as

if (C2Int.uge(C2Width))
  return SDValue();
uint64_t C2Value = C2Int.getZExtValue();

[Peter]

Update and remove getZExtValue()

llvm/lib/Target/ARM/ARMISelLowering.cpp

13778

I 100% agree with you. I meant to change that in another patch as it could be some other issue, but I'll just do it here.
However, I couldn't find a file that can crash it, probably legalizer has already split i128 into smaller types prior to this, so I'll just update the code without new tests.

getZExtValue and getSExtValue in APInt can be particularly problematic, I see people use it without checking if it is allowed and crash the compiler.
I have found another bug with it that actually fits your concern. (https://github.com/llvm/llvm-project/issues/59316)
Probably we should move this discussion to there.

[dmgreen]

The change sounds good to me, LGTM.

llvm/lib/Target/ARM/ARMISelLowering.cpp

static SDValue PerformSHLSimplify(SDNode *N,
 
   auto *C1ShlC2 = dyn_cast<ConstantSDNode>(N->getOperand(1));
   auto *C2 = dyn_cast<ConstantSDNode>(SHL.getOperand(1));
   if (!C1ShlC2 || !C2)
     return SDValue();
 
   APInt C2Int = C2->getAPIntValue();
   APInt C1Int = C1ShlC2->getAPIntValue();
+  unsigned C2Width = C2Int.getBitWidth();
+  if (C2Int.uge(C2Width))
+    return SDValue();
+  uint64_t C2Value = C2Int.getZExtValue();
 
   // Check that performing a lshr will not lose any information.
-  APInt Mask = APInt::getHighBitsSet(C2Int.getBitWidth(),
-                                     C2Int.getBitWidth() - C2->getZExtValue());
+  APInt Mask = APInt::getHighBitsSet(C2Width, C2Width - C2Value);
   if ((C1Int & Mask) != C1Int)
     return SDValue();
 
   // Shift the first constant.
   C1Int.lshrInPlace(C2Int);
 
   // The immediates are encoded as an 8-bit value that can be rotated.
   auto LargeImm = [](const APInt &Imm) {

llvm/test/CodeGen/ARM/pr59317.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=arm %s -o - | FileCheck --check-prefix=arm %s
+; RUN: llc -mtriple=armeb %s -o - | FileCheck --check-prefix=armeb %s
+
+define i1 @pr59317(i16 %F) {
+; arm-LABEL: pr59317:
+; arm:       @ %bb.0: @ %BB
+; arm-NEXT:    sub sp, sp, #8
+; arm-NEXT:    mov r0, #0
+; arm-NEXT:    add sp, sp, #8
+; arm-NEXT:    mov pc, lr
+;
+; armeb-LABEL: pr59317:
+; armeb:       @ %bb.0: @ %BB
+; armeb-NEXT:    sub sp, sp, #8
+; armeb-NEXT:    mov r0, #0
+; armeb-NEXT:    add sp, sp, #8
+; armeb-NEXT:    mov pc, lr
+BB:
+  %E = extractelement <1 x i16> <i16 -1>, i16 %F
+  %RP = alloca i64, align 8
+  %B = shl i16 %E, %E
+  %C1 = icmp ugt i16 %B, %F
+  ret i1 %C1
+}