Skip to content

Optimize num_traits::Float::powi to use GLSL.std.450 Pow #518

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
LegNeato wants to merge 1 commit into
base: main
Choose a base branch
from
+43 −0
Open

Optimize num_traits::Float::powi to use GLSL.std.450 Pow #518

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions crates/rustc_codegen_spirv/src/builder/libm_intrinsics.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ pub enum LibmCustomIntrinsic {
Tgamma,
Log1p,
NextAfter,
Powi,
Remainder,
RemQuo,
Scalbn,
Expand Down Expand Up @@ -157,6 +158,7 @@ pub const TABLE: &[(&str, LibmIntrinsic)] = &[
),
("pow", LibmIntrinsic::GLOp(GLOp::Pow)),
("powf", LibmIntrinsic::GLOp(GLOp::Pow)),
("powi", LibmIntrinsic::Custom(LibmCustomIntrinsic::Powi)),
(
"remainder",
LibmIntrinsic::Custom(LibmCustomIntrinsic::Remainder),
Expand Down Expand Up @@ -306,6 +308,12 @@ impl Builder<'_, '_> {
LibmIntrinsic::Custom(LibmCustomIntrinsic::NextAfter) => {
self.undef_zombie(result_type, "NextAfter not supported yet")
}
LibmIntrinsic::Custom(LibmCustomIntrinsic::Powi) => {
assert_eq!(args.len(), 2);
// Convert integer exponent to float, then use GLOp::Pow
let float_exp = self.sitofp(args[1], args[0].ty);
self.gl_op(GLOp::Pow, result_type, [args[0], float_exp])
}
LibmIntrinsic::Custom(LibmCustomIntrinsic::Remainder) => {
self.undef_zombie(result_type, "Remainder not supported yet")
}
Expand Down
8 changes: 8 additions & 0 deletions crates/rustc_codegen_spirv/src/codegen_cx/declare.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,14 @@ impl<'tcx> CodegenCx<'tcx> {
}
}

// Check for usage of `num_traits` intrinsics (like Float::powi) that we can optimize
if self.tcx.crate_name(def_id.krate) == self.sym.num_traits && !def_id.is_local() {
let item_name = self.tcx.item_name(def_id);
if let Some(&intrinsic) = self.sym.libm_intrinsics.get(&item_name) {
self.libm_intrinsics.borrow_mut().insert(def_id, intrinsic);
}
}
Comment on lines +169 to +175
Copy link
Contributor

@nazar-pc nazar-pc Jan 21, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have no idea how things work more generally, but why are things like this even necessary? I'd expect the generic machinery of the compiler to optimize it to some canonical form that the target would like the most.

So it shouldn't make any difference if one writes x.powi(2) or x * x or uses a few layers of zero-cost abstractions in the process, should all result in identical SPIR-V, just like it does with LLVM. What am I missing?

Copy link
Collaborator Author

@LegNeato LegNeato Jan 22, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The problem is we are using num_traits to get the APIs we need off primitives, and num_traits knows nothing about the intrinsics available on the platform. It is using a loop as impl https://github.com/rust-num/num-traits/blob/master/src/pow.rs#L179. Pretty much no compiler would be able to tell that code can be replaced with a cast and a powi, and rustc doesn't. Instead somewhere needs to map the call / the semantics to the underlying fast impl, which is what the backend is supposed to do. It already does it for many ops, we just missed a spot (I guess? This all predates me).

One can argue we shouldn't use num_traits and instead should just hang our own impls off the types, I have no clue why it was chosen vs that route. Probably a make it work first then make it fast tradeoff (as the num_traits impl does work on spirv).

I have a new optimizer where we can control things much more rather than relying on spirv-opt, but this replacement happens before the optimizer.

Copy link
Contributor

@nazar-pc nazar-pc Jan 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ouch, that definitely looks like a problem. Maybe worth creating an issue to look into getting rid of num_traits? Doing things like looping instead of native exponentiation is definitely going to hurt and there might be more examples like it.


// Check if this is a From trait implementation
if let Some(impl_def_id) = self.tcx.impl_of_method(def_id)
&& let Some(trait_ref) = self.tcx.impl_trait_ref(impl_def_id)
Expand Down
2 changes: 2 additions & 0 deletions crates/rustc_codegen_spirv/src/symbols.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ pub struct Symbols {
pub vector: Symbol,
pub v1: Symbol,
pub libm: Symbol,
pub num_traits: Symbol,
pub entry_point_name: Symbol,
pub spv_khr_vulkan_memory_model: Symbol,

Expand Down Expand Up @@ -416,6 +417,7 @@ impl Symbols {
vector: Symbol::intern("vector"),
v1: Symbol::intern("v1"),
libm: Symbol::intern("libm"),
num_traits: Symbol::intern("num_traits"),
entry_point_name: Symbol::intern("entry_point_name"),
spv_khr_vulkan_memory_model: Symbol::intern("SPV_KHR_vulkan_memory_model"),

Expand Down
13 changes: 13 additions & 0 deletions tests/compiletests/ui/dis/powi.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
// Test that `Float::powi` uses GLSL.std.450 Pow instead of a loop-based implementation.
// See https://github.com/Rust-GPU/rust-gpu/issues/516

// build-pass
// compile-flags: -C llvm-args=--disassemble-entry=main

use spirv_std::num_traits::Float;
use spirv_std::spirv;

#[spirv(fragment)]
pub fn main(input: f32, output: &mut f32) {
*output = input.powi(2);
}
12 changes: 12 additions & 0 deletions tests/compiletests/ui/dis/powi.stderr
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
%1 = OpFunction %2 None %3
%4 = OpLabel
OpLine %5 11 12
%6 = OpLoad %7 %8
OpLine %5 12 20
%9 = OpConvertSToF %7 %10
%11 = OpExtInst %7 %12 26 %6 %9
OpLine %5 12 4
OpStore %13 %11
OpNoLine
OpReturn
OpFunctionEnd
Loading