Optimize constant casts to avoid unnecessary capabilities

Add constant folding for widening casts (u8->u32, f32->f64) to avoid creating
intermediate types. Remove capability checks during type creation and add
bidirectional validation in linker to remove unused type capabilities.

Fixes #300

Author: LegNeato

crates/rustc_codegen_spirv/src/builder/builder_methods.rs

@@ -2033,6 +2033,31 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
         if val.ty == dest_ty {
             val
         } else {
+            // If casting a constant, directly create a constant of the target type.
+            // This avoids creating intermediate types that might require additional
+            // capabilities. For example, casting a f16 constant to f32 will directly
+            // create a f32 constant, avoiding the need for Float16 capability if it is
+            // not used elsewhere.
+            if let Some(const_val) = self.builder.lookup_const_scalar(val) {
+                if let (SpirvType::Float(src_width), SpirvType::Float(dst_width)) =
+                    (self.lookup_type(val.ty), self.lookup_type(dest_ty))
+                {
+                    if src_width < dst_width {
+                        // Convert the bit representation to the actual float value
+                        let float_val = match src_width {
+                            32 => Some(f32::from_bits(const_val as u32) as f64),
+                            64 => Some(f64::from_bits(const_val as u64)),
+                            _ => None,
+                        };
+
+                        if let Some(val) = float_val {
+                            return self.constant_float(dest_ty, val);
+                        }
+                    }
+                }
+            }
+
+            // Regular conversion
             self.emit()
                 .f_convert(dest_ty, None, val.def(self))
                 .unwrap()
@@ -2198,6 +2223,46 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
             // I guess?
             return val;
         }
