Is it possible to tell compiler to treat an atomic value as immutable?

Imagine code like this:

pub static FOO: AtomicUsize = AtomicUsize::new(0);

pub fn foo(buf: &[u8]) {
    for val in buf {
        if FOO.load(Relaxed) == 1 {
            foo1(val);
        } else {
            foo2(val);
        }
    }
}

Compiler as expected will generate code which checks value behind FOO on each iteration. But what if we can guarantee that value behind FOO does not change? Can we somehow convince compiler to treat FOO as &usize so it will generate code like here (note that compiler checks value behind foo only once)? Simply casting &FOO to &usize has no effect. Using static mut FOO: usize instead of the atomic has no effect as well.

To give a practical example, in the cpufeatures crate under the hood we have atomic variables which get initialized only once per program execution (well, strictly speaking they may get initialized several times, but always to the same value) and we use ZST tokens which prove that we have indeed checked that initialization has been performed. This token allows compiler to remove branch responsible for initialization, but it also proves that value behind the underlying atomic will not change. Unfortunately, I couldn't find a way to convince compiler to trust me about this.

Note that in real life code the loop and the FOO check reside in separate crates, so I can't simply move the check outside of the loop manually.

This would have been my best guess. It is generally guaranteed that a value does not change between any two uses of the same immutable reference.

Is there a reason InitToken has to be a ZST? It would be simpler if it just contained the u8 directly.

1 Like

Alright, here's an interesting look:
[Playground - Do LLVM IR output on release]

There is something interesting to note here:

pub fn foo_test1(buf: &[u8], discriminant: &usize) {
    for val in buf {
        if *discriminant == 1 {
            foo1(val);
        } else {
            foo2(val);
        }
    }
}

This is optimized as intended, with two loops where the check is done beforehand.

pub fn foo_test3(buf: &[u8]) {
    let ptr = unsafe {
        &*(&FOO as *const AtomicUsize as *const usize)
    };
    foo_test1(buf, ptr);
}

This does get optimized too, we observe two loops, with the check on FOO done beforehand.
However

pub fn foo_test2(buf: &[u8]) {
    let ptr = unsafe {
        &*(&FOO as *const AtomicUsize as *const usize)
    };
    for val in buf {
        if *ptr == 1 {
            foo1(val);
        } else {
            foo2(val);
        }
    }
}

Inlining the foo_test1 causes LLVM to still view ptr as an atomic variable, and therefore check on each iteration.

3 Likes

IIRC we hit a similar thing in glib and investigated using the GCC attribute((pure)) which is almost what's desired here. I think the problem is it also declares the function won't mutate global state.

What we really want here is a slight weakening like __hoistable__ or something that just says it's OK to hoist multiple calls out of a loop.

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.