🦀 Mapping Rust Shaders to SPIR-V (via rust-gpu): This Gist demonstrates how a simple Rust fragment shader is compiled into SPIR-V using rust-gpu. It includes the original Rust code, the corresponding SPIR-V output, and a concise mapping between the two. It also outlines how the rustc_codegen_spirv backend translates MIR to SPIR-V and how SPIR-T …

rust_to_spirv.md

Mapping Rust Shader to SPIR-V

#[spirv(fragment)]				
pub fn main_fs(output: &mut Vec4) {
 *output = vec4(1.0, 0.0, 0.0, 1.0);
}

; SPIR-V
; Version: 1.3
; Generator: Google rspirv; 0
; Bound: 15
; Schema: 0

 OpCapability Shader
 OpCapability VulkanMemoryModel
 OpExtension "SPV_KHR_vulkan_memory_model"
 OpMemoryModel Logical Vulkan
 OpEntryPoint Fragment %1 "main_fs" %2
 OpExecutionMode %1 OriginUpperLeft
 OpDecorate %2 Location 0
 %float = OpTypeFloat 32
 %v4float = OpTypeVector %float 4
%_ptr_Output_v4float = OpTypePointer Output %v4float
 %void = OpTypeVoid
 %9 = OpTypeFunction %void
 %float_1 = OpConstant %float 1
 %float_0 = OpConstant %float 0
 %2 = OpVariable %_ptr_Output_v4float Output
 %14 = OpConstantComposite %v4float %float_1 %float_0 %float_0 %float_1
 %1 = OpFunction %void None %9
 %12 = OpLabel
       OpStore %2 %14
       OpReturn
       OpFunctionEnd

How each part of our Rust shader is mapped to SPIR-V by the rustc_codegen_spirv backend:

1. Header and Capabilities

OpCapability Shader
OpCapability VulkanMemoryModel
OpExtension "SPV_KHR_vulkan_memory_model"
OpMemoryModel Logical Vulkan

These are generated based on the target environment. The backend adds:

• Shader capability for fragment shaders
• VulkanMemoryModel capability and extension for Vulkan targets
• Logical memory model for SPIR-V

2. Entry Point Declaration

#[spirv(fragment)]
pub fn main_fs(output: &mut Vec4) {

Maps to:

OpEntryPoint Fragment %1 "main_fs" %2
OpExecutionMode %1 OriginUpperLeft

• #[spirv(fragment)] → Fragment execution model in OpEntryPoint
• main_fs function name → "main_fs" string literal
• %1 is the function ID for main_fs
• %2 is the variable ID for output parameter
• OriginUpperLeft execution mode is added automatically for fragment shaders

3. Interface Variable Decoration

output: &mut Vec4

Maps to:

OpDecorate %2 Location 0
%_ptr_Output_v4float = OpTypePointer Output %v4float
%2 = OpVariable %_ptr_Output_v4float Output

• &mut Vec4 → Output storage class with pointer to v4float
• Location 0 is assigned automatically (fragment shader outputs start at location 0)

4. Type Declarations

Vec4

Maps to:

%float = OpTypeFloat 32
%v4float = OpTypeVector %float 4
%void = OpTypeVoid
%9 = OpTypeFunction %void

• Vec4 → %v4float (vector of 4 32-bit floats)
• Function return type () → %void
• Function type with void return and no parameters → %9

5. Constants

vec4(1.0, 0.0, 0.0, 1.0)

Maps to:

%float_1 = OpConstant %float 1
%float_0 = OpConstant %float 0
%14 = OpConstantComposite %v4float %float_1 %float_0 %float_0 %float_1

• Individual float literals → OpConstant instructions
• vec4(...) constructor → OpConstantComposite instruction

6. Function Definition

pub fn main_fs(output: &mut Vec4) {
    *output = vec4(1.0, 0.0, 0.0, 1.0);
}

Maps to:

%1 = OpFunction %void None %9
%12 = OpLabel
      OpStore %2 %14
      OpReturn
      OpFunctionEnd

• Function declaration → OpFunction with void return type
• Function body starts with OpLabel
• Assignment *output = ... → OpStore instruction
• End of function → OpReturn followed by OpFunctionEnd

Key Transformations

1. Storage Classes: &mut Vec4 is recognized as an output parameter and gets the Output storage class
2. Entry Point Interface: The output parameter is added to the entry point interface list
3. Decorations: Location decorations are added automatically
4. Execution Modes: Fragment shader execution modes are added automatically
5. Dereferencing: The *output = ... dereference is translated to an OpStore instruction
6. Vector Construction: The vec4(...) call is converted to a constant composite

This demonstrates how the rustc_codegen_spirv backend handles the translation from high-level Rust shader code to low-level SPIR-V instructions, managing everything from type declarations to memory operations.

MIR to SPIR-V Conversion Components

The conversion from MIR to SPIR-V happens across multiple files in the rustc_codegen_spirv crate, not just in abi.rs.

1. Type Conversion (primarily in abi.rs and spirv_type.rs):

   • abi.rs handles how Rust types map to SPIR-V types, especially focusing on ABI concerns
   • spirv_type.rs defines the SpirvType enum and implementation details for type conversion

2. Function Conversion (across multiple files):

   • codegen_cx/declare.rs: Handles function declarations and signatures
   • builder/builder_methods.rs: Implements the actual codegen for function bodies
   • builder/mod.rs: Contains the Builder implementation for emitting SPIR-V instructions

3. Global Variables (in codegen_cx/declare.rs and codegen_cx/constant.rs):

   • Declaration and initialization of global variables
   • Handling of constants and static items

4. Control Flow (in builder/builder_methods.rs):

   • Translation of MIR basic blocks to SPIR-V blocks
   • Implementation of branching, loops, and other control flow structures

5. Entry Points and Shader Interface (in codegen_cx/entry.rs):

   • Processing of shader entry points
   • Setting up shader interface variables (inputs/outputs)

6. Intrinsics and Special Functions (in builder/intrinsics.rs):

   • Implementation of Rust intrinsics as SPIR-V instructions
   • Special handling for built-in functions

Key Files for MIR to SPIR-V Conversion

// Handles type conversion and ABI concerns
impl<'tcx> ConvSpirvType<'tcx> for TyAndLayout<'tcx> {
    fn spirv_type(&self, mut span: Span, cx: &CodegenCx<'tcx>) -> Word {
        // Type conversion logic...
    }
}

