[InstCombine] Preserve nsw constraint via @llvm.assume when transforming umin + sub nsw to usub.sat to enable further folding

[ParkHanbum]

Description

I am proposing an optimization to merge two range checks into a single comparison.
Specifically, the pattern (X < C) || (X - C < C2) can be folded into X < C + C2 under certain conditions (positive constants and no signed wrap for the sum).

However, the current InstCombine pipeline transforms the original pattern (umin + sub nsw) into a form (equivalent to usub.sat) that drops the implicit nsw constraint. Consequently, the final optimization becomes invalid because it lacks the necessary range information for %x.

1. The Goal

We want to optimize the range check logic into a single comparison.
Preconditions: C >= 0, C2 >= 0, C + C2 does not overflow (signed), and the predicate is slt.

Original Source (Optimization Valid):
In this original IR, the nsw flag on sub implies that the operation does not overflow. This implicitly restricts the range of %x (e.g., preventing negative values that would cause underflow), making the transformation to @tgt valid.

    
define i1 @src_origin(i8 noundef %x, i8 noundef %C, i8 %C2) {
    %assume = icmp slt i8 0, %C2
    call void @llvm.assume(i1 %assume)
    %C_is_pos = icmp sge i8 %C, 0
    call void @llvm.assume(i1 %C_is_pos)
    %sum = add i8 %C, %C2
    %sum_is_valid = icmp sge i8 %sum, 0
    call void @llvm.assume(i1 %sum_is_valid)

    %v0 = tail call i8 @llvm.umin.i8(i8 %x, i8 %C)
    ; The 'nsw' here is crucial. It makes UB if x causes underflow.
    %v1 = sub nsw i8 %x, %v0
    %v2 = icmp slt i8 %v1, %C2
    ret i1 %v2
}

define i1 @tgt_origin(i8 noundef %x, i8 noundef %C, i8 %C2) {
    ; Same assumptions as above...
    %sum = add i8 %C, %C2
    ; This optimization is valid because 'nsw' in src refined the domain of %x.
    %cmp = icmp slt i8 %x, %sum
    ret i1 %cmp
}

2. The Problem: Information Loss in Current Transformation

Currently, LLVM canonicalizes umin + sub patterns into usub.sat (or an equivalent OR-pattern: (X u< C) || (X - C < C2)).
During this step, the nsw information is discarded because usub.sat is well-defined for all inputs.

Once converted to usub.sat, the final optimization is no longer valid without explicit constraints on %x.

Current IR State (Optimization Invalid):

define i1 @src_current_transformed(i8 noundef %x, i8 noundef %C, i8 %C2) {
    ; Assumptions on C and C2 exist, BUT NO assumption on %x.
    %assume = icmp slt i8 0, %C2
    call void @llvm.assume(i1 %assume)
    %C_is_pos = icmp sge i8 %C, 0
    call void @llvm.assume(i1 %C_is_pos)
    %sum = add i8 %C, %C2
    %sum_is_valid = icmp sge i8 %sum, 0
    call void @llvm.assume(i1 %sum_is_valid)

    ; 'nsw' is gone. usub.sat is defined even if x is negative (large unsigned).
    %sub = call i8 @llvm.usub.sat.i8(i8 %x, i8 %C)
    %cmp = icmp slt i8 %sub, %C2
    ret i1 %cmp
}

define i1 @tgt_desired(i8 noundef %x, i8 noundef %C, i8 %C2) {
    ; ... assumptions ...
    %sum = add i8 %C, %C2
    %cmp = icmp slt i8 %x, %sum
    ret i1 %cmp
}

Reason for failure: If %x is a negative value (e.g., 0x80 / -128), @src_current_transformed evaluates to false (because usub.sat returns a positive result, and slt checks it against C2), whereas @tgt_desired evaluates to true (because -128 is less than Sum).

The original @src_origin allowed this because %x = -128 would result in poison due to sub nsw, making the refinement valid.

3. Proposed Solution

To enable this optimization, we must preserve the information provided by nsw before it is lost during the umin + sub nsw
usub.sat transformation.

I suggest that when InstCombine transforms sub nsw patterns into usub.sat, it should materialize the implicit nsw constraint into an explicit @llvm.assume (e.g., asserting %x >= 0 or relevant range checks).

By keeping the range hint for %x, we can safely fold the logic into the efficient icmp slt %x, %sum form in subsequent steps.

Alive2 Proof : https://alive2.llvm.org/ce/z/BAVye9

associated with : #167055

[ParkHanbum]

@nikic @dtcxzyw Could you please review whether the content I wrote is valid? I'm not sure if I made the proposal correctly.

[dtcxzyw]

IIRC this specific case demonstrates the problem: https://alive2.llvm.org/ce/z/odonBA

However the violation of llvm.assume is immediate UB, which is stronger than poison return value in the src function (aka deferred UB). So we cannot simply use llvm.assume to salvage information from poison-generating flags. Fortunately as we create a new intrinsic call to usub.sat, we can preserve the range information in the range attribute of the first parameter (maybe it also applies to the retval's attributes):

%v0 = tail call i64 @llvm.umin.i64(i64 %arg0, i64 13)
 %v1 = sub nsw i64 %arg0, %v0
->
%v1 = tail call i64 @llvm.usub.sat.i64(i64 range(i64 INT_MIN-13, INT_MIN) %arg0, i64 13)

Then we can make the transformation legal: https://alive2.llvm.org/ce/z/rbTvmb (Alive2 doesn't support range attributes on intrinsic calls so I encode the range attribute on the caller's declaration).

Although this way suffers the risk of miscompilation (See #115988), it has been adopted to allow better codegen of isPow2 idiom (See #100899). In the future fshl/fshr may also have the range attribute on shamts: #118200.

So I think it is fair enough to use the range attribute. Generalized proof: https://alive2.llvm.org/ce/z/B-wt7l