burn_cubecl/
lib.rs

1#![warn(missing_docs)]
2#![cfg_attr(docsrs, feature(doc_auto_cfg))]
3
4//! Burn JIT Backend
5
6#[macro_use]
7extern crate derive_new;
8extern crate alloc;
9
10/// Utilities for implementing JIT kernels
11pub mod ops;
12
13/// Kernel module
14pub mod kernel;
15/// Tensor module.
16pub mod tensor;
17
18/// Elements for JIT backend
19pub mod element;
20
21use cubecl::{Feature, Runtime, compute::CubeTask};
22pub use element::{BoolElement, CubeElement, FloatElement, IntElement};
23
24mod backend;
25
26pub use backend::*;
27
28// Re-export cubecl.
29pub use cubecl;
30
31mod tune_key;
32pub use tune_key::CubeAutotuneKey;
33
34#[cfg(any(feature = "fusion", test))]
35/// Module for interacting with fusion
36pub mod fusion;
37
38#[cfg(feature = "template")]
39/// Module for compiling custom non-jit kernels
40pub mod template;
41
42#[cfg(feature = "export_tests")]
43pub mod tests;
44
45/// Just-in-Time runtime extending the [cube runtime](Runtime).
46pub trait CubeRuntime: Runtime<Device = Self::CubeDevice, Server = Self::CubeServer> {
47    /// The device that should also implement [burn_tensor::backend::DeviceOps].
48    type CubeDevice: burn_tensor::backend::DeviceOps;
49    /// The cube server with the [CubeAutotuneKey].
50    type CubeServer: cubecl::server::ComputeServer<Kernel = Box<dyn CubeTask<Self::Compiler>>, Feature = Feature>;
51}
52
53pub use cubecl::CubeTuneId;
burn_cubecl/lib.rs

burn_cubecl/
lib.rs
