[ConstraintElim] Decompose shl nsw for signed predicates (#76961)

shl nsw x, shift can be interpreted as mul nsw x, (1<<shift), except
when shift is bw-1 (https://alive2.llvm.org/ce/z/vDh2xT). Use this when
decomposing shl. The equivalent decomposition for the unsigned case
already exists.

llvm/lib/Transforms/Scalar/ConstraintElimination.cpp

@@ -517,6 +517,18 @@ static Decomposition decompose(Value *V,
       return Result;
     }
 
+    // (shl nsw x, shift) is (mul nsw x, (1<<shift)), with the exception of
+    // shift == bw-1.
+    if (match(V, m_NSWShl(m_Value(Op0), m_ConstantInt(CI)))) {
+      uint64_t Shift = CI->getValue().getLimitedValue();
+      if (Shift < Ty->getIntegerBitWidth() - 1) {
+        assert(Shift < 64 && "Would overflow");
+        auto Result = decompose(Op0, Preconditions, IsSigned, DL);
+        Result.mul(int64_t(1) << Shift);
+        return Result;
+      }
+    }
+
     return V;
   }
 

llvm/test/Transforms/ConstraintElimination/shl.ll

@@ -1283,8 +1283,7 @@ define i1 @shl_nsw_x8_slt_x7(i8 %start, i8 %high) {
 ; CHECK-NEXT:    [[C_1:%.*]] = icmp slt i8 [[START_SHL_3]], [[HIGH]]
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[C_1]])
 ; CHECK-NEXT:    [[START_MUL_7:%.*]] = mul nsw i8 [[START]], 7
-; CHECK-NEXT:    [[T_1:%.*]] = icmp slt i8 [[START_MUL_7]], [[HIGH]]
-; CHECK-NEXT:    ret i1 [[T_1]]
+; CHECK-NEXT:    ret i1 true
 ;
   %c.0 = icmp sge i8 %high, 0
   call void @llvm.assume(i1 %c.0)
@@ -1327,11 +1326,9 @@ define i1 @shl_nsw_sign_implication(i8 %x) {
 ; CHECK-NEXT:    [[CMP1:%.*]] = icmp slt i8 [[X]], 0
 ; CHECK-NEXT:    br i1 [[CMP1]], label [[IF:%.*]], label [[ELSE:%.*]]
 ; CHECK:       if:
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp slt i8 [[SHL]], 0
-; CHECK-NEXT:    ret i1 [[CMP2]]
+; CHECK-NEXT:    ret i1 true
 ; CHECK:       else:
-; CHECK-NEXT:    [[CMP3:%.*]] = icmp sge i8 [[SHL]], 0
-; CHECK-NEXT:    ret i1 [[CMP3]]
+; CHECK-NEXT:    ret i1 true
 ;
   %shl = shl nsw i8 %x, 2
   %cmp1 = icmp slt i8 %x, 0
@@ -1361,3 +1358,37 @@ define i1 @shl_nsw_by_bw_minus_1(i64 %x) {
   %t.1 = icmp slt i64 %x, 0
   ret i1 %t.1
 }
+
+; Shift returns poison in this case, just make sure we don't crash.
+define i1 @shl_nsw_by_bw(i64 %x) {
+; CHECK-LABEL: @shl_nsw_by_bw(
+; CHECK-NEXT:    [[X_SHL:%.*]] = shl nsw i64 [[X:%.*]], 64
+; CHECK-NEXT:    [[C_1:%.*]] = icmp slt i64 [[X_SHL]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[C_1]])
+; CHECK-NEXT:    [[T_1:%.*]] = icmp slt i64 [[X]], 0
+; CHECK-NEXT:    ret i1 [[T_1]]
+;
+  %x.shl = shl nsw i64 %x, 64
+  %c.1 = icmp slt i64 %x.shl, 0
+  call void @llvm.assume(i1 %c.1)
+
+  %t.1 = icmp slt i64 %x, 0
+  ret i1 %t.1
+}
+
+; Shift returns poison in this case, just make sure we don't crash.
+define i1 @shl_nsw_by_bw_plus_1(i64 %x) {
+; CHECK-LABEL: @shl_nsw_by_bw_plus_1(
+; CHECK-NEXT:    [[X_SHL:%.*]] = shl nsw i64 [[X:%.*]], 65
+; CHECK-NEXT:    [[C_1:%.*]] = icmp slt i64 [[X_SHL]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[C_1]])
+; CHECK-NEXT:    [[T_1:%.*]] = icmp slt i64 [[X]], 0
+; CHECK-NEXT:    ret i1 [[T_1]]
+;
+  %x.shl = shl nsw i64 %x, 65
+  %c.1 = icmp slt i64 %x.shl, 0
+  call void @llvm.assume(i1 %c.1)
+
+  %t.1 = icmp slt i64 %x, 0
+  ret i1 %t.1
+}