Shuffle#

The shuffle primitive shuffles data along the shuffle axis (here is designated as \(C\)) with the group parameter \(G\). Namely, the shuffle axis is thought to be a 2D tensor of size \((\frac{C}{G} \times G)\) and it is being transposed to \((G \times \frac{C}{G})\). Variable names follow the standard Conventions.

The formal definition is shown below:

Forward#

\[\dst(\overline{ou}, c, \overline{in}) = \src(\overline{ou}, c', \overline{in})\]

where

  • \(c\) dimension is called a shuffle axis,

  • \(G\) is a group_size,

  • \(\overline{ou}\) is the outermost indices (to the left from shuffle axis),

  • \(\overline{in}\) is the innermost indices (to the right from shuffle axis), and

  • \(c'\) and \(c\) relate to each other as define by the system:

\[\begin{split}\begin{cases} c &= u + v \cdot \frac{C}{G}, \\ c' &= u \cdot G + v, \\ \end{cases}\end{split}\]

Here, \(0 \leq u < \frac{C}{G}\) and \(0 \leq v < G\).

Difference Between Forward Training and Forward Inference#

There is no difference between the forward_training and forward_inference propagation kinds.

Backward#

The backward propagation computes \(\diffsrc(ou, c, in)\), based on \(\diffdst(ou, c, in)\).

Essentially, backward propagation is the same as forward propagation with \(g\) replaced by \(C / g\).

Execution Arguments#

When executed, the inputs and outputs should be mapped to an execution argument index as specified by the following table.

Primitive input/output

Execution argument index

\(\src\)

DNNL_ARG_SRC

\(\dst\)

DNNL_ARG_DST

\(\diffsrc\)

DNNL_ARG_DIFF_SRC

\(\diffdst\)

DNNL_ARG_DIFF_DST

Operation Details#

Data Types Support#

The shuffle primitive supports the following combinations of data types:

Note

Here we abbreviate data types names for readability. For example, dnnl::memory::data_type::f32 is abbreviated to f32.

Propagation

Source / Destination

forward / backward

f32, bf16

forward

s32, s8, u8

Data Layouts#

The shuffle primitive works with arbitrary data tensors. There is no special meaning associated with any logical dimensions. However, the shuffle axis is typically referred to as channels (hence in formulas we use \(c\)).

Shuffle operation typically appear in CNN topologies. Hence, in the library the shuffle primitive is optimized for the corresponding memory formats:

Spatial

Logical tensor

Shuffle Axis

Implementations optimized for memory formats

2D

NCHW

1 (C)

nchw (abcd), nhwc (acdb), optimized^

3D

NCDHW

1 (C)

ncdhw (abcde), ndhwc (acdeb), optimized^

Here optimized^ means the format that comes out of any preceding compute-intensive primitive.

Post-ops and Attributes#

The shuffle primitive does not have to support any post-ops or attributes.

API#

struct shuffle_forward : public dnnl::primitive#

Shuffle forward propagation primitive.

Public Functions

shuffle_forward()#

Default constructor. Produces an empty object.

shuffle_forward(const primitive_desc &pd)#

Constructs a shuffle forward propagation primitive.

Parameters:

pd – Primitive descriptor for a shuffle forward propagation primitive.

struct primitive_desc : public dnnl::primitive_desc#

Primitive descriptor for a shuffle forward propagation primitive.

Public Functions

primitive_desc() = default#

Default constructor. Produces an empty object.

primitive_desc(const engine &aengine, prop_kind aprop_kind, const memory::desc &src_desc, const memory::desc &dst_desc, int axis, int group_size, const primitive_attr &attr = default_attr(), bool allow_empty = false)#

Constructs a primitive descriptor for a shuffle forward propagation primitive.

Parameters:

  • aengine – Engine to use.

  • aprop_kind – Propagation kind. Possible values are dnnl::prop_kind::forward_training, and dnnl::prop_kind::forward_inference.

  • src_desc – Source memory descriptor.

  • dst_desc – Destination memory descriptor.

  • axis – The axis along which the data is shuffled.

  • group_size – Shuffle group size.

  • attr – Primitive attributes to use. Attributes are optional and default to empty attributes.

  • allow_empty – A flag signifying whether construction is allowed to fail without throwing an exception. In this case an empty object will be produced. This flag is optional and defaults to false.

memory::desc src_desc() const#

Returns a source memory descriptor.

Returns:

Source memory descriptor.

Returns:

A zero memory descriptor if the primitive does not have a source parameter.

memory::desc dst_desc() const#

Returns a destination memory descriptor.

Returns:

Destination memory descriptor.

Returns:

A zero memory descriptor if the primitive does not have a destination parameter.

prop_kind get_prop_kind() const#

Returns a propagation kind.

Returns:

A propagation kind.

Returns:

dnnl::prop_kind::undef if the primitive does not have a propagation parameter.

int get_axis() const#

Returns an axis.

Returns:

An axis.

Returns:

A negative number if the primitive does not have an axis parameter.

memory::dim get_group_size() const#

Returns a shuffle group size parameter.

Returns:

A shuffle group size parameter.

Returns:

Zero if the primitive does not have a shuffle group size parameter.

struct shuffle_backward : public dnnl::primitive#

Shuffle backward propagation primitive.

Public Functions

shuffle_backward()#

Default constructor. Produces an empty object.

shuffle_backward(const primitive_desc &pd)#

Constructs a shuffle backward propagation primitive.

Parameters:

pd – Primitive descriptor for a shuffle backward propagation primitive.

struct primitive_desc : public dnnl::primitive_desc#

Primitive descriptor for a shuffle backward propagation primitive.

Public Functions

primitive_desc() = default#

Default constructor. Produces an empty object.

primitive_desc(const engine &aengine, const memory::desc &diff_src_desc, const memory::desc &diff_dst_desc, int axis, int group_size, const shuffle_forward::primitive_desc &hint_fwd_pd, const primitive_attr &attr = default_attr(), bool allow_empty = false)#

Constructs a primitive descriptor for a shuffle backward propagation primitive.

Parameters:

  • aengine – Engine to use.

  • diff_src_desc – Diff source memory descriptor.

  • diff_dst_desc – Diff destination memory descriptor.

  • axis – The axis along which the data is shuffled.

  • group_size – Shuffle group size.

  • hint_fwd_pd – Primitive descriptor for a shuffle forward propagation primitive. It is used as a hint for deciding which memory format to use.

  • attr – Primitive attributes to use. Attributes are optional and default to empty attributes.

  • allow_empty – A flag signifying whether construction is allowed to fail without throwing an exception. In this case an empty object will be produced. This flag is optional and defaults to false.

memory::desc diff_src_desc() const#

Returns a diff source memory descriptor.

Returns:

Diff source memory descriptor.

Returns:

A zero memory descriptor if the primitive does not have a diff source memory with.

memory::desc diff_dst_desc() const#

Returns a diff destination memory descriptor.

Returns:

Diff destination memory descriptor.

Returns:

A zero memory descriptor if the primitive does not have a diff destination parameter.

prop_kind get_prop_kind() const#

Returns a propagation kind.

Returns:

A propagation kind.

Returns:

dnnl::prop_kind::undef if the primitive does not have a propagation parameter.

int get_axis() const#

Returns an axis.

Returns:

An axis.

Returns:

A negative number if the primitive does not have an axis parameter.

memory::dim get_group_size() const#

Returns a shuffle group size parameter.

Returns:

A shuffle group size parameter.

Returns:

Zero if the primitive does not have a shuffle group size parameter.