baracuda_kernels/elementwise/
mod.rs

1//! Elementwise op family — unified plan-based API.
2//!
3//! Phase 3 of the baracuda-kernels comprehensive plan: every PyTorch
4//! `torch.<op>` and JAX `jax.numpy.<op>` / `jax.lax.<op>` for elementwise
5//! / shape / layout primitives lives here.
6//!
7//! Today this module hosts the trailblazer dispatcher [`BinaryPlan`] —
8//! one Plan type per op category (unary / binary / ternary / gated /
9//! shape-layout), each parameterized on the scalar element type `T` and
10//! the tensor rank `N`. The op identity within a category is a runtime
11//! enum field on the descriptor (e.g. [`BinaryKind`]), dispatching to
12//! per-(op, dtype) kernel SKUs in `baracuda-kernels-sys` — exactly the
13//! shape used by the GEMM family's [`crate::EpilogueKind`].
14//!
15//! Future modules: `unary` (categories B + B'), `ternary` (category D),
16//! `gated` (category C'), `shape_layout` (category N).
17
18pub mod affine;
19pub mod binary;
20pub mod binary_backward;
21pub mod binary_cmp;
22pub mod binary_param;
23pub mod binary_param_backward;
24pub mod cast;
25pub mod cast_subbyte;
26pub mod gated_activation;
27pub mod gated_activation_backward;
28pub mod prelu;
29pub mod prelu_backward;
30pub mod ternary;
31pub mod ternary_backward;
32pub mod unary;
33pub mod unary_backward;
34pub mod unary_param;
35pub mod unary_param_backward;
36pub mod where_backward;
37pub mod where_op;
38
39pub use affine::{AffineArgs, AffineDescriptor, AffinePlan};
40pub use binary::{BinaryArgs, BinaryDescriptor, BinaryPlan};
41pub use cast::{CastArgs, CastDescriptor, CastPlan};
42pub use cast_subbyte::{CastSubByteArgs, CastSubByteDescriptor, CastSubBytePlan};
43pub use binary_backward::{BinaryBackwardArgs, BinaryBackwardDescriptor, BinaryBackwardPlan};
44pub use binary_cmp::{BinaryCmpArgs, BinaryCmpDescriptor, BinaryCmpPlan};
45pub use binary_param::{BinaryParamArgs, BinaryParamDescriptor, BinaryParamPlan};
46pub use binary_param_backward::{
47    BinaryParamBackwardArgs, BinaryParamBackwardDescriptor, BinaryParamBackwardPlan,
48};
49pub use gated_activation::{
50    GatedActivationArgs, GatedActivationDescriptor, GatedActivationPlan,
51};
52pub use gated_activation_backward::{
53    GatedActivationBackwardArgs, GatedActivationBackwardDescriptor, GatedActivationBackwardPlan,
54};
55pub use prelu::{PReluArgs, PReluDescriptor, PReluPlan};
56pub use prelu_backward::{PReluBackwardArgs, PReluBackwardDescriptor, PReluBackwardPlan};
57pub use ternary::{TernaryArgs, TernaryDescriptor, TernaryPlan};
58pub use ternary_backward::{
59    TernaryBackwardArgs, TernaryBackwardDescriptor, TernaryBackwardPlan,
60};
61pub use unary::{UnaryArgs, UnaryDescriptor, UnaryPlan};
62pub use unary_backward::{UnaryBackwardArgs, UnaryBackwardDescriptor, UnaryBackwardPlan};
63pub use unary_param::{UnaryParamArgs, UnaryParamDescriptor, UnaryParamPlan};
64pub use unary_param_backward::{
65    UnaryParamBackwardArgs, UnaryParamBackwardDescriptor, UnaryParamBackwardPlan,
66};
67pub use where_backward::{WhereBackwardArgs, WhereBackwardDescriptor, WhereBackwardPlan};
68pub use where_op::{WhereArgs, WhereDescriptor, WherePlan};
baracuda_kernels/elementwise/mod.rs

baracuda_kernels/elementwise/
mod.rs