+
+        // If casting a constant, directly create a constant of the target type. This
+        // avoids creating intermediate types that might require additional
+        // capabilities. For example, casting a u8 constant to u32 will directly create
+        // a u32 constant, avoiding the need for Int8 capability if it is not used
+        // elsewhere.
+        if let Some(const_val) = self.builder.lookup_const_scalar(val) {
+            let src_ty = self.lookup_type(val.ty);
+            let dst_ty_spv = self.lookup_type(dest_ty);
+
+            // Try to optimize the constant cast
+            let optimized_result = match (src_ty, dst_ty_spv) {
+                // Integer to integer cast
+                (SpirvType::Integer(src_width, _), SpirvType::Integer(dst_width, _)) => {
+                    // Only optimize if we're widening or keeping the same width.
+                    // This avoids creating the source type when it's safe to do so.
+                    // For narrowing casts (e.g., u32 as u8), we need the proper truncation
+                    // behavior that the regular cast provides.
+                    if src_width <= dst_width {
+                        Some(self.constant_int(dest_ty, const_val))
+                    } else {
+                        None
+                    }
+                }
+                // Bool to integer cast - const_val will be 0 or 1
+                (SpirvType::Bool, SpirvType::Integer(_, _)) => {
+                    Some(self.constant_int(dest_ty, const_val))
+                }
+                // Integer to bool cast - compare with zero
+                (SpirvType::Integer(_, _), SpirvType::Bool) => {
+                    Some(self.constant_bool(self.span(), const_val != 0))
+                }
+                _ => None,
+            };
+
+            if let Some(result) = optimized_result {
+                return result;
+            }
+        }
+
         match (self.lookup_type(val.ty), self.lookup_type(dest_ty)) {
             // sign change
             (

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -477,6 +477,16 @@ pub fn link(
         simple_passes::remove_non_uniform_decorations(sess, &mut output)?;
     }
 
+    {
+        let _timer = sess.timer("link_remove_unused_type_capabilities");
+        simple_passes::remove_unused_type_capabilities(&mut output);
+    }
+
+    {
+        let _timer = sess.timer("link_type_capability_check");
+        simple_passes::check_type_capabilities(sess, &output)?;
+    }
+
     // NOTE(eddyb) SPIR-T pipeline is entirely limited to this block.
     {
         let (spv_words, module_or_err, lower_from_spv_timer) =

crates/rustc_codegen_spirv/src/linker/simple_passes.rs

@@ -7,6 +7,25 @@ use rustc_session::Session;
 use std::iter::once;
 use std::mem::take;
 
+/// Returns the capability required for an integer type of the given width, if any.
+fn capability_for_int_width(width: u32) -> Option<rspirv::spirv::Capability> {
+    match width {
+        8 => Some(rspirv::spirv::Capability::Int8),
+        16 => Some(rspirv::spirv::Capability::Int16),
+        64 => Some(rspirv::spirv::Capability::Int64),
+        _ => None,
+    }
+}
+
+/// Returns the capability required for a float type of the given width, if any.
+fn capability_for_float_width(width: u32) -> Option<rspirv::spirv::Capability> {
+    match width {
+        16 => Some(rspirv::spirv::Capability::Float16),
+        64 => Some(rspirv::spirv::Capability::Float64),
+        _ => None,
+    }
+}
+
 pub fn shift_ids(module: &mut Module, add: u32) {
     module.all_inst_iter_mut().for_each(|inst| {
         if let Some(ref mut result_id) = &mut inst.result_id {
@@ -266,6 +285,111 @@ pub fn check_fragment_insts(sess: &Session, module: &Module) -> Result<()> {
     }
 }
 
+/// Check that types requiring specific capabilities have those capabilities declared.
+///
+/// This function validates that if a module uses types like u8/i8 (requiring Int8),
+/// u16/i16 (requiring Int16), etc., the corresponding capabilities are declared.
+pub fn check_type_capabilities(sess: &Session, module: &Module) -> Result<()> {
+    use rspirv::spirv::Capability;
+
+    // Collect declared capabilities
+    let declared_capabilities: FxHashSet<Capability> = module
+        .capabilities
+        .iter()
+        .map(|inst| inst.operands[0].unwrap_capability())
+        .collect();
+
+    let mut errors = Vec::new();
+
+    for inst in &module.types_global_values {
+        match inst.class.opcode {
+            Op::TypeInt => {
+                let width = inst.operands[0].unwrap_literal_bit32();
+                let signedness = inst.operands[1].unwrap_literal_bit32() != 0;
+                let type_name = if signedness { "i" } else { "u" };
+
+                if let Some(required_cap) = capability_for_int_width(width) {
+                    if !declared_capabilities.contains(&required_cap) {
+                        errors.push(format!(
+                            "`{type_name}{width}` type used without `OpCapability {required_cap:?}`"
+                        ));
+                    }
+                }
+            }
+            Op::TypeFloat => {
+                let width = inst.operands[0].unwrap_literal_bit32();
+
+                if let Some(required_cap) = capability_for_float_width(width) {
+                    if !declared_capabilities.contains(&required_cap) {
+                        errors.push(format!(
+                            "`f{width}` type used without `OpCapability {required_cap:?}`"
+                        ));
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+
+    if !errors.is_empty() {
+        let mut err = sess
+            .dcx()
+            .struct_err("Missing required capabilities for types");
+        for error in errors {
+            err = err.with_note(error);
+        }
+        Err(err.emit())
+    } else {
+        Ok(())
+    }
+}
+
+/// Remove type-related capabilities that are not required by any types in the module.
+///
+/// This function specifically targets Int8, Int16, Int64, Float16, and Float64 capabilities,
+/// removing them if no types in the module require them. All other capabilities are preserved.
+/// This is part of the fix for issue #300 where constant casts were creating unnecessary types.
+pub fn remove_unused_type_capabilities(module: &mut Module) {
+    use rspirv::spirv::Capability;
+
+    // Collect type-related capabilities that are actually needed
+    let mut needed_type_capabilities = FxHashSet::default();
+
+    // Scan all types to determine which type-related capabilities are needed
+    for inst in &module.types_global_values {
+        match inst.class.opcode {
+            Op::TypeInt => {
+                let width = inst.operands[0].unwrap_literal_bit32();
+                if let Some(cap) = capability_for_int_width(width) {
+                    needed_type_capabilities.insert(cap);
+                }
+            }
+            Op::TypeFloat => {
+                let width = inst.operands[0].unwrap_literal_bit32();
+                if let Some(cap) = capability_for_float_width(width) {
+                    needed_type_capabilities.insert(cap);
+                }
+            }
+            _ => {}
+        }
+    }
+
+    // Remove only type-related capabilities that aren't needed
+    module.capabilities.retain(|inst| {
+        let cap = inst.operands[0].unwrap_capability();
+        match cap {
+            // Only remove these type-related capabilities if they're not used
+            Capability::Int8
+            | Capability::Int16
+            | Capability::Int64
+            | Capability::Float16
+            | Capability::Float64 => needed_type_capabilities.contains(&cap),
+            // Keep all other capabilities
+            _ => true,
+        }
+    });
+}
+
 /// Remove all [`Decoration::NonUniform`] if this module does *not* have [`Capability::ShaderNonUniform`].
 /// This allows image asm to always declare `NonUniform` and not worry about conditional compilation.
 pub fn remove_non_uniform_decorations(_sess: &Session, module: &mut Module) -> Result<()> {

crates/rustc_codegen_spirv/src/spirv_type.rs

@@ -3,7 +3,7 @@ use crate::builder_spirv::SpirvValue;
 use crate::codegen_cx::CodegenCx;
 use indexmap::IndexSet;
 use rspirv::dr::Operand;
-use rspirv::spirv::{Capability, Decoration, Dim, ImageFormat, StorageClass, Word};
+use rspirv::spirv::{Decoration, Dim, ImageFormat, StorageClass, Word};
 use rustc_data_structures::fx::FxHashMap;
 use rustc_middle::span_bug;
 use rustc_span::def_id::DefId;
@@ -105,21 +105,6 @@ impl SpirvType<'_> {
                 let result = cx.emit_global().type_int_id(id, width, signedness as u32);
                 let u_or_i = if signedness { "i" } else { "u" };
                 match width {
-                    8 if !cx.builder.has_capability(Capability::Int8) => cx.zombie_with_span(
-                        result,
-                        def_span,
-                        &format!("`{u_or_i}8` without `OpCapability Int8`"),
-                    ),
-                    16 if !cx.builder.has_capability(Capability::Int16) => cx.zombie_with_span(
-                        result,
-                        def_span,
-                        &format!("`{u_or_i}16` without `OpCapability Int16`"),
-                    ),
-                    64 if !cx.builder.has_capability(Capability::Int64) => cx.zombie_with_span(
-                        result,
-                        def_span,
-                        &format!("`{u_or_i}64` without `OpCapability Int64`"),
-                    ),
                     8 | 16 | 32 | 64 => {}
                     w => cx.zombie_with_span(
                         result,
@@ -132,16 +117,6 @@ impl SpirvType<'_> {
             Self::Float(width) => {
                 let result = cx.emit_global().type_float_id(id, width);
                 match width {
-                    16 if !cx.builder.has_capability(Capability::Float16) => cx.zombie_with_span(
-                        result,
-                        def_span,
-                        "`f16` without `OpCapability Float16`",
-                    ),
-                    64 if !cx.builder.has_capability(Capability::Float64) => cx.zombie_with_span(
-                        result,
-                        def_span,
-                        "`f64` without `OpCapability Float64`",
-                    ),
                     16 | 32 | 64 => (),
                     other => cx.zombie_with_span(
                         result,

tests/compiletests/ui/dis/const-float-cast-optimized.rs

@@ -0,0 +1,19 @@
+// Test that constant float widening casts are optimized to avoid creating
+// the smaller float type when not needed elsewhere.
+
+// 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;
+
+#[spirv(fragment)]
+pub fn main(output: &mut f64) {
+    // This should optimize away the f32 type since it's widening
+    const SMALL: f32 = 20.5;
+    let widened = SMALL as f64;
+    *output = widened;
+}

tests/compiletests/ui/dis/const-float-cast-optimized.stderr

@@ -0,0 +1,16 @@
+OpCapability Shader
+OpCapability Float64
+OpCapability ShaderClockKHR
+OpExtension "SPV_KHR_shader_clock"
+OpMemoryModel Logical Simple
+OpEntryPoint Fragment %1 "main" %2
+OpExecutionMode %1 OriginUpperLeft
+%3 = OpString "$OPSTRING_FILENAME/const-float-cast-optimized.rs"
+OpName %2 "output"
+OpDecorate %2 Location 0
+%4 = OpTypeFloat 64
+%5 = OpTypePointer Output %4
+%6 = OpTypeVoid
+%7 = OpTypeFunction %6
+%2 = OpVariable  %5  Output
+%8 = OpConstant  %4  4626463454704697344

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

@@ -0,0 +1,31 @@
+// Test whether float constant casts need optimization
+
+// 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;
+
+#[spirv(fragment)]
+pub fn main(output: &mut f32) {
+    // Test f64 to f32 (narrowing)
+    const BIG: f64 = 123.456;
+    let narrowed = BIG as f32;
+    *output = narrowed;
+
+    // Test f32 to f64 (widening) - this might create f32 type unnecessarily
+    const SMALL: f32 = 20.5;
+    let widened = SMALL as f64;
+    *output += widened as f32;
+
+    let kept: f32 = 1.0 + SMALL;
+    *output += kept;
+
+    // Test integer to float
+    const INT: u32 = 42;
+    let as_float = INT as f32;
+    *output += as_float;
+}

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

@@ -0,0 +1,22 @@
+OpCapability Shader
+OpCapability Float64
+OpCapability ShaderClockKHR
+OpExtension "SPV_KHR_shader_clock"
+OpMemoryModel Logical Simple
+OpEntryPoint Fragment %1 "main" %2
+OpExecutionMode %1 OriginUpperLeft
+%3 = OpString "$OPSTRING_FILENAME/const-float-cast.rs"
+OpName %2 "output"
+OpDecorate %2 Location 0
+%4 = OpTypeFloat 32
+%5 = OpTypePointer Output %4
+%6 = OpTypeVoid
+%7 = OpTypeFunction %6
+%8 = OpTypeFloat 64
+%9 = OpConstant  %8  4638387860618067575
+%2 = OpVariable  %5  Output
+%10 = OpConstant  %8  4626463454704697344
+%11 = OpConstant  %4  1065353216
+%12 = OpConstant  %4  1101266944
+%13 = OpTypeInt 32 0
+%14 = OpConstant  %13  42

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

@@ -0,0 +1,22 @@
+// Test that constant integer 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 = 20;
+
+#[spirv(fragment)]
+pub fn main(output: &mut u32) {
+    let position = 2u32;
+    // This cast should be optimized to directly create a u32 constant with value 20,
+    // avoiding the creation of a u8 type that would require Int8 capability
+    let global_y_offset_bits = position * K as u32;
+    *output = global_y_offset_bits;
+}