Crate cubecl_core 

Re-exports

pub use frontend::cmma;

pub use frontend::wgmma;

pub use prelude::Assign;

pub use prelude::IntoRuntime;

pub use cubecl_zspace as zspace;

pub use cubecl_environment as environment;

pub use cubecl_ir as ir;

pub use num_traits;

pub use codegen::*;

Modules

arena

A circular, allocation-free arena for reusable memory blocks.

benchmark

Module for benchmark timings

bytes

Utilities module to manipulate bytes.

client

Compute client module.

codegen

compute

device

Device module.

device_handle

Device handle module.

format

Format utilities.

frontend

Cube Frontend Types.

future

Future utils with a compatible API for native, non-std and wasm environments.

hash

Hashing helper for stable, collision resistant hashes

io

Input Output utilities.

pool

A dynamically-growing pool that leases exclusive, reusable single-cell items. A dynamically-growing pool of single-cell items handed out as exclusive, non-cloneable handles.

post_processing

prelude

profile

Module for profiling any executable part

quant

Quantization primitives required outside of cubecl-quant

server

Compute server module.

throughput

Throughput related utilities.

tune

Autotune module

work

The work a benchmark or a kernel performs, for scoring against measured peaks.

Macros

__expand_debug_print

Print a formatted message using the target’s debug print facilities. The format string is target specific, but Vulkan and CUDA both use the C++ conventions. WGSL isn’t currently supported.

__expand_seq

comment

Insert a literal comment into the kernel source code.

comptime

Mark the contents of this macro as compile time values, turning off all expansion for this code and using it verbatim

comptime_type

Makes the function return a compile time value Useful in a cube trait to have a part of the trait return comptime values

debug_print

Print a formatted message using the target’s debug print facilities. The format string is target specific, but Vulkan and CUDA both use the C++ conventions. WGSL isn’t currently supported.

define

let define!(T) = elem; defines a type T naming elem to the end of its block. A function re-entered on another element gives T back to its outer call when the inner one returns.

define_scalar

Define a custom type to be used for a comptime scalar type. Useful for cases where generics can’t work.

define_size

Define a custom type to be used for a comptime size. Useful for cases where generics can’t work.

expand_assert

expand_error

gpu_asm

GPU version of asm. Currently parses all the same options, but most are not applicable to GPU assembly architectures. Should validate and give proper errors at some point. Also adds a new register spec: the inferred register specifier (_). This is because the specifier isn’t actually meaningful in PTX and is currently ignored, with constraints being inferred from the value type. The reason it’s still present is because we may want to add an explicit mem specifier to allow fine-grained memory clobbering, or support for other assembly formats that do use different register types.

intrinsic

Mark the contents of this macro as an intrinsic, turning off all expansion for this code and calling it with the scope

obfuscate

Generate a type-erased, inline wrapper for a single value.

seq

size

let size!(N) = size; defines a size N to the end of its block. A function re-entered on another size gives N back to its outer call when the inner one returns.

terminate

Terminate the execution of the kernel for the current unit.

unexpanded

Structs

AuxiliaryMemoryReport

A memory a runtime keeps for its own use, and what it holds.

ComptimeFloat

A finite float usable as a comptime kernel parameter.

CubeDim

The number of units across all 3 axis totalling to the number of working units in a cube.

CubeTuneId

ID used to identify a Just-in-Time environment.

InvalidComptimeFloat

A float value that cannot be used as a ComptimeFloat because it is infinite or NaN.

MemoryPoolReport

A structured snapshot of one memory pool: its shape, its current usage, and the high-water marks a memory plan is derived from.

MemoryPoolsReport

What a memory holds, pool by pool.

MemoryReport

Everything the memory of a MemoryScope holds, stream by stream: the single place memory is read from. A caller that wants the totals asks for the usage.

MemoryUsage

Amount of memory in use by this allocator and statistics on how much memory is reserved and wasted in total.

