Skip to main content

OpSpec

atomr_accel_cuda::kernel::cudnn::graph

Enum OpSpec

pub enum OpSpec {
    ConvFwd {
        x: i64,
        w: i64,
        y: i64,
        spatial_dims: usize,
        pre_padding: Vec<i64>,
        post_padding: Vec<i64>,
        stride: Vec<i64>,
        dilation: Vec<i64>,
        compute_dtype: DtypeTag,
        alpha: f64,
        beta: f64,
    },
    ConvBwdData {
        dy: i64,
        w: i64,
        dx: i64,
        spatial_dims: usize,
        pre_padding: Vec<i64>,
        post_padding: Vec<i64>,
        stride: Vec<i64>,
        dilation: Vec<i64>,
        compute_dtype: DtypeTag,
        alpha: f64,
        beta: f64,
    },
    ConvBwdFilter {
        x: i64,
        dy: i64,
        dw: i64,
        spatial_dims: usize,
        pre_padding: Vec<i64>,
        post_padding: Vec<i64>,
        stride: Vec<i64>,
        dilation: Vec<i64>,
        compute_dtype: DtypeTag,
        alpha: f64,
        beta: f64,
    },
    Pointwise {
        mode: PointwiseMode,
        x: i64,
        b: Option<i64>,
        y: i64,
        compute_dtype: DtypeTag,
        alpha1: f64,
        alpha2: f64,
    },
    PoolFwd {
        kind: PoolKind,
        x: i64,
        y: i64,
        window: Vec<i64>,
        pre_padding: Vec<i64>,
        post_padding: Vec<i64>,
        stride: Vec<i64>,
        compute_dtype: DtypeTag,
    },
    PoolBwd {
        kind: PoolKind,
        dy: i64,
        x: i64,
        y: i64,
        dx: i64,
        window: Vec<i64>,
        pre_padding: Vec<i64>,
        post_padding: Vec<i64>,
        stride: Vec<i64>,
        compute_dtype: DtypeTag,
    },
    NormFwd {
        mode: NormMode,
        phase: NormPhase,
        x: i64,
        scale: i64,
        bias: i64,
        mean: Option<i64>,
        var: Option<i64>,
        y: i64,
        compute_dtype: DtypeTag,
        epsilon: f64,
        exp_avg_factor: f64,
    },
    NormBwd {
        mode: NormMode,
        x: i64,
        dy: i64,
        scale: i64,
        mean: i64,
        var: i64,
        dx: i64,
        dscale: i64,
        dbias: i64,
        compute_dtype: DtypeTag,
    },
    Matmul {
        a: i64,
        b: i64,
        c: i64,
        compute_dtype: DtypeTag,
    },
    Reduce {
        op: ReduceOp,
        x: i64,
        y: i64,
        compute_dtype: DtypeTag,
    },
    Reshape {
        x: i64,
        y: i64,
    },
}

Expand description

One op in a backend graph. Each OpSpec references TensorSpecs by their uid; the actual TensorSpec values live on the parent OperationGraphSpec.

Hash is implemented manually so that float fields participate via their bit-pattern (so two specs with alpha = 0.0 hash equal even though f64: !Eq).

Variants§

ConvFwd

Convolution forward: y = conv(x, w).

Fields

§x: i64

§w: i64

§y: i64

§spatial_dims: usize

§pre_padding: Vec<i64>

§post_padding: Vec<i64>

§stride: Vec<i64>

§dilation: Vec<i64>

§compute_dtype: DtypeTag

§alpha: f64

§beta: f64

ConvBwdData

Convolution backward data: dx = conv_bwd_data(w, dy).

Fields

§dy: i64

§w: i64

§dx: i64

§spatial_dims: usize

§pre_padding: Vec<i64>

§post_padding: Vec<i64>

§stride: Vec<i64>

§dilation: Vec<i64>

§compute_dtype: DtypeTag

§alpha: f64

§beta: f64

ConvBwdFilter

Convolution backward filter: dw = conv_bwd_filter(x, dy).

Fields

§x: i64

§dy: i64

§dw: i64

§spatial_dims: usize

§pre_padding: Vec<i64>

§post_padding: Vec<i64>

§stride: Vec<i64>

§dilation: Vec<i64>

§compute_dtype: DtypeTag

§alpha: f64

§beta: f64

Pointwise

Pointwise op (activation, scale, bias-add, …).

Fields

§mode: PointwiseMode

§x: i64

§b: Option<i64>

§y: i64

