Skip to main content

cubecl_tile/
lib.rs

1//! The axis-agnostic tile DSL engine: a [`Space`] cut into [`Level`]s forms a [`Partitioning`],
2//! whose loops hand out [`Region`]s; the [`Launcher`] reads the grid off the same levels.
3
4// `#[cube]` expands to `cubecl::` paths, which the language crate serves.
5extern crate cubecl_core as cubecl;
6
7pub(crate) mod algebra;
8mod comptime_only;
9pub mod dialect;
10pub mod launch;
11pub mod layout;
12pub mod ops;
13pub mod space;
14pub mod stage;
15pub(crate) mod tile;
16
17/// What a [`Tile`] can be in the kernel.
18pub mod kind {
19    pub use crate::tile::accumulator::planes_output::PlanesOutput;
20    pub use crate::tile::accumulator::smem_accumulation::SmemAccumulation;
21    pub use crate::tile::accumulator::smem_cyclic::SmemCyclicAccumulation;
22    pub use crate::tile::accumulator::smem_slots::SmemSlotsAccumulation;
23    pub use crate::tile::kind::memory::base::{Boundary, Schedule, Storage, Write};
24    pub use crate::tile::kind::memory::global::GlobalOperand;
25    pub use crate::tile::kind::memory::view::decoded::Decoded;
26    pub use crate::tile::kind::memory::view::masked::{Masked, MaskedMut};
27    pub use crate::tile::packing::base::Field;
28    /// A field's codes in and out of words, exposed for the round-trip tests.
29    #[doc(hidden)]
30    pub use crate::tile::packing::view::{pack_line, unpack_line};
31}
32
33/// Tiles whose values are computed from their coordinates.
34pub mod procedural {
35    pub use crate::tile::kind::procedural::affine::{AffineCoordinate, affine_along};
36    pub use crate::tile::kind::procedural::base::{Reads, Recipe, RecipeCoords, RecipeExpand};
37    pub use crate::tile::kind::procedural::constant::Constant;
38    pub use crate::tile::kind::procedural::kind::Procedural;
39    pub use crate::tile::kind::procedural::memory_values::{MemoryValues, MemoryValuesExpand};
40    pub use crate::tile::kind::procedural::normalization::{DivGuard, TapSupport};
41    pub use crate::tile::kind::procedural::phase::Phase;
42    pub use crate::tile::kind::procedural::product::{Product, product_of};
43    pub use crate::tile::kind::procedural::separable::Factors;
44    pub use crate::tile::kind::procedural::sum::{Sum, sum_of};
45}
46
47pub(crate) use comptime_only::comptime_only;
48
49// Flat crate namespace for `use crate::*`.
50#[allow(unused_imports)]
51pub(crate) use {algebra::*, launch::*, layout::*, ops::*, space::*, stage::*, tile::*};
52
53pub use algebra::monoid::{Carrier, Monoid};
54pub use algebra::semiring::Semiring;
55pub use launch::arg::tensor::{TileArg, TileArgLaunch};
56pub use launch::base::Launcher;
57pub use launch::spec::TileSpec;
58pub use layout::projection::Projection;
59pub use layout::storage::Geometry;
60pub use ops::matmul::instruction::{Instruction, MmaOperands};
61pub use ops::matmul::register_block::RegisterBlock;
62pub use space::axis::Axis;
63pub use space::base::Space;
64pub use space::partition::base::Partitioning;
65pub use space::partition::level::Level;
66pub use space::partition::levels::Levels;
67pub use space::region::{Leaves, Planes, Region};
68pub use space::walk::base::Walk;
69pub use stage::base::StageStorage;
70pub use stage::pipeline::base::Stages;
71pub use tile::accumulator::base::Scratch;
72pub use tile::base::{Tile, TileExpand};
73pub use tile::factor::{Factor, FactorExpand};
74pub use tile::index::Indexing;
75pub use tile::slices::{AxisSlices, AxisSlicesExpand};