Const-first functions

language design
1

Today, as far as I'm aware, Rust has two variants of const functions:

What I often want though is something I would call const-first function. It is a function very similar to const fn, but such that is always evaluated at compile time, assuming all arguments are known at compile time. It could have been the default, but I am aware of open quesstions in Tracking Issue for const_eval_select · Issue #124625 · rust-lang/rust · GitHub, and it seems unlikely because of that. And even if it was, there would still be a need for some kind of opt-out for size-constrained environments.

I have two related examples for where this is valuable, both involving heavy use of const generics.

Miri

I have implemented an algorithm that uses a few global constants and a few const generics. Once parametrized the code compiles to an efficient and mostly branchless code by LLVM.

Here is an example of such a function:

Every branch there is eliminated at compile time, every shit is known at compile time too. The issue is that the resolution of these values seems to happen too late.

This is likely why running tests that call this function under Miri take an unacceptably long time in CI. It does finish eventually, but take way too long to be usable.

My current understanding is that this is happening due to Miri needing to evaluate a lot of expressions, function calls and conditions in runtime despite almost everything here known at compile time. A lot of runtime compute could have been saved if as much as possible was optimized early.

rust-gpu

I have also implemented the same algorithm on the GPU too using rust-gpu, here is an equivalent function there:

While not exactly the same, I did hit an issue of a similar nature with it, where, initially, compiled shader required Int8 capability (ability to use 8-bit integers, which is not the default capability for all GPUs). The reason for that was again computation that was fully known at compile time technically, but not evaluated at compile time by Rustc due to lack of const-first functions.

I do not know enough about rust-gpu internals, but it is clear that whatever stage it takes the input from was not optimized basically at all (I have more experience with it pointing to this being the case). My workaround was to use const {} in some places and some ugly unrolling of expressions in others.

The annoying thing is that while I can use const {} once, after assigning it to a variable I can't use variable's output in another const {} block. I believe generic_const_args might have a bit nicer ergonomics in this regard than generic_const_exprs linked examples uses, but it is still a major hit to ergonomics forcing me to write unnatural code.

Ideal solution

Ideally, I'd like to be able to have a const function that always evaluates at compile time as long as all inputs are known at compile time. Moreover, if an expression is composed only from inputs that are known at compile time, compute it recursively too. I don't think this requires any sophisticated analysis and should be easily parallelizable.

For example, foo2 would be reduced to a constant by Rustc and involve zero compute for Miri or Int8 capability for rust-gpu in the following example:

const A: u8 = 5;
const B: u8 = 3;

let foo1 = u32::from(5) << B;
let foo2 = foo1 * 16;

...

I believe most of the const functions in the standard library and otherwise will be const-first. Only rare cases when the function derives some large amount of data from small amount of data (increasing binary size) or the computed function itself needs heavy optimizations (impact compile time too much otherwise), it would be preferred to not compute it at compile time.

Today this seems to be done on LLVM level (I might be totally wrong about this though), but I think it would be beneficial, at least for both examples above, to do that much, much earlier.

2

If this can be const and only has const inputs, most likely LLVM will be able to constant-fold it. So if the only motivation is performance, it is unlikely to be strong enough.

3

Yes, it is already done by LLVM, my motivation is not that. I know that the post is a bit long and takes a while to read, but even LLVM is already mentioned there twice.

4

I've read that, and it's not convincing. Optimizing for Miri execution is not something we should do, especially when you give up (some) detection of UB by that. If that works for you, you can enable MIR optimizations. It's also not clear to me why rust-gpu doesn't optimize that, but that is pretty clearly a bug in rust-gpu and does not justify a language feature.

5

I'm not aware of a way to enable MIR optimizations in a way that benefits Miri performance. And it is not like I'm asking something that can't be done today or something LLVM doesn't do anyway, it just has very bad ergonomics today.

I wish I had knowledge of Rust internals to add optimization pass before Miri and see the impact of such a change on interpreter performance.

I'm also wondering if this has the potential of increasing overall compiler performance for const-heavy code. LLVM does optimize it, but Rustc has more context and can likely do it much earlier and produce less output for later compiler stages in an obvious way.

6

Just enabling MIR optimizations can benefit Miri performance, i.e. RUSTFLAGS="-Zmir-opt-level=1,2,3". This even used to be documented in Miri's README, not sure why it was removed. And it can make Miri miss UB, as I said.

7

What does this mean? Can you describe, in terms of what the program actually does (i.e. the observable behavior of the program in terms of the I/O it performs), how such a function differs from a fn?

Or are you describing something that's more of an optimization hint, like #[inline]? It would be good to be explicit about this as "new language feature for different program semantics" vs "new optimization hint" are very different kinds of discussions.

Why would knowing which branch to take speed up Miri? Dead branches cost basically no time.

I think we have an XY problem where you have an incorrect idea of how Miri works and then are proposing new features based on that.

If there are values that you can pre-compute, you can put them in const blocks. They will then only be computed once even if you call the function many times, also in Miri.

const fn compute_thing(a: i32) -> i32 {
    a * 42
}

pub fn my_fun<const N: i32>(x: i32) -> i32 {
    let thing = const { compute_thing(N) };
    thing * x
}

It usually doesn't help performance very much. And as you said, it makes Miri miss UB, so it kind of defeats the purpose of the tool.

1 Like
8

