1#![allow(unknown_lints, unnecessary_transmutes)]
2
3use std::{
4 fmt::{Debug, Display},
5 sync::Arc,
6};
7
8use cubecl_core::prelude::Visibility;
9use cubecl_opt::Optimizer;
10use rspirv::{binary::Disassemble, dr::Module};
11
12mod arithmetic;
13mod atomic;
14mod bitwise;
15mod branch;
16mod cmma;
17mod compiler;
18mod debug;
19mod extensions;
20mod globals;
21mod instruction;
22mod item;
23mod lookups;
24mod metadata;
25mod subgroup;
26mod sync;
27mod target;
28mod tensor_indexing;
29mod transformers;
30mod value;
31
32pub use compiler::*;
33use serde::{Deserialize, Serialize};
34pub use target::*;
35
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct SpirvKernel {
38 #[serde(skip)]
39 pub module: Option<Arc<Module>>,
40 #[serde(skip)]
41 pub optimizer: Option<Arc<Optimizer>>,
42
43 pub assembled_module: Vec<u32>,
44 pub bindings: Vec<Visibility>,
45 pub shared_size: usize,
46 pub immediate_size: Option<usize>,
47 pub info_visibility: Visibility,
48}
49
50impl Eq for SpirvKernel {}
51impl PartialEq for SpirvKernel {
52 fn eq(&self, other: &Self) -> bool {
53 self.assembled_module == other.assembled_module
54 }
55}
56
57#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
58pub struct SpirvCacheEntry {
59 pub entrypoint_name: String,
60 pub kernel: SpirvKernel,
61}
62
63impl SpirvCacheEntry {
64 pub fn new(entrypoint_name: String, kernel: SpirvKernel) -> Self {
65 SpirvCacheEntry {
66 entrypoint_name,
67 kernel,
68 }
69 }
70}
71
72impl Display for SpirvKernel {
73 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74 if let Some(module) = &self.module {
75 write!(f, "{}", module.disassemble())
76 } else {
77 f.write_str("SPIR-V")
78 }
79 }
80}