Optimize From::from calls with constant arguments

Author: LegNeato

crates/rustc_codegen_spirv/src/builder/builder_methods.rs

@@ -3193,6 +3193,8 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
             .and_then(|def_id| self.buffer_store_intrinsics.borrow().get(&def_id).copied());
         let is_panic_entry_point = instance_def_id
             .is_some_and(|def_id| self.panic_entry_points.borrow().contains(&def_id));
+        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 {
             let result = self.call_libm_intrinsic(libm_intrinsic, result_type, args);
@@ -3204,8 +3206,10 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
                     self.debug_type(result.ty),
                 );
             }
-            result
-        } else if is_panic_entry_point {
+            return result;
+        }
+
+        if is_panic_entry_point {
             // HACK(eddyb) Rust 2021 `panic!` always uses `format_args!`, even
             // in the simple case that used to pass a `&str` constant, which
             // would not remain reachable in the SPIR-V - but `format_args!` is
@@ -3678,24 +3682,75 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
             // HACK(eddyb) redirect any possible panic call to an abort, to avoid
             // needing to materialize `&core::panic::Location` or `format_args!`.
             self.abort_with_kind_and_message_debug_printf("panic", message, debug_printf_args);
-            self.undef(result_type)
-        } else if let Some(mode) = buffer_load_intrinsic {
-            self.codegen_buffer_load_intrinsic(result_type, args, mode)
-        } else if let Some(mode) = buffer_store_intrinsic {
+            return self.undef(result_type);
+        }
+
+        if let Some(mode) = buffer_load_intrinsic {
+            return self.codegen_buffer_load_intrinsic(result_type, args, mode);
+        }
+
+        if let Some(mode) = buffer_store_intrinsic {
             self.codegen_buffer_store_intrinsic(args, mode);
 
             let void_ty = SpirvType::Void.def(rustc_span::DUMMY_SP, self);
-            SpirvValue {
+            return SpirvValue {
                 kind: SpirvValueKind::IllegalTypeUsed(void_ty),
                 ty: void_ty,
+            };
+        }
+
+        if let Some((source_ty, target_ty)) = from_trait_impl {
+            // Optimize From::from calls with constant arguments to avoid creating intermediate types.
+            if let [arg] = args {
+                if let Some(const_val) = self.builder.lookup_const_scalar(*arg) {
+                    use rustc_middle::ty::{FloatTy, IntTy, UintTy};
+
+                    let optimized_result = match (source_ty.kind(), target_ty.kind()) {
+                        // Unsigned integer widening conversions
+                        (
+                            ty::Uint(UintTy::U8),
+                            ty::Uint(UintTy::U16 | UintTy::U32 | UintTy::U64 | UintTy::U128),
+                        )
+                        | (
+                            ty::Uint(UintTy::U16),
+                            ty::Uint(UintTy::U32 | UintTy::U64 | UintTy::U128),
+                        )
+                        | (ty::Uint(UintTy::U32), ty::Uint(UintTy::U64 | UintTy::U128))
+                        | (ty::Uint(UintTy::U64), ty::Uint(UintTy::U128))
+                        // Signed integer widening conversions
+                        | (
+                            ty::Int(IntTy::I8),
+                            ty::Int(IntTy::I16 | IntTy::I32 | IntTy::I64 | IntTy::I128),
+                        )
+                        | (ty::Int(IntTy::I16), ty::Int(IntTy::I32 | IntTy::I64 | IntTy::I128))
+                        | (ty::Int(IntTy::I32), ty::Int(IntTy::I64 | IntTy::I128))
+                        | (ty::Int(IntTy::I64), ty::Int(IntTy::I128)) => {
+                            Some(self.constant_int(result_type, const_val))
+                        }
+
+                        // Float widening conversions: f32->f64
+                        (ty::Float(FloatTy::F32), ty::Float(FloatTy::F64)) => {
+                            let float_val = f32::from_bits(const_val as u32) as f64;
+                            Some(self.constant_float(result_type, float_val))
+                        }
+
+                        // No optimization for narrowing conversions or unsupported types
+                        _ => None,
+                    };
+
+                    if let Some(result) = optimized_result {
+                        return result;
+                    }
+                }
             }
-        } else {
-            let args = args.iter().map(|arg| arg.def(self)).collect::<Vec<_>>();
-            self.emit()
-                .function_call(result_type, None, callee_val, args)
-                .unwrap()
-                .with_type(result_type)
         }
+
+        // Default: emit a regular function call
+        let args = args.iter().map(|arg| arg.def(self)).collect::<Vec<_>>();
+        self.emit()
+            .function_call(result_type, None, callee_val, args)
+            .unwrap()
+            .with_type(result_type)
     }
 
     fn zext(&mut self, val: Self::Value, dest_ty: Self::Type) -> Self::Value {

crates/rustc_codegen_spirv/src/codegen_cx/declare.rs

@@ -172,6 +172,30 @@ impl<'tcx> CodegenCx<'tcx> {
             }
         }
 