§compute_dtype: DtypeTag

§alpha1: f64

§alpha2: f64

PoolFwd

Pooling/resample forward.

Fields

§kind: PoolKind

§x: i64

§y: i64

§window: Vec<i64>

§pre_padding: Vec<i64>

§post_padding: Vec<i64>

§stride: Vec<i64>

§compute_dtype: DtypeTag

PoolBwd

Pooling/resample backward.

Fields

§kind: PoolKind

§dy: i64

§x: i64

§y: i64

§dx: i64

§window: Vec<i64>

§pre_padding: Vec<i64>

§post_padding: Vec<i64>

§stride: Vec<i64>

§compute_dtype: DtypeTag

NormFwd

Normalisation forward (batch / layer / instance / group).

Fields

§mode: NormMode

§phase: NormPhase

§x: i64

§scale: i64

§bias: i64

§mean: Option<i64>

§var: Option<i64>

§y: i64

§compute_dtype: DtypeTag

§epsilon: f64

§exp_avg_factor: f64

NormBwd

Normalisation backward.

Fields

§mode: NormMode

§x: i64

§dy: i64

§scale: i64

§mean: i64

§var: i64

§dx: i64

§dscale: i64

§dbias: i64

§compute_dtype: DtypeTag

Matmul

Matmul (2D) — used by attention fusion.

Fields

§a: i64

§b: i64

§c: i64

§compute_dtype: DtypeTag

Reduce

Reduction (sum / max / min / mul / norm).

Fields

§op: ReduceOp

§x: i64

§y: i64

§compute_dtype: DtypeTag

Reshape

Reshape (no-copy view change).

Fields

§x: i64

§y: i64

Trait Implementations§

impl Clone for OpSpec

fn clone(&self) -> OpSpec

Returns a duplicate of the value. Read more

1.0.0 · Source§

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source. Read more

impl Debug for OpSpec

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter. Read more

impl Hash for OpSpec

fn hash<H: Hasher>(&self, h: &mut H)

Feeds this value into the given Hasher. Read more

1.3.0 · Source§

fn hash_slice<H>(data: &[Self], state: &mut H)
where H: Hasher, Self: Sized,

Feeds a slice of this type into the given Hasher. Read more

Auto Trait Implementations§

impl Freeze for OpSpec

impl RefUnwindSafe for OpSpec

impl Send for OpSpec

impl Sync for OpSpec

impl Unpin for OpSpec

impl UnsafeUnpin for OpSpec

impl UnwindSafe for OpSpec

Blanket Implementations§

impl<T> Any for T
where T: 'static + ?Sized,

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more

impl<T> Borrow<T> for T
where T: ?Sized,

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more

impl<T> BorrowMut<T> for T
where T: ?Sized,

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more

impl<T> CloneToUninit for T
where T: Clone,

unsafe fn clone_to_uninit(&self, dest: *mut u8)

🔬This is a nightly-only experimental API. (clone_to_uninit)

Performs copy-assignment from self to dest. Read more

impl<T> From<T> for T

fn from(t: T) -> T

Returns the argument unchanged.

impl<T> Instrument for T

fn instrument(self, span: Span) -> Instrumented<Self>

Instruments this type with the provided [Span], returning an Instrumented wrapper. Read more

fn in_current_span(self) -> Instrumented<Self>

Instruments this type with the current Span, returning an Instrumented wrapper. Read more

impl<T, U> Into<U> for T
where U: From<T>,

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

impl<T> ToOwned for T
where T: Clone,

type Owned = T

The resulting type after obtaining ownership.

fn to_owned(&self) -> T

Creates owned data from borrowed data, usually by cloning. Read more

fn clone_into(&self, target: &mut T)

Uses borrowed data to replace owned data, usually by cloning. Read more

impl<T, U> TryFrom<U> for T
where U: Into<T>,

type Error = Infallible

The type returned in the event of a conversion error.

fn try_from(value: U) -> Result<T, <T as TryFrom<U>>::Error>

Performs the conversion.

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.

impl<T> WithSubscriber for T

fn with_subscriber<S>(self, subscriber: S) -> WithDispatch<Self>
where S: Into<Dispatch>,

Attaches the provided Subscriber to this type, returning a [WithDispatch] wrapper. Read more

fn with_current_subscriber(self) -> WithDispatch<Self>

Attaches the current default Subscriber to this type, returning a [WithDispatch] wrapper. Read more

impl<T> Extension for T
where T: Any + Send + Sync,