Skip to content

Optimize constant casts to avoid unnecessary capabilities #302

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.

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
65 changes: 65 additions & 0 deletions crates/rustc_codegen_spirv/src/builder/builder_methods.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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. 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
(
Expand Down
10 changes: 10 additions & 0 deletions crates/rustc_codegen_spirv/src/linker/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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) =
Expand Down
124 changes: 124 additions & 0 deletions crates/rustc_codegen_spirv/src/linker/simple_passes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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<()> {
Expand Down
27 changes: 1 addition & 26 deletions crates/rustc_codegen_spirv/src/spirv_type.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
4 changes: 0 additions & 4 deletions tests/compiletests/ui/dis/asm_op_decorate.stderr
Original file line number Diff line number Diff line change
@@ -1,8 +1,4 @@
OpCapability Shader
OpCapability Float64
OpCapability Int64
OpCapability Int16
OpCapability Int8
OpCapability ShaderClockKHR
OpCapability RuntimeDescriptorArray
OpExtension "SPV_EXT_descriptor_indexing"
Expand Down
19 changes: 19 additions & 0 deletions tests/compiletests/ui/dis/const-float-cast-optimized.rs
Original file line number Diff line number Diff line change
@@ -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;
}
16 changes: 16 additions & 0 deletions tests/compiletests/ui/dis/const-float-cast-optimized.stderr
Original file line number Diff line number Diff line change
@@ -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
31 changes: 31 additions & 0 deletions tests/compiletests/ui/dis/const-float-cast.rs
Original file line number Diff line number Diff line change
@@ -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;
}
22 changes: 22 additions & 0 deletions tests/compiletests/ui/dis/const-float-cast.stderr
Original file line number Diff line number Diff line change
@@ -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
Loading