Missed NonZero optimization opportunity

[mtak-]

Godbolt link

When a NonZeroUsize is created through new_unchecked and wrapped in a Some, the optimizer does not recognize that the value is always Some. This can be fixed by adding a std::intrinsics::assume(value != 0) line before creating the NonZero type.

use std::num::NonZeroUsize;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::Relaxed;

pub static X: AtomicUsize = AtomicUsize::new(1);

pub unsafe fn get() -> Option<NonZeroUsize> {
    let x = X.load(Relaxed);
    Some(NonZeroUsize::new_unchecked(x))
}

pub unsafe fn get2() -> usize {
    match get() {
        Some(x) => x.get(),
        None => unreachable!(), // not optimized out
    }
}

Lots of panicking code gets generated in this example.

Fixed by:

#![feature(core_intrinsics)]

use std::num::NonZeroUsize;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::Relaxed;

pub static X: AtomicUsize = AtomicUsize::new(1);

pub unsafe fn get() -> Option<NonZeroUsize> {
    let x = X.load(Relaxed);
    std::intrinsics::assume(x != 0); // gets rid of the panicking code
    Some(NonZeroUsize::new_unchecked(x))
}

pub unsafe fn get2() -> usize {
    match get() {
        Some(x) => x.get(),
        None => unreachable!(), // optimized away
    }
}

I would guess that adding assume inside of new_unchecked would be ok, since it's already UB if it's not zero.

Related: #51346. Posted as a new issue, because it can be fixed without any LLVM pass changes.

[scottmcm]

Related: #49572

Also, assume can have detrimental effects on downstream codegen, so isn't always a good idea.

[mtak-]

Ah missed #49572. Maybe this is mostly a dupe of existing issues.

[steveklabnik]

Triage: no changes

[bugadani]

I'd like to point out, that this is not caused by NonZero*, but Atomic*:
https://rust.godbolt.org/z/czeWsc