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

crates/rustc_codegen_spirv/src/builder/libm_intrinsics.rs

@@ -30,6 +30,7 @@ pub enum LibmCustomIntrinsic {
     Tgamma,
     Log1p,
     NextAfter,
+    Powi,
     Remainder,
     RemQuo,
     Scalbn,
@@ -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),
@@ -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")
             }

crates/rustc_codegen_spirv/src/codegen_cx/declare.rs

@@ -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);
+            }
+        }
+
         // 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)

crates/rustc_codegen_spirv/src/symbols.rs

@@ -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,
 
@@ -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"),
 

tests/compiletests/ui/dis/powi.rs

@@ -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);
+}

tests/compiletests/ui/dis/powi.stderr

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