PrecompiledSource

A hand-written kernel’s own compiled text, standing in for what the compiler would have produced from a KernelDefinition.

Ratio

An exact ratio of two integers, reduced to lowest terms on construction.

StreamMemoryReport

What one stream’s memory holds.

e2m1

A 4-bit floating point type with 2 exponent bits and 1 mantissa bit.

e2m3

A 6-bit floating point type with 2 exponent bits and 3 mantissa bits.

e2m1x2

A 4-bit floating point type with 2 exponent bits and 1 mantissa bit. Packed with two elements per value, to allow for conversion to/from bytes. Care must be taken to ensure the shape is adjusted appropriately.

e3m2

A 6-bit floating point type with 3 exponent bits and 2 mantissa bits.

e4m3

A 8-bit floating point type with 4 exponent bits and 3 mantissa bits.

e5m2

A 8-bit floating point type with 5 exponent bits and 2 mantissa bits.

flex32

A floating point type with relaxed precision, minimum f16, max f32.

tf32

A 19-bit floating point type implementing the tfloat32 format.

ue8m0

An 8-bit unsigned floating point type with 8 exponent bits and no mantissa bits. Used for scaling factors.

Enums

CompilationError

JIT compilation error.

CubeCount

Specifieds the number of cubes to be dispatched for a kernel.

ExecutionMode

The kind of execution to be performed.

MemoryConfiguration

High level configuration of memory management.

MemoryPoolKind

The pool shape a MemoryPoolReport describes, carrying the pool’s effective configuration (after alignment rounding and page-size shrinking).

MemoryScope

Which memory a MemoryReport covers.

VectorizationError

Traits

Compiler

Compiles the representation into its own representation that can be formatted into tokens.

CubeElement

The base element trait for the jit backend.

CubeKernel

Kernel that can be defined

FloatBits

A float type whose bit representation can stand in for PartialEq/Eq/Hash.

ScalarArgType

Functions

calculate_cube_count_elemwise

Calculate the number of cubes required to execute an operation where one cube unit is assigned to one element.

define_kernel

kernel’s definition, or the panic its expansion raised, as a CompilationError.

powi_portable

Rust’s powi returns mathematically incorrect results on Windows for large negative exponents. This is an implementation of the standard exponentiation by squaring algorithm used by compilers on most platforms to implement accurate powi. Should be bit-identical to the one in LLVM’s compiler-rt.

tensor_vector_size_parallel

Find the maximum vector size usable for parallel vectorization along the given axis from the supported vector sizes or return 1 if vectorization is impossible.

tensor_vector_size_perpendicular

Find the maximum vector size usable for perpendicular vectorization along the given axis from the supported vector sizes or return 1 if vectorization is impossible.

tensor_vectorization

tensor_vectorization_factor

try_tensor_vector_size_parallel

Like try_tensor_vector_size_parallel but does not assume 1 is supported

try_tensor_vector_sizes_perpendicular

Like tensor_vector_sizes_perpendicular but does not assume 1 is supported

Type Aliases

ExpandType

RuntimeArg

Runtime arguments to launch a kernel.

Attribute Macros

cube

Mark a cube function, trait or implementation for expansion.

derive_cube_comptime

Attribute macro to define a type that can be used as a kernel comptime argument This derive Debug, Hash, PartialEq, Eq, Clone, Copy

derive_expand

Attribute macro to derive cube traits for existing structs, without redefining that struct.

Derive Macros

AutotuneKey

Implements display and initialization for autotune keys.

CubeHarness

Derive macro to let a cube type cross between the host and a cube function’s harness, as an input or an output. The type must also derive CubeType and CubeLaunch.

CubeLaunch

Derive macro to define a cube type that is launched with a kernel

CubeType

Derive macro to define a cube type that is not launched

CubeTypeMut

Implements mutability for a CubeType

IntoRuntime

Implements IntoRuntime for a CubeType