+        // Check if this is a From trait implementation
+        if let Some(impl_def_id) = self.tcx.impl_of_method(def_id) {
+            if let Some(trait_ref) = self.tcx.impl_trait_ref(impl_def_id) {
+                let trait_def_id = trait_ref.skip_binder().def_id;
+
+                // Check if this is the From trait.
+                let trait_path = self.tcx.def_path_str(trait_def_id);
+                if matches!(
+                    trait_path.as_str(),
+                    "core::convert::From" | "std::convert::From"
+                ) {
+                    // Extract the source and target types from the trait substitutions
+                    let trait_args = trait_ref.skip_binder().args;
+                    if let (Some(target_ty), Some(source_ty)) =
+                        (trait_args.types().nth(0), trait_args.types().nth(1))
+                    {
+                        self.from_trait_impls
+                            .borrow_mut()
+                            .insert(def_id, (source_ty, target_ty));
+                    }
+                }
+            }
+        }
+
         if [
             self.tcx.lang_items().panic_fn(),
             self.tcx.lang_items().panic_fmt(),

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -84,6 +84,10 @@ pub struct CodegenCx<'tcx> {
     /// Intrinsic for storing a `<T>` into a `&[u32]`. The `PassMode` is the mode of the `<T>`.
     pub buffer_store_intrinsics: RefCell<FxHashMap<DefId, &'tcx PassMode>>,
 
+    /// Maps `DefId`s of `From::from` method implementations to their source and target types.
+    /// Used to optimize constant conversions like `u32::from(42u8)` to avoid creating the source type.
+    pub from_trait_impls: RefCell<FxHashMap<DefId, (Ty<'tcx>, Ty<'tcx>)>>,
+
     /// Some runtimes (e.g. intel-compute-runtime) disallow atomics on i8 and i16, even though it's allowed by the spec.
     /// This enables/disables them.
     pub i8_i16_atomics_allowed: bool,
@@ -203,6 +207,7 @@ impl<'tcx> CodegenCx<'tcx> {
             fmt_rt_arg_new_fn_ids_to_ty_and_spec: Default::default(),
             buffer_load_intrinsics: Default::default(),
             buffer_store_intrinsics: Default::default(),
+            from_trait_impls: Default::default(),
             i8_i16_atomics_allowed: false,
             codegen_args,
         }

tests/compiletests/ui/dis/const-from-cast.rs

@@ -0,0 +1,22 @@
+// Test that constant integer from casts are optimized to avoid creating intermediate
+// types that would require additional capabilities (e.g., Int8 capability for u8).
+
+// build-pass
+// compile-flags: -C llvm-args=--disassemble-globals
+// normalize-stderr-test "OpCapability VulkanMemoryModel\n" -> ""
+// normalize-stderr-test "OpSource .*\n" -> ""
+// normalize-stderr-test "OpExtension .SPV_KHR_vulkan_memory_model.\n" -> ""
+// normalize-stderr-test "OpMemoryModel Logical Vulkan" -> "OpMemoryModel Logical Simple"
+
+use spirv_std::spirv;
+
+const K: u8 = 42;
+
+#[spirv(fragment)]
+pub fn main(output: &mut u32) {
+    let position = 2u32;
+    // This cast should be optimized to directly create a u32 constant with value 42,
+    // avoiding the creation of a u8 type that would require Int8 capability
+    let global_y_offset_bits = u32::from(K);
+    *output = global_y_offset_bits;
+}

tests/compiletests/ui/dis/const-from-cast.stderr

@@ -0,0 +1,15 @@
+OpCapability Shader
+OpCapability ShaderClockKHR
+OpExtension "SPV_KHR_shader_clock"
+OpMemoryModel Logical Simple
+OpEntryPoint Fragment %1 "main" %2
+OpExecutionMode %1 OriginUpperLeft
+%3 = OpString "$OPSTRING_FILENAME/const-from-cast.rs"
+OpName %2 "output"
+OpDecorate %2 Location 0
+%4 = OpTypeInt 32 0
+%5 = OpTypePointer Output %4
+%6 = OpTypeVoid
+%7 = OpTypeFunction %6
+%2 = OpVariable  %5  Output
+%8 = OpConstant  %4  42

tests/compiletests/ui/lang/consts/u32-from-u64-fail.rs

@@ -0,0 +1,16 @@
+// Test that u32::from(u64) fails to compile since From<u64> is not implemented for u32
+// This ensures our From trait optimization doesn't accidentally allow invalid conversions
+
+// build-fail
+
+use spirv_std::spirv;
+
+const K: u64 = 42;
+
+#[spirv(fragment)]
+pub fn main(output: &mut u32) {
+    // This should fail to compile because From<u64> is not implemented for u32
+    // (u64 to u32 is a narrowing conversion that could lose data)
+    let value = u32::from(K);
+    *output = value;
+}

tests/compiletests/ui/lang/consts/u32-from-u64-fail.stderr

@@ -0,0 +1,17 @@
+error[E0277]: the trait bound `u32: From<u64>` is not satisfied
+  --> $DIR/u32-from-u64-fail.rs:14:17
+   |
+14 |     let value = u32::from(K);
+   |                 ^^^ the trait `From<u64>` is not implemented for `u32`
+   |
+   = help: the following other types implement trait `From<T>`:
+             `u32` implements `From<Char>`
+             `u32` implements `From<Ipv4Addr>`
+             `u32` implements `From<bool>`
+             `u32` implements `From<char>`
+             `u32` implements `From<u16>`
+             `u32` implements `From<u8>`
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0277`.