// Implements actual codegen for MIR function calls
impl<'a, 'tcx> BuilderMethods<'tcx> for Builder<'a, 'tcx> {
    fn call(
        &mut self,
        callee_ty: Self::Type,
        _fn_attrs: Option<&CodegenFnAttrs>,
        _fn_abi: Option<&FnAbi<'tcx, Ty<'tcx>>>,
        callee: Self::Value,
        args: &[Self::Value],
        funclet: Option<&Self::Funclet>,
        instance: Option<ty::Instance<'tcx>>,
    ) -> Self::Value {
        let span = tracing::span!(tracing::Level::DEBUG, "call");
        let _enter = span.enter();

        if funclet.is_some() {
            self.fatal("TODO: Funclets are not supported");
        }

        // NOTE(eddyb) see the comment on `SpirvValueKind::FnAddr`, this should
        // be fixed upstream, so we never see any "function pointer" values being
        // created just to perform direct calls.
        let (callee_val, result_type, argument_types) = match self.lookup_type(callee.ty) {
            // HACK(eddyb) this seems to be needed, but it's not what `get_fn_addr`
            // produces, are these coming from inside `rustc_codegen_spirv`?
            SpirvType::Function {
                return_type,
                arguments,
            } => {
                assert_ty_eq!(self, callee_ty, callee.ty);
                (callee.def(self), return_type, arguments)
            }
		...
		...
    
    // Many other methods for different MIR operations...
}

// Handles shader entry points and interface
impl<'tcx> CodegenCx<'tcx> {
    // Entry points declare their "interface" (all uniforms, inputs, outputs, etc.) as parameters.
    // spir-v uses globals to declare the interface. So, we need to generate a lil stub for the
    // "real" main that collects all those global variables and calls the user-defined main
    // function.
    pub fn entry_stub(
        &self,
        instance: &Instance<'_>,
        fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
        entry_func: SpirvValue,
        name: String,
        entry: Entry,
    ) {
        let span = self
            .tcx
            .def_ident_span(instance.def_id())
            .unwrap_or_else(|| self.tcx.def_span(instance.def_id()));

        // Entry point processing...
    }
}

Summary

The conversion from MIR to SPIR-V is distributed across multiple files in the codebase, each handling different aspects:

1. abi.rs is crucial for type conversion and ABI concerns
2. builder_methods.rs handles most of the actual instruction emission
3. entry.rs manages shader-specific concerns like entry points
4. declare.rs handles function and global declarations
5. intrinsics.rs implements special functions and operations

This modular approach allows the codebase to separate concerns while maintaining a cohesive translation process from MIR to SPIR-V.

SPIR-T in the rust-gpu Pipeline

SPIR-T (SPIR-V Transformer) is an addition to the rust-gpu pipeline that sits between the initial SPIR-V generation and the final output. Here's where it fits in:

SPIR-T Integration

pub fn link(
    sess: &Session,
    mut inputs: Vec<Module>,
    opts: &Options,
    outputs: &OutputFilenames,
    disambiguated_crate_name_for_dumps: &OsStr,
) -> Result<LinkResult> {
    // HACK(eddyb) this is defined here to allow SPIR-T pretty-printing to apply
    // to SPIR-V being dumped, outside of e.g. `--dump-spirt-passes`.
    let spv_module_to_spv_words_and_spirt_module = |spv_module: &Module| {
        let spv_words;
        let spv_bytes = {
            let _timer = sess.timer("assemble-to-spv_bytes-for-spirt");
            spv_words = spv_module.assemble();
            // FIXME(eddyb) this is wastefully cloning all the bytes, but also
            // `spirt::Module` should have a method that takes `Vec<u32>`.
            spirv_tools::binary::from_binary(&spv_words).to_vec()
        };

        // FIXME(eddyb) should've really been "spirt::Module::lower_from_spv_bytes".
        let lower_from_spv_timer = sess.timer("spirt::Module::lower_from_spv_file");
        let cx = std::rc::Rc::new(spirt::Context::new());
        crate::custom_insts::register_to_spirt_context(&cx);
        (
            spv_words,
            spirt::Module::lower_from_spv_bytes(cx, spv_bytes),
            lower_from_spv_timer,
        )
    };
    
    // ...
    
    // NOTE(eddyb) SPIR-T pipeline is entirely limited to this block.
    {
        let (spv_words, module_or_err, lower_from_spv_timer) =
            spv_module_to_spv_words_and_spirt_module(&output);
        let module = &mut module_or_err.map_err(|e| {
            // Error handling...
        })?;
        
        // Run SPIR-T passes if specified
        if !opts.spirt_passes.is_empty() {
            spirt_passes::run_func_passes(
                module,
                &opts.spirt_passes,
                |name, _module| before_pass(name),
                after_pass,
            );
        }
        
        // ...
        
        // Convert back to SPIR-V
        let spv_words = {
            let _timer = before_pass("spirt::Module::lift_to_spv_module_emitter");
            module.lift_to_spv_module_emitter().unwrap().words
        };
        output = {
            let _timer = sess.timer("parse-spv_words-from-spirt");
            let mut loader = Loader::new();
            rspirv::binary::parse_words(&spv_words, &mut loader).unwrap();
            loader.module()
        };
    }
    // ...
}

The Complete Pipeline with SPIR-T

1. Initial SPIR-V Generation:

   • MIR is translated to SPIR-V using CodegenCx and BuilderSpirv
   • This produces an initial SPIR-V module

2. SPIR-T Transformation (no longer opt-in):

   • The SPIR-V module is converted to SPIR-T representation
   • Various optimization and transformation passes can be applied
   • From the code, we can see passes like reduce, fuse_selects, etc.

3. Back to SPIR-V:

   • The transformed SPIR-T module is converted back to SPIR-V
   • This becomes the final output

Why SPIR-T?

SPIR-T was introduced to provide a more flexible intermediate representation for transformations. From the code and comments, we can see:

1. Up until a certain version SPIR-T was opt-in via RUSTGPU_CODEGEN_ARGS=--spirt (mentioned in CHANGELOG.md) but it no longer is now.
2. It provides optimization passes like reduce which can simplify expressions
3. It has better support for control flow transformations

SPIR-T is particularly useful for:

• Optimizing shader code beyond what the initial codegen can do
• Implementing transformations that are difficult in raw SPIR-V
• Providing a more ergonomic representation for analysis and transformation

Note: SPIR-V uses structured control flow = single entry and exit for each control structure, with no goto statements. This is why SPIR-V requires explicit structuring and why compiling from languages with arbitrary control flow (like Rust's MIR) to SPIR-V requires special transformation passes to convert unstructured control flow into structured form.

Another word for this is structurization.

The SPIR-T passes are defined in crates/rustc_codegen_spirv/src/linker/spirt_passes/ and include functionality for control flow transformations, debug info handling, diagnostics, and optimizations like fusing select operations and reducing expressions.