transformer_engine.h

Base classes and functions of Transformer Engine API.

Typedefs

typedef void *NVTETensor

TE Tensor type.

NVTETensor is a contiguous tensor type storing a pointer to data of a given shape and type. It does not own the memory it points to.

Enums

enum NVTEDType

TE datatype.

Values:

enumerator kNVTEByte

Byte

enumerator kNVTEInt32

32-bit integer

enumerator kNVTEFloat32

32-bit float

enumerator kNVTEFloat16

16-bit float (E5M10)

enumerator kNVTEBFloat16

16-bit bfloat (E8M7)

enumerator kNVTEFloat8E4M3

8-bit float (E4M3)

enumerator kNVTEFloat8E5M2

8-bit float (E5M2)

enumerator kNVTENumTypes

Number of supported types

Functions

NVTETensor nvte_create_tensor(void *dptr, const NVTEShape shape, const NVTEDType dtype, float *amax_dptr, float *scale_dptr, float *scale_inv_dptr)

Create a new TE tensor.

Create a new TE tensor with a given shape, datatype and data. TE tensors are just wrappers on top of raw data and do not own memory.

Parameters:

  • dptr – [in] Pointer to the tensor data.

  • shape – [in] Shape of the tensor.

  • dtype – [in] Data type of the tensor.

  • amax_dptr – [in] Pointer to the AMAX value.

  • scale_dptr – [in] Pointer to the scale value.

  • scale_inv_dptr – [in] Pointer to the inverse of scale value.

Returns:

A new TE tensor.

void nvte_destroy_tensor(NVTETensor tensor)

Destroy a TE tensor.

Since the TE tensor does not own memory, the underlying data is not freed during this operation.

Parameters:

tensor – [in] Tensor to be destroyed.

NVTEDType nvte_tensor_type(const NVTETensor tensor)

Get a tensor’s data type.

Parameters:

tensor – [in] Tensor.

Returns:

A data type of the input tensor.

NVTEShape nvte_tensor_shape(const NVTETensor tensor)

Get a tensor’s data shape.

Parameters:

tensor – [in] Tensor.

Returns:

A shape of the input tensor.

void *nvte_tensor_data(const NVTETensor tensor)

Get a raw pointer to the tensor’s data.

Parameters:

tensor – [in] Tensor.

Returns:

A raw pointer to tensor’s data.

float *nvte_tensor_amax(const NVTETensor tensor)

Get a pointer to the tensor’s amax data.

Parameters:

tensor – [in] Tensor.

Returns:

A pointer to tensor’s amax data.

float *nvte_tensor_scale(const NVTETensor tensor)

Get a pointer to the tensor’s scale data.

Parameters:

tensor – [in] Tensor.

Returns:

A pointer to tensor’s scale data.

float *nvte_tensor_scale_inv(const NVTETensor tensor)

Get a pointer to the tensor’s inverse of scale data.

Parameters:

tensor – [in] Tensor.

Returns:

A pointer to tensor’s inverse of scale data.

struct NVTEShape
#include <transformer_engine.h>

Shape of the tensor.

Public Members

const size_t *data

Shape data, of size ndim.

size_t ndim

Number of dimensions.

namespace transformer_engine

Namespace containing C++ API of Transformer Engine.

Enums

enum class DType

TE datatype.

Values:

enumerator kByte
enumerator kInt32
enumerator kFloat32
enumerator kFloat16
enumerator kBFloat16
enumerator kFloat8E4M3
enumerator kFloat8E5M2
enumerator kNumTypes
struct TensorWrapper
#include <transformer_engine.h>

C++ wrapper for the NVTETensor class.

Public Functions

inline TensorWrapper(void *dptr, const NVTEShape &shape, const DType dtype, float *amax_dptr = nullptr, float *scale_dptr = nullptr, float *scale_inv_dptr = nullptr)

Constructs new TensorWrapper.

Create a new TE tensor with a given shape, datatype and data. TE tensors are just wrappers on top of raw data and do not own memory.

Parameters:

  • dptr – [in] Pointer to the tensor data.

  • shape – [in] Shape of the tensor.

  • dtype – [in] Data type of the tensor.

  • amax_dptr – [in] Pointer to the AMAX value.

  • scale_dptr – [in] Pointer to the scale value.

  • scale_inv_dptr – [in] Pointer to the inverse of scale value.

inline TensorWrapper(void *dptr, const std::vector<size_t> &shape, const DType dtype, float *amax_dptr = nullptr, float *scale_dptr = nullptr, float *scale_inv_dptr = nullptr)

Constructs new TensorWrapper.

Create a new TE tensor with a given shape, datatype and data. TE tensors are just wrappers on top of raw data and do not own memory.

Parameters:

  • dptr – [in] Pointer to the tensor data.

  • shape – [in] Shape of the tensor.

  • dtype – [in] Data type of the tensor.

  • amax_dptr – [in] Pointer to the AMAX value.

  • scale_dptr – [in] Pointer to the scale value.

  • scale_inv_dptr – [in] Pointer to the inverse of scale value.

inline TensorWrapper()

Constructs new empty TensorWrapper.

Create a new empty TE tensor which holds nothing.

inline ~TensorWrapper()

TensorWrapper destructor.

TensorWrapper &operator=(const TensorWrapper &other) = delete
TensorWrapper(const TensorWrapper &other) = delete
inline TensorWrapper(TensorWrapper &&other)

Constructs new TensorWrapper from existing TensorWrapper.

Pass an existing TE tensor to a new TensorWrapper.

Parameters:

other – [inout] The source of the data.

inline TensorWrapper &operator=(TensorWrapper &&other)

Assign the data from existing TensorWrapper.

Change ownership of an existing TE tensor.

Parameters:

other – [inout] The source of the data.

inline NVTETensor data() const noexcept

Get an underlying NVTETensor.

Returns:

NVTETensor held by this TensorWrapper.

inline const NVTEShape shape() const noexcept

Get the shape of this TensorWrapper.

Returns:

Shape of this TensorWrapper.

inline DType dtype() const noexcept

Get the data type of this TensorWrapper.

Returns:

Data type of this TensorWrapper.

inline void *dptr() const noexcept

Get a raw pointer to the tensor’s data.

Returns:

A raw pointer to tensor’s data.

inline float *amax() const noexcept

Get a pointer to the tensor’s amax data.

Returns:

A pointer to tensor’s amax data.

inline float *scale() const noexcept

Get a pointer to the tensor’s scale data.

Returns:

A pointer to tensor’s scale data.

inline float *scale_inv() const noexcept

Get a pointer to the tensor’s inverse of scale data.

Returns:

A pointer to tensor’s inverse of scale data.

Private Members

NVTETensor tensor_ = nullptr

Wrapped NVTETensor.