cubecl_wgpu/compiler/wgsl/
compiler.rs

1use super::shader::ComputeShader;
2use crate::compiler::wgsl::{
3    self, EnableFeaturesPass, builtin::LowerBuiltinsPass, lower::LowerOpsWgslPass,
4    metadata::declare_info, rewrite_args, shared_memory_size, types,
5};
6use cubecl_server::kernel::BufferIOAttr;
7
8use cubecl_core::{
9    WgpuCompilationOptions,
10    post_processing::{
11        checked_io::{CheckedIo, CheckedIoPass},
12        minifloat::{
13            Fp8Container, LowerMinifloatCast, LowerMinifloatCastPass, LowerMinifloatCompare,
14            LowerMinifloatComparePass,
15        },
16        saturating::LowerSaturatingArithmeticPass,
17        unroll::UnrollPass,
18    },
19};
20use cubecl_environment::backtrace::BackTrace;
21use cubecl_ir::{
22    ContextExt,
23    features::EnumSet,
24    pliron::{
25        builtin::ops::{FuncOp, ModuleOp},
26        operation::verify_operation,
27        opts::{dce::DCEPass, mem2reg::Mem2RegPass},
28    },
29    prelude::{AnalysisManager, NestedOpsPass, Op, OpPass, PMConfig, Pass, Passes},
30    rewrite::SimplifyOpsPass,
31    settings::Dim3,
32};
33use cubecl_opt::passes::{
34    annotate_buffer_visibility::AnnotateGlobalVisibilityPass, inst_combine::InstCombinePass,
35    sccp::SCCPPass, simple_cse::SimpleCSEPass, sroa::SROAPass,
36};
37use cubecl_server::compiler::CompilationError;
38use cubecl_server::kernel;
39
40const MAX_VECTOR_SIZE: usize = 4;
41
42pub struct KernelInfo {
43    pub cube_dim: Dim3,
44}
45
46/// Wgsl Compiler.
47#[derive(Clone, Default)]
48pub struct WgslCompiler;
49
50impl core::fmt::Debug for WgslCompiler {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        f.write_str("WgslCompiler")
53    }
54}
55
56impl cubecl_core::Compiler for WgslCompiler {
57    type Representation = ComputeShader;
58    type CompilationOptions = WgpuCompilationOptions;
59
60    fn buffer_io(repr: &Self::Representation) -> Option<Vec<BufferIOAttr>> {
61        Some(repr.io.clone())
62    }
63
64    fn compile(
65        &mut self,
66        shader: kernel::KernelDefinition,
67        compilation_options: &Self::CompilationOptions,
68    ) -> Result<Self::Representation, CompilationError> {
69        self.compile_shader(shader, compilation_options)
70    }
71
72    fn extension(&self) -> &'static str {
73        "wgsl"
74    }
75
76    fn lang_tag(&self) -> &'static str {
77        "wgsl"
78    }
79}
80
81impl WgslCompiler {
82    fn compile_shader(
83        &mut self,
84        value: kernel::KernelDefinition,
85        compilation_options: &WgpuCompilationOptions,
86    ) -> Result<wgsl::ComputeShader, CompilationError> {
87        #[cfg(target_family = "wasm")]
88        crate::compiler::wgsl::ops::wasm_inventory_root();
89
90        let errors = value.body.pop_errors();
91        if !errors.is_empty() {
92            let mut reason = "Can't compile wgsl kernel".to_string();
93            for error in errors {
94                reason += error.as_str();
95                reason += "\n";
96            }
97
98            return Err(CompilationError::Validation {
99                reason,
100                backtrace: BackTrace::capture(),
101            });
102        }
103
104        #[cfg(feature = "pliron-dump")]
105        let ir_printing_dir = kernel_dir_name(&value.settings.kernel_name);
106
107        let module = value.body.state().module;
108        let entry_func = value.body.state().entry_func;
109        let module_op = module.get_operation();
110        let mut ctx = value.body.into_context().expect("Should be unique");
111        ctx.set_aux_ty(value.info);
112        ctx.set_aux_ty(KernelInfo {
113            cube_dim: value.settings.cube_dim,
114        });
115        ctx.set_aux_ty(*compilation_options);
116
117        #[cfg(feature = "pliron-dump")]
118        if let Some(print_dir) = &ir_printing_dir {
119            use pliron::printable::Printable;
120            let str = std::format!("{}", module_op.disp(&ctx));
121            std::fs::write(print_dir.join("initial.plir"), &str).unwrap();
122        }
123
124        verify_operation(module_op, &ctx)?;
125        types::check_fp8_lanes(&ctx, module_op)?;
126
127        let config = PMConfig {
128            #[cfg(feature = "pliron-dump")]
129            ir_printing_dir,
130            print_after_all: cfg!(feature = "pliron-dump"),
131            ..Default::default()
132        };
133
134        let mut analyses = AnalysisManager::default();
135        analyses.set_config(config);
136
137        let mut passes = OpPass::<ModuleOp, Passes>::default();
138        let mut func_passes = OpPass::<FuncOp, Passes>::default();
139
140        func_passes.add_pass(SROAPass);
141        func_passes.add_pass(CheckedIoPass::new(CheckedIo::new(
142            value.settings.execution_mode,
143            value.settings.kernel_name.clone(),
144        )));
145        func_passes.add_pass(UnrollPass::new(MAX_VECTOR_SIZE));
146        // After the unroll so that an fp8 vector is at most one word, see `types.rs`.
147        func_passes.add_pass(LowerMinifloatCastPass::new(LowerMinifloatCast::new(
148            EnumSet::empty(),
149            Fp8Container::Words,
150        )));
151        func_passes.add_pass(LowerMinifloatComparePass::new(LowerMinifloatCompare::new(
152            Fp8Container::Words,
153        )));
154
155        func_passes.add_pass(LowerOpsWgslPass::default());
156        func_passes.add_pass(LowerSaturatingArithmeticPass::default());
157        func_passes.add_pass(LowerBuiltinsPass);
158
159        func_passes.add_pass(SCCPPass);
160        func_passes.add_pass(InstCombinePass::default());
161        func_passes.add_pass(SimpleCSEPass::without_memory());
162        func_passes.add_pass(SimplifyOpsPass::default());
163        func_passes.add_pass(DCEPass);
164        func_passes.add_pass(SROAPass);
165
166        // SCCP/DCE may unlock more mem2reg opportunities, and vice versa. So we do a sandwich.
167        func_passes.add_pass(Mem2RegPass);
168
169        func_passes.add_pass(SROAPass);
170        func_passes.add_pass(SCCPPass);
171        func_passes.add_pass(SimpleCSEPass::with_memory());
172        func_passes.add_pass(SimplifyOpsPass::default());
173        func_passes.add_pass(DCEPass);
174
175        passes.add_pass(NestedOpsPass::new(func_passes));
176        passes.add_pass(AnnotateGlobalVisibilityPass);
177        passes.add_pass(EnableFeaturesPass);
178
179        passes.run(module_op, &mut ctx, &mut analyses)?;
180
181        // Before `rewrite_args` erases the arguments and widens their
182        // visibility for the shader.
183        let io = cubecl_core::ir::attributes::buffer_io_by_position(&ctx, entry_func)
184            .into_iter()
185            .collect();
186        let buffers = rewrite_args(&mut ctx, entry_func);
187        declare_info(&mut ctx, entry_func, buffers.len());
188        let shared_memory_size = shared_memory_size(&ctx, module_op);
189
190        verify_operation(module.get_operation(), &ctx)?;
191
192        Ok(ComputeShader {
193            buffers,
194            io,
195            shared_memory_size,
196            ctx,
197        })
198    }
199}
200
201#[cfg(feature = "pliron-dump")]
202pub fn kernel_dir_name(name: &str) -> Option<std::path::PathBuf> {
203    if let Ok(dir) = std::env::var("CUBECL_DEBUG_PLIRON") {
204        let path = sanitize_filename::sanitize_with_options(
205            name,
206            sanitize_filename::Options {
207                replacement: "_",
208                ..Default::default()
209            },
210        );
211        let dir = std::path::PathBuf::from(dir).join(&path);
212        std::fs::create_dir_all(&dir).unwrap();
213        Some(dir)
214    } else {
215        None
216    }
217}
cubecl_wgpu/compiler/wgsl/compiler.rs

cubecl_wgpu/compiler/wgsl/
compiler.rs
