burn_cubecl/
lib.rs

1#![warn(missing_docs)]
2#![cfg_attr(docsrs, feature(doc_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::{CubeTask, Runtime};
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/// Just-in-Time runtime extending the [cube runtime](Runtime).
43pub trait CubeRuntime: Runtime<Device = Self::CubeDevice, Server = Self::CubeServer> {
44    /// The device that should also implement [burn_backend::backend::DeviceOps].
45    type CubeDevice: burn_backend::DeviceOps;
46    /// The cube server with the [CubeAutotuneKey].
47    type CubeServer: cubecl::server::ComputeServer<Kernel = Box<dyn CubeTask<Self::Compiler>>>;
48}
49
50pub use cubecl::CubeTuneId;