The net effect is very similar to wrapping it const {} without the need to slap const {} everywhere. So I guess it is kind of like optimization hint, but I'd rather for it to be more of a requirement.

In the example provided there would be no branches and most function calls will be eliminated too. Just a few bit operations with offsets known at compile time and blake3 hash. With this being a very hot function, I'm extrapolating that this would have non-negligible performance impact, but I don't have hard numbers.

I can, but:

  1. This if is verbose, for functions like u32::from() that should, arguably, be the default behavior without explicit const {}
  2. It doesn't really help with the examples provided since the result is first stored in a regular variable and then mixed with other variables. However, those variables are all derived from constant expressions. For example:
// const fn with constants known at compile time:
let parent_metadata_bits = metadata_size_bits(K, PARENT_TABLE_NUMBER);
// Combined with another const fn, with constants also known at compile time
let y_and_left_bits = y_size_bits(K) + parent_metadata_bits;
// More of the same again
let right_bits_start_offset = u128::BITS as usize - parent_metadata_bits;
// More still
let num_bytes_with_data =
    (y_size_bits(K) + parent_metadata_bits * 2).div_ceil(u8::BITS as usize);
// etc.

I could, of course, rewrite it like this:

let y_and_left_bits = const { y_size_bits(K) + metadata_size_bits(K, PARENT_TABLE_NUMBER) };
let right_bits_start_offset = const { u128::BITS as usize - metadata_size_bits(K, PARENT_TABLE_NUMBER) };
let num_bytes_with_data =
    (y_size_bits(K) + metadata_size_bits(K, PARENT_TABLE_NUMBER) * 2).div_ceil(u8::BITS as usize);

But it quickly becomes hard to maintain. Really annoying when Rustc doesn't evaluate const functions with const arguments early at compile time.

9

It would help if you said why. Performance is, today, always a Quality of Implementation issue. What makes this particular thing different?

10

I don't think I have much to add beyond what I already provided in above examples.

I value Rust because I can write readable code that is also fast most of the time. When need to do something awkward (like inlining variables into explicit const {} blocks above) I consider it an area that could be improved rather than accepting that "it has always been this way and always will be".

fn is runtime-only. Current const fn is runtime-first, compile-time-capable. Comptime is compile-time-only. I want compile-time-first, runtime-capable flavor.

11

How would you express the requirement? "The final assembly has to contain just a constant"? We generally don't specify anything about the final assembly so that won't work.

There's a reason we require const blocks to be explicit. We have one case where we implicitly turn expressions into compile-time expressions, which is const promotion -- if you write &(1 + 2), we turn that into a constant very early during compilation. This caused a whole metric ton of problems because what if the expression panics or loops forever or does other fun things? So while it may sound easy and obvious to say "yeah just const-eval this if the arguments are const", it's really neither easy nor obvious.

From what I understand, rustc -O makes those things constants just fine. Usually there's no reason to care about where during compilation these optimizations happen. So far we're still trying to figure out why you care about that. You mentioned Miri performance and something about rust-gpu, but I wasn't able to follow the GPU part so I am focusing on Miri.

If what you really want is "faster Miri by doing more MIR optimizations", then I think you got the wrong hammer. We currently turn off all optimizations because they might remove UB. It's very hard to audit optimizations for their UB-removal potential. But maybe we could let you mark some functions as "it's fine if you miss some UB here just make it faster please".

You can try yourself what the ceiling is for Miri speedups via MIR optimizations, by running things with RUSTFLAGS="-Zmir-opt-level=3" cargo miri test. So far people typically report only minor speedups.

You also said something about rust-gpu but I wasn't really able to follow that part.

1 Like
12

I'm completely out of my area of expertize here, so my assumption is basically "do what const {} does", but implicitly. Panic at during const-evaluation means compile-time error, at least that would be the most natural outcome from my perspective. I assume execution loops are handled by the compiler in such cases somehow already too. So this isn't something fundamentally new that is being introduced here.

I saw some discussions about this, but mostly abstract rather specific cases of this causing issues.

Just tried this on one slow test, the change was indeed less dramatic than I expected: 896.28s vs 801.45s.

13

you could rewrite it like this:

let (y_and_left_bits, right_bits_start_offset, num_bytes_with_data) = const {
    // const fn with constants known at compile time:
    let parent_metadata_bits = metadata_size_bits(K, PARENT_TABLE_NUMBER);
    // Combined with another const fn, with constants also known at compile time
    let y_and_left_bits = y_size_bits(K) + parent_metadata_bits;
    // More of the same again
    let right_bits_start_offset = u128::BITS as usize - parent_metadata_bits;
    // More still
    let num_bytes_with_data =
        (y_size_bits(K) + parent_metadata_bits * 2).div_ceil(u8::BITS as usize);
    // etc.
    (y_and_left_bits, right_bits_start_offset, num_bytes_with_data)
}
14

From my reading of the suggested change, it would automatically wrap all expressions in const { } blocks that it can, recursively.

One issue I can see with this (contrasting to current solution of manually wrapping in const blocks) is that it'd be very fragile. Somewhere deep in the call stack, changing something from a const-friendly function to non-const would silently change the runtime characteristics (miri runtime duration, for example).

The benefit of const blocks in the source code is that you're documenting where in the code everything must be const. If you call a non-const function by mistake, the compiler helpfully shows you where you defined const and where you violated that constraint.