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

Merged
Firestar99 merged 2 commits into from
Feb 18, 2026
Merged

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
2 commits
Select commit Hold shift + click to select a range
c602fec
Optimize num_traits::Float::powi to use GLSL.std.450 Pow
LegNeato Jan 21, 2026
2ba2af5
Separate tables for `libm` and `num_traits` overwrites
Firestar99 Feb 2, 2026
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
4 changes: 3 additions & 1 deletion crates/rustc_codegen_spirv/src/builder/builder_methods.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3319,6 +3319,8 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {

let libm_intrinsic =
instance_def_id.and_then(|def_id| self.libm_intrinsics.borrow().get(&def_id).copied());
let num_traits_intrinsics = instance_def_id
.and_then(|def_id| self.num_traits_intrinsics.borrow().get(&def_id).copied());
let buffer_load_intrinsic = instance_def_id
.is_some_and(|def_id| self.buffer_load_intrinsics.borrow().contains(&def_id));
let buffer_store_intrinsic = instance_def_id
Expand All @@ -3328,7 +3330,7 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
let from_trait_impl =
instance_def_id.and_then(|def_id| self.from_trait_impls.borrow().get(&def_id).copied());

if let Some(libm_intrinsic) = libm_intrinsic {
if let Some(libm_intrinsic) = libm_intrinsic.or(num_traits_intrinsics) {
let result = self.call_libm_intrinsic(libm_intrinsic, result_type, args);
if result_type != result.ty {
bug!(
Expand Down
12 changes: 11 additions & 1 deletion 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 All @@ -42,7 +43,7 @@ pub enum LibmIntrinsic {
Custom(LibmCustomIntrinsic),
}

pub const TABLE: &[(&str, LibmIntrinsic)] = &[
pub const LIBM_TABLE: &[(&str, LibmIntrinsic)] = &[
("acos", LibmIntrinsic::GLOp(GLOp::Acos)),
("acosf", LibmIntrinsic::GLOp(GLOp::Acos)),
("acosh", LibmIntrinsic::GLOp(GLOp::Acosh)),
Expand Down Expand Up @@ -198,6 +199,9 @@ pub const TABLE: &[(&str, LibmIntrinsic)] = &[
("truncf", LibmIntrinsic::GLOp(GLOp::Trunc)),
];

pub const NUM_TRAITS_TABLE: &[(&str, LibmIntrinsic)] =
&[("powi", LibmIntrinsic::Custom(LibmCustomIntrinsic::Powi))];

impl Builder<'_, '_> {
pub fn call_libm_intrinsic(
&mut self,
Expand Down Expand Up @@ -306,6 +310,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])
}
Comment on lines +313 to +318
Copy link

@foxtran foxtran Feb 17, 2026

Choose a reason for hiding this comment

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

I think for known low powers it would be better to produce explicit sequence of multiplications (or division + sequence of multiplications for negate powers)

Copy link
Member

@Firestar99 Firestar99 Feb 18, 2026

Choose a reason for hiding this comment

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

While generally correct, SPIR-V isn't final machine code. Most drivers will spin up an entire LLVM compiler to turn that SPIR-V into whatever machine code they actually run, which are going to be much better at emitting properly optimized Pow. So the best thing we can do is tell them directly it's a Pow, and not needlessly complicating it.

LibmIntrinsic::Custom(LibmCustomIntrinsic::Remainder) => {
self.undef_zombie(result_type, "Remainder not supported yet")
}
Expand Down
10 changes: 10 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,16 @@ 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.num_traits_intrinsics.get(&item_name) {
self.num_traits_intrinsics
.borrow_mut()
.insert(def_id, intrinsic);
}
}
Comment on lines 169 to 177
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.

Copy link
Member

@Firestar99 Firestar99 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.

We can't get rid of num_traits for silly rust reasons:
core::f32::powf is marked as #[rustc_allow_incoherent_impl], meaning it doesn't actually "exist" unless some other crate declares it. I'm not too familiar on the details, but it seems like when you have std it exists and if you only have core it's declared but doesn't exist and using it will fail with:

  error[E0599]: no function or associated item named `powf` found for type `f32` in the current scope
     --> crates/restir-shader/src/material/pbr/eval.rs:159:25
      |
  159 |     f0 + (1.0 - f0) * f32::powf(f32::clamp(1.0 - cos_theta, 0.0, 1.0), 5.0)
      |                            ^^^^ function or associated item not found in `f32`
      |
      = help: items from traits can only be used if the trait is in scope

see rust-lang/rust#149347 on this sillyness

The workaround has usually been this use statement:

#[cfg(target_arch = "spirv")]
use spirv_std::num_traits::Float;

num_traits with std feature will call std functions and with libm feature / in #[no_std] environments call the appropriate libm function. For SPIR-V, our spirv-std configures the feature for you and we then intercept those libm calls and replace them with intrinsics.

While researching this, I've noticed this has been discussed in 2021 on the embark repo already, and their conclusion was:

I think this is a fully external issue, then - rust-gpu is working "as intended", faithfully compiling the code that it is given, which doesn't use an intrinsic. The way to fix this, then, would be to change upstream code to be less branchy (e.g. requesting that libm adds a powi intrinsic, and having num-traits call it).

I don't mind if we want to merge this anyway, so it's not a trap to end users @LegNeato.

But num-traits isn't a perfect solution either: Technically, it doesn't actually reimplement all f32 functions. For example, it excludes euclid, for which you need to use the Euclid trait, which takes args by reference instead of by value, so you'll get awful code like this (which initially sparked my investigation into this last summer):

#[cfg(feature = "std")]
let rem = |x: f32| x.rem_euclid(1.);
#[cfg(not(feature = "std"))]
let rem = |x: f32| x.rem_euclid(&1.);

Copy link

@voidentente voidentente Jan 22, 2026

Choose a reason for hiding this comment

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

I'm also in favor of replacing num-traits with an in-house solution in the future. It would give Rust-GPU the necessary stability it needs. If num-traits does some change in the future it could silently break Rust-GPU in subtle ways that aren't immediately obvious, which is poison.

Copy link
Collaborator Author

@LegNeato LegNeato Jan 23, 2026

Choose a reason for hiding this comment

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

@Firestar99 created an issue for this, let's just merge for now as the current state is obviously bad and this makes it (slightly) better?


// 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
3 changes: 3 additions & 0 deletions crates/rustc_codegen_spirv/src/codegen_cx/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,8 @@ pub struct CodegenCx<'tcx> {
// FIXME(eddyb) should the maps exist at all, now that the `DefId` is known
// at `call` time, and presumably its high-level details can be looked up?
pub libm_intrinsics: RefCell<FxHashMap<DefId, super::builder::libm_intrinsics::LibmIntrinsic>>,
pub num_traits_intrinsics:
RefCell<FxHashMap<DefId, super::builder::libm_intrinsics::LibmIntrinsic>>,

/// All `panic!(...)`s and builtin panics (from MIR `Assert`s) call into one
/// of these lang items, which we always replace with an "abort".
Expand Down Expand Up @@ -202,6 +204,7 @@ impl<'tcx> CodegenCx<'tcx> {
sym,
instruction_table: InstructionTable::new(),
libm_intrinsics: Default::default(),
num_traits_intrinsics: Default::default(),
panic_entry_points: Default::default(),
fmt_args_new_fn_ids: Default::default(),
fmt_rt_arg_new_fn_ids_to_ty_and_spec: Default::default(),
Expand Down
13 changes: 12 additions & 1 deletion 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 All @@ -33,6 +34,7 @@ pub struct Symbols {
pub attributes: FxHashMap<Symbol, SpirvAttribute>,
pub execution_modes: FxHashMap<Symbol, (ExecutionMode, ExecutionModeExtraDim)>,
pub libm_intrinsics: FxHashMap<Symbol, libm_intrinsics::LibmIntrinsic>,
pub num_traits_intrinsics: FxHashMap<Symbol, libm_intrinsics::LibmIntrinsic>,
}

const BUILTINS: &[(&str, BuiltIn)] = {
Expand Down Expand Up @@ -405,17 +407,25 @@ impl Symbols {
}

let mut libm_intrinsics = FxHashMap::default();
for &(a, b) in libm_intrinsics::TABLE {
for &(a, b) in libm_intrinsics::LIBM_TABLE {
let old = libm_intrinsics.insert(Symbol::intern(a), b);
assert!(old.is_none());
}

let mut num_traits_intrinsics = FxHashMap::default();
for &(a, b) in libm_intrinsics::NUM_TRAITS_TABLE {
let old = num_traits_intrinsics.insert(Symbol::intern(a), b);
assert!(old.is_none());
}

Self {
discriminant: Symbol::intern("discriminant"),
rust_gpu: Symbol::intern("rust_gpu"),
spirv_attr_with_version: Symbol::intern(&spirv_attr_with_version()),
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 All @@ -431,6 +441,7 @@ impl Symbols {
attributes,
execution_modes,
libm_intrinsics,
num_traits_intrinsics,
}
}

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