Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
152 changes: 143 additions & 9 deletions include/matx/core/dlpack.h
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,11 @@
#define DLPACK_EXTERN_C
#endif

/*! \brief The current version of dlpack */
#define DLPACK_VERSION 80
/*! \brief The current major version of dlpack */
#define DLPACK_MAJOR_VERSION 1

/*! \brief The current ABI version of dlpack */
#define DLPACK_ABI_VERSION 1
/*! \brief The current minor version of dlpack */
#define DLPACK_MINOR_VERSION 1

/*! \brief DLPACK_DLL prefix for windows */
#ifdef _WIN32
Expand All @@ -38,6 +38,33 @@
#ifdef __cplusplus
extern "C" {
#endif

/*!
* \brief The DLPack version.
*
* A change in major version indicates that we have changed the
* data layout of the ABI - DLManagedTensorVersioned.
*
* A change in minor version indicates that we have added new
* code, such as a new device type, but the ABI is kept the same.
*
* If an obtained DLPack tensor has a major version that disagrees
* with the version number specified in this header file
* (i.e. major != DLPACK_MAJOR_VERSION), the consumer must call the deleter
* (and it is safe to do so). It is not safe to access any other fields
* as the memory layout will have changed.
*
* In the case of a minor version mismatch, the tensor can be safely used as
* long as the consumer knows how to interpret all fields. Minor version
* updates indicate the addition of enumeration values.
*/
typedef struct {
/*! \brief DLPack major version. */
uint32_t major;
/*! \brief DLPack minor version. */
uint32_t minor;
} DLPackVersion;

/*!
* \brief The device type in DLDevice.
*/
Expand Down Expand Up @@ -89,6 +116,8 @@ typedef enum {
kDLWebGPU = 15,
/*! \brief Qualcomm Hexagon DSP */
kDLHexagon = 16,
/*! \brief Microsoft MAIA devices */
kDLMAIA = 17,
} DLDeviceType;

/*!
Expand Down Expand Up @@ -128,6 +157,26 @@ typedef enum {
kDLComplex = 5U,
/*! \brief boolean */
kDLBool = 6U,
/*! \brief FP8 data types */
kDLFloat8_e3m4 = 7U,
kDLFloat8_e4m3 = 8U,
kDLFloat8_e4m3b11fnuz = 9U,
kDLFloat8_e4m3fn = 10U,
kDLFloat8_e4m3fnuz = 11U,
kDLFloat8_e5m2 = 12U,
kDLFloat8_e5m2fnuz = 13U,
kDLFloat8_e8m0fnu = 14U,
/*! \brief FP6 data types
* Setting bits != 6 is currently unspecified, and the producer must ensure it is set
* while the consumer must stop importing if the value is unexpected.
*/
kDLFloat6_e2m3fn = 15U,
kDLFloat6_e3m2fn = 16U,
/*! \brief FP4 data types
* Setting bits != 4 is currently unspecified, and the producer must ensure it is set
* while the consumer must stop importing if the value is unexpected.
*/
kDLFloat4_e2m1fn = 17U,
} DLDataTypeCode;

