Separate tables for libm and num_traits overwrites

Author: Firestar99

crates/rustc_codegen_spirv/src/builder/builder_methods.rs

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

crates/rustc_codegen_spirv/src/builder/libm_intrinsics.rs

@@ -43,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)),
@@ -158,7 +158,6 @@ 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),
@@ -200,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,

crates/rustc_codegen_spirv/src/codegen_cx/declare.rs

@@ -169,8 +169,10 @@ 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);
+            if let Some(&intrinsic) = self.sym.num_traits_intrinsics.get(&item_name) {
+                self.num_traits_intrinsics
+                    .borrow_mut()
+                    .insert(def_id, intrinsic);
             }
         }
 

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

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

crates/rustc_codegen_spirv/src/symbols.rs

@@ -34,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)] = {
@@ -406,10 +407,17 @@ 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"),
@@ -433,6 +441,7 @@ impl Symbols {
             attributes,
             execution_modes,
             libm_intrinsics,
+            num_traits_intrinsics,
         }
     }