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 kNVTEInt64¶
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
-
enumerator kNVTEByte¶
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.
-
void nvte_tensor_pack_create(NVTETensorPack *pack)¶
Create
tensors
in NVTETensorPack.
-
void nvte_tensor_pack_destroy(NVTETensorPack *pack)¶
Destroy
tensors
in NVTETensorPack.
-
struct NVTEShape¶
- #include <transformer_engine.h>
Shape of the tensor.
-
struct NVTETensorPack¶
- #include <transformer_engine.h>
Pack of tensors, generally used for auxiliary outputs.
Public Members
-
NVTETensor tensors[MAX_SIZE]¶
Wrappers of tensors. They do not hold the associated memory.
-
size_t size = 0¶
Actual number of tensors in the pack, 0 <= size <= MAX_SIZE.
Public Static Attributes
-
static const int MAX_SIZE = 10¶
Max number of tensors in the pack. Assumed <= 10.
-
NVTETensor tensors[MAX_SIZE]¶
-
namespace transformer_engine
Namespace containing C++ API of Transformer Engine.
Enums
-
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.
-
inline TensorWrapper(void *dptr, const NVTEShape &shape, const DType dtype, float *amax_dptr = nullptr, float *scale_dptr = nullptr, float *scale_inv_dptr = nullptr)¶
-
struct TensorWrapper¶