/*!
Expand All @@ -141,6 +190,12 @@ typedef enum {
* - int8: type_code = 0, bits = 8, lanes = 1
* - std::complex<float>: type_code = 5, bits = 64, lanes = 1
* - bool: type_code = 6, bits = 8, lanes = 1 (as per common array library convention, the underlying storage size of bool is 8 bits)
* - float8_e4m3: type_code = 8, bits = 8, lanes = 1 (packed in memory)
* - float6_e3m2fn: type_code = 16, bits = 6, lanes = 1 (packed in memory)
* - float4_e2m1fn: type_code = 17, bits = 4, lanes = 1 (packed in memory)
*
* When a sub-byte type is packed, DLPack requires the data to be in little bit-endian, i.e.,
* for a packed data set D ((D >> (i * bits)) && bit_mask) stores the i-th element.
*/
typedef struct {
/*!
Expand Down Expand Up @@ -186,6 +241,9 @@ typedef struct {
* return size;
* }
* \endcode
*
* Note that if the tensor is of size zero, then the data pointer should be
* set to `NULL`.
*/
void* data;
/*! \brief The device of the tensor */
Expand All @@ -211,6 +269,13 @@ typedef struct {
* not meant to transfer the tensor. When the borrowing framework doesn't need
* the tensor, it should call the deleter to notify the host that the resource
* is no longer needed.
*
* \note This data structure is used as Legacy DLManagedTensor
* in DLPack exchange and is deprecated after DLPack v0.8
* Use DLManagedTensorVersioned instead.
* This data structure may get renamed or deleted in future versions.
*
* \sa DLManagedTensorVersioned
*/
typedef struct DLManagedTensor {
/*! \brief DLTensor which is being memory managed */
Expand All @@ -219,14 +284,83 @@ typedef struct DLManagedTensor {
* which DLManagedTensor is used in the framework. It can also be NULL.
*/
void * manager_ctx;
/*! \brief Destructor signature void (*)(void*) - this should be called
* to destruct manager_ctx which holds the DLManagedTensor. It can be NULL
* if there is no way for the caller to provide a reasonable destructor.
* The destructors deletes the argument self as well.
/*!
* \brief Destructor - this should be called
* to destruct the manager_ctx which backs the DLManagedTensor. It can be
* NULL if there is no way for the caller to provide a reasonable destructor.
* The destructor deletes the argument self as well.
*/
void (*deleter)(struct DLManagedTensor * self);
} DLManagedTensor;

// bit masks used in in the DLManagedTensorVersioned

/*! \brief bit mask to indicate that the tensor is read only. */
#define DLPACK_FLAG_BITMASK_READ_ONLY (1UL << 0UL)

/*!
* \brief bit mask to indicate that the tensor is a copy made by the producer.
*
* If set, the tensor is considered solely owned throughout its lifetime by the
* consumer, until the producer-provided deleter is invoked.
*/
#define DLPACK_FLAG_BITMASK_IS_COPIED (1UL << 1UL)

/*
* \brief bit mask to indicate that whether a sub-byte type is packed or padded.
*
* The default for sub-byte types (ex: fp4/fp6) is assumed packed. This flag can
* be set by the producer to signal that a tensor of sub-byte type is padded.
*/
#define DLPACK_FLAG_BITMASK_IS_SUBBYTE_TYPE_PADDED (1UL << 2UL)

/*!
* \brief A versioned and managed C Tensor object, manage memory of DLTensor.
*
* This data structure is intended to facilitate the borrowing of DLTensor by
* another framework. It is not meant to transfer the tensor. When the borrowing
* framework doesn't need the tensor, it should call the deleter to notify the
* host that the resource is no longer needed.
*
* \note This is the current standard DLPack exchange data structure.
*/
struct DLManagedTensorVersioned {
/*!
* \brief The API and ABI version of the current managed Tensor
*/
DLPackVersion version;
/*!
* \brief the context of the original host framework.
*
* Stores DLManagedTensorVersioned is used in the
* framework. It can also be NULL.
*/
void *manager_ctx;
/*!
* \brief Destructor.
*
* This should be called to destruct manager_ctx which holds the DLManagedTensorVersioned.
* It can be NULL if there is no way for the caller to provide a reasonable
* destructor. The destructor deletes the argument self as well.
*/
void (*deleter)(struct DLManagedTensorVersioned *self);
/*!
* \brief Additional bitmask flags information about the tensor.
*
* By default the flags should be set to 0.
*
* \note Future ABI changes should keep everything until this field
* stable, to ensure that deleter can be correctly called.
*
* \sa DLPACK_FLAG_BITMASK_READ_ONLY
* \sa DLPACK_FLAG_BITMASK_IS_COPIED
*/
uint64_t flags;
/*! \brief DLTensor which is being memory managed */
DLTensor dl_tensor;
};

#ifdef __cplusplus
} // DLPACK_EXTERN_C
#endif
#endif // DLPACK_DLPACK_H_
#endif // DLPACK_DLPACK_H_