LLVM failed to use the knowledge from a never-overflow assumption

[dtcxzyw]

Hi all, I am an LLVM developer working on the middle-end optimizations.

Recently I found an anti-pattern in some Rust applications from my benchmark. All of them come from hashbrown::raw::RawTableInner::free_buckets:

hashbrown/src/raw/mod.rs

Lines 3290 to 3294 in 274c7bb

	// Avoid `Option::unwrap_or_else` because it bloats LLVM IR.
	let (layout, ctrl_offset) = match table_layout.calculate_layout_for(self.buckets()) {
	Some(lco) => lco,
	None => unsafe { hint::unreachable_unchecked() },
	};

LLVM IR:

define i32 @src(i32 %x, i32 %y) {
  %res = call { i32, i1 } @llvm.umul.with.overflow.i32(i32 %x, i32 %y)
  %ov = extractvalue { i32, i1 } %res, 1
  %nowrap = xor i1 %ov, true
  tail call void @llvm.assume(i1 %nowrap)
  %val = extractvalue { i32, i1 } %res, 0
  ret i32 %val
}
define i32 @tgt(i32 %x, i32 %y) {
  %res = mul nuw i32 %x, %y
  ret i32 %res
}

Alive2: https://alive2.llvm.org/ce/z/YBRjxc

Currently LLVM cannot use the knowledge from a never-overflow assumption. I have a draft patch for the fold and it enables more optimizations in downstream users of hashbrown. But I am not willing to do an application-specific optimization in LLVM. So I file this issue to see whether we can fix it in hashbrown.

cc @nikic @dianqk

[dianqk]

Could you further explain why alive2 proves the correctness of the transformation, but we cannot transform it in LLVM?
BTW, we should probably find a reduce example to report at rust. But I can also do this later.

[dtcxzyw]

but we cannot transform it in LLVM?

It is just not implemented. But I am not willing to do this in LLVM. We should fix it in hashbrown or rustc.

[dtcxzyw]

An example from tokio-rs: https://godbolt.org/z/Gr55dfMc8

[nikic]

Looks like the same as llvm/llvm-project#80637. As I noted there, this technically loses information, but I guess in practice it's a lot better to replace with mul nuw here. I wouldn't be opposed to doing that in LLVM, as this kind of pattern is probably somewhat common. It arises whenever you have uses of checked_* in some method, and then it gets used in multiple places, some of them checking the result, and the others doing something unwrap_unchecked style. It avoids the duplication of implementing it both with checked_* and unchecked_* arithmetic.

@dtcxzyw Do you have some example where the replacement enables follow-up optimization in practice? As far as I can tell your examples will be lowered to the same asm after assumes are dropped in CGP.

[dtcxzyw]

@nikic See the IR diff in https://github.com/dtcxzyw/llvm-opt-benchmark/pull/289/files:
bench/diesel-rs/optimized/1hskzwx2vflsavf7.ll: https://godbolt.org/z/8e6qd9Gjo
bench/image-rs/optimized/2mngkegtim1o10y3.ll: enables factorization: (%2505 + 1) * 392 + 15 --> %2505 * 392 + 407

[dianqk]

It arises whenever you have uses of checked_* in some method, and then it gets used in multiple places, some of them checking the result, and the others doing something unwrap_unchecked style. It avoids the duplication of implementing it both with checked_* and unchecked_* arithmetic.

As @nikic said, we want to reuse calculate_layout_for.

hashbrown/src/raw/mod.rs

Lines 1671 to 1674 in 274c7bb

	match Self::TABLE_LAYOUT.calculate_layout_for(self.table.buckets()) {
	Some(lco) => lco,
	None => unsafe { hint::unreachable_unchecked() },
	};

hashbrown/src/raw/mod.rs

Lines 1745 to 1748 in 274c7bb

	let (layout, ctrl_offset) = match table_layout.calculate_layout_for(buckets) {
	Some(lco) => lco,
	None => return Err(fallibility.capacity_overflow()),
	};

[RalfJung]

But I am not willing to do an application-specific optimization in LLVM. So I file this issue to see whether we can fix it in hashbrown.

So as argued above by others, this is not an application-specific pattern. It is a pretty natural pattern for Rust programs: one has a checked operation that can fail in a safe way, and then one turns that into the unchecked variant (with UB on failure) by doing unwrap_unchecked() (which expands to the match statement from the issue description).

I'm not sure how else you are suggesting that code should be written, but it'd clearly be bad if calculate_layout_for had to be written twice.

[workingjubilee]

Even aside from the pattern overall being visible in Rust code beyond hashbrown, most of the code in hashbrown is transitively included in every single Rust program that uses a hashmap, so "application-specific" is an amusing way to describe "probably 20% of all Rust programs ever? higher, even?"

[workingjubilee]

@dtcxzyw Most of the callers that eventually arrived in hashbrown::raw::RawTableInner::free_buckets had std::collections::HashMap::some_fn higher on the stack, right? (I don't know what your tooling is like as to whether it can make this kind of thing obvious)

[dtcxzyw]

@dtcxzyw Most of the callers that eventually arrived in hashbrown::raw::RawTableInner::free_buckets had std::collections::HashMap::some_fn higher on the stack, right? (I don't know what your tooling is like as to whether it can make this kind of thing obvious)

Yeah.

Even aside from the pattern overall being visible in Rust code beyond hashbrown, most of the code in hashbrown is transitively included in every single Rust program that uses a hashmap, so "application-specific" is an amusing way to describe "probably 20% of all Rust programs ever? higher, even?"

I believe there is a better way to fix this issue by adding some MIR optimizations or modifying hashbrown itself. It is tricky to handle it in LLVM :(
See llvm/llvm-project#84016.

[nikic]

Even aside from the pattern overall being visible in Rust code beyond hashbrown, most of the code in hashbrown is transitively included in every single Rust program that uses a hashmap, so "application-specific" is an amusing way to describe "probably 20% of all Rust programs ever? higher, even?"

I think you're missing the point here. When it comes to optimization issues, in Rust we are in the very fortunate position that the issue can be addressed either by adding new optimizations, or by changing library code to be more amenable to optimization. The fact that we can change the hashbrown implementation and have ~all Rust code benefit is the point, not an analysis mistake.

<rant> This is contrary to the stupidity prevalent in C/C++, where an incredible amount of compiler engineering resources go into implementing completely ridiculous optimizations that target specific patterns in SPEC benchmarks, where the proper thing to do would have been to just change the benchmark code. </rant>

(This is just for context -- as said above, I'm not convinced that modifying hashbrown is indeed the best solution here.)