28 #ifndef MXNET_TENSOR_BLOB_H_ 29 #define MXNET_TENSOR_BLOB_H_ 31 #include <dmlc/logging.h> 89 template<
typename DType>
91 : dptr_(dptr), shape_(shape),
92 type_flag_(
mshadow::DataType<DType>::kFlag) {
93 SetDLTensor(dev_mask,
dev_id);
104 : dptr_(dptr), shape_(shape), type_flag_(type_flag) {
105 SetDLTensor(dev_mask,
dev_id);
112 : dptr_(dltensor.data),
114 type_flag_(DLDataTypeTransform(dltensor.dtype)),
115 dltensor_(dltensor) {
117 if (dltensor.
strides !=
nullptr) {
120 const int64_t *shape = dltensor.
shape;
121 const int64_t *strides = dltensor.
strides;
124 if (strides[ndim - 1] != 1) {
127 for (
int i = ndim - 2; i >= 0; --i) {
128 if (strides[i] != shape[i + 1] * strides[i + 1]) {
135 LOG(FATAL) <<
"Unsupported DLPack because MXNet only support compact tensor now";
147 template<
typename Device,
int dim,
typename DType>
155 TBlob(
const TBlob &src): dptr_(src.dptr_), shape_(src.shape_), type_flag_(src.type_flag_) {
166 template<
typename Device,
int dim,
typename DType>
171 SetDLTensor(Device::kDevMask, -1);
198 CHECK_EQ(this->shape_.
Size(), shape.
Size()) <<
"Shape size mismatch " 199 << this->shape_.
Size() <<
" v.s. " << shape.
Size();
210 template<
typename Device,
typename DType>
213 CHECK(Device::kDevMask == this->
dev_mask())
214 <<
"TBlob.get: device type do not match specified type";
216 <<
"TBlob.get_with_shape: data type do not match specified type." 229 template<
typename Device,
typename DType>
232 return this->get_with_shape<Device, 1, DType>(
237 return shape_.
ndim();
249 inline size_t Size(
void)
const {
250 return shape_.
Size();
253 template<
typename DType>
256 <<
"TBlob.get_with_shape: data type do not match specified type." 258 return static_cast<DType*
>(
dptr_);
285 template<
typename Device,
int dim,
typename DType>
287 CHECK(Device::kDevMask == this->
dev_mask())
288 <<
"TBlob.get: device type do not match specified type";
290 shape_.get<dim>(), shape_[shape_.
ndim() - 1], stream);
302 template<
typename Device,
int dim,
typename DType>
306 CHECK(Device::kDevMask == this->
dev_mask())
307 <<
"TBlob.get: device type do not match specified type";
308 CHECK_EQ(this->
CheckContiguous(),
true) <<
"TBlob.get_reshape: must be contiguous";
309 CHECK_EQ(this->shape_.
Size(),
static_cast<size_t>(shape.
Size()))
310 <<
"TBlob.get_with_shape: new and old shape do not match total elements";
312 shape[dim - 1], stream);
323 template<
typename Device,
typename DType>
326 return this->get_with_shape<Device, 3, DType>(
327 this->shape_.FlatTo3D(axis), stream);
339 template<
typename Device,
typename DType>
341 int axis_begin,
int axis_end,
343 return this->get_with_shape<Device, 3, DType>(
344 this->shape_.FlatTo3D(axis_begin, axis_end), stream);
355 template<
typename Device,
int dim,
typename DType>
361 for (
int i = 0; i < dim -
ndim(); ++i) {
365 for (
int i = 0; i <
ndim() - dim + 1; ++i) {
366 shape[0] *= shape_[i];
369 for (
int i = std::max(0,
ndim() - dim + 1); i <
ndim(); ++i) {
370 shape[i -
ndim() + dim] = shape_[i];
372 return this->get_with_shape<Device, dim, DType>(shape, stream);
376 static DLDataType DTypeTransform(
int type_flag) {
387 LOG(FATAL) <<
"Unknown type_flag=" << type_flag;
392 static int DLDataTypeTransform(
DLDataType dldata_type) {
393 if (dldata_type.
lanes != 1) {
394 LOG(FATAL) <<
"Unsupported DLDataType whose lanes != 1";
396 switch (dldata_type.
code) {
398 switch (dldata_type.
bits) {
405 switch (dldata_type.
bits) {
410 switch (dldata_type.
bits) {
417 LOG(FATAL) <<
"Unknown DLDataType{" << dldata_type.
code 418 <<
", " << dldata_type.
bits 419 <<
", " << dldata_type.
lanes <<
"}";
427 dltensor_.
dtype = DTypeTransform(type_flag_);
447 namespace parameter {
451 :
public FieldEntryBase<FieldEntry<mxnet::TShape>, mxnet::TShape> {
457 virtual void Check(
void *head)
const {
460 if (expect_ndim_ != 0 && v.
ndim() != expect_ndim_) {
461 std::ostringstream os;
462 os <<
"value " << v <<
"for Parameter " << this->key_
463 <<
" has wrong dimensions, expected dimension=" << expect_ndim_;
464 throw dmlc::ParamError(os.str());
466 if (enforce_nonzero_) {
467 for (
int i = 0; i < v.
ndim(); ++i) {
469 std::ostringstream os;
470 os <<
"value " << v <<
"for Parameter " << this->key_
471 <<
" is invalid, the input shape must be nonzero in all dimensions";
472 throw dmlc::ParamError(os.str());
478 this->enforce_nonzero_ =
true;
488 bool enforce_nonzero_;
496 #endif // MXNET_TENSOR_BLOB_H_ #define DMLC_DECLARE_TYPE_NAME(Type, Name)
macro to quickly declare traits information
Definition: type_traits.h:133
TBlob & operator=(const mshadow::Tensor< Device, dim, DType > &src)
assignment from tensor
Definition: tensor_blob.h:167
The common header of DLPack.
MSHADOW_XINLINE index_t Size(void) const
Definition: tensor.h:126
mxnet::TShape shape_
shape of the tensor
Definition: tensor_blob.h:72
constexpr const int kTVMNDArrayTypeCode
Definition: tensor_blob.h:49
DType * dptr_
pointer to the data
Definition: tensor.h:416
TBlob(const DLTensor &dltensor)
constructor that construct TBlob from DLTensor
Definition: tensor_blob.h:111
namespace of mxnet
Definition: base.h:89
A Device context for Tensor and operator.
Definition: dlpack.h:69
Shape< dimension > shape_
shape of the tensor
Definition: tensor.h:418
int64_t * strides
strides of the tensor (in number of elements, not bytes) can be NULL, indicating tensor is compact an...
Definition: dlpack.h:144
mshadow::default_real_t real_t
data type that will be used to store ndarray
Definition: base.h:97
TBlob(void)
default constructor, default copy assign will work
Definition: tensor_blob.h:77
Definition: tensor_blob.h:450
int type_flag_
type flag of the tensor blob
Definition: tensor_blob.h:74
mshadow::Tensor< Device, dim, DType > get_with_shape(const mshadow::Shape< dim > &shape, mshadow::Stream< Device > *stream=NULL) const
fetch a tensor in given shape If size do not match the stored size, an error will be issued ...
Definition: tensor_blob.h:303
FieldEntry< mxnet::TShape > & set_expect_ndim(int ndim)
Definition: tensor_blob.h:481
FieldEntry< mxnet::TShape > & enforce_nonzero()
Definition: tensor_blob.h:477
FieldEntryBase< FieldEntry< mxnet::TShape >, mxnet::TShape > Parent
Definition: tensor_blob.h:455
mshadow::Tensor< Device, 1, DType > FlatTo1D(mshadow::Stream< Device > *stream=NULL) const
flatten the tensor to 1 dimension, collapse all the dimensions together.
Definition: tensor_blob.h:230
const dim_t * data() const
Definition: tuple.h:523
TBlob(void *dptr, const mxnet::TShape &shape, int dev_mask, int type_flag, int dev_id=-1)
constructor that construct TBlob from contiguous memory
Definition: tensor_blob.h:103
uint8_t code
Type code of base types. We keep it uint8_t instead of DLDataTypeCode for minimal memory footprint...
Definition: dlpack.h:99
int device_id
The device index.
Definition: dlpack.h:73
constexpr const int kGPU
Definition: tensor_blob.h:44
Lightweight JSON Reader/Writer that read save into C++ data structs. This includes STL composites and...
CPU device.
Definition: dlpack.h:40
index_t size(index_t idx) const
return size of i-th dimension, start counting from highest dimension. return type needs to be a signe...
Definition: tensor_blob.h:245
size_t Size() const
Definition: tuple.h:494
void * dptr_
pointer to the data
Definition: tensor_blob.h:70
int64_t * shape
The shape of the tensor.
Definition: dlpack.h:139
namespace for dmlc
Definition: array_view.h:12
uint8_t bits
Number of bits, common choices are 8, 16, 32.
Definition: dlpack.h:103
void * data
The opaque data pointer points to the allocated data. This will be CUDA device pointer or cl_mem hand...
Definition: dlpack.h:131
DLDataType dtype
The data type of the pointer.
Definition: dlpack.h:137
DType * dptr() const
get pointer in dtype
Definition: tensor_blob.h:254
DLDeviceType
The device type in DLContext.
Definition: dlpack.h:38
int ndim(void) const
return number of dimension of the tensor inside
Definition: tensor_blob.h:236
DLDeviceType device_type
The device type used in the device.
Definition: dlpack.h:71
TBlob reshape(const mxnet::TShape &shape) const
reshape to shape
Definition: tensor_blob.h:197
MSHADOW_XINLINE Shape< 1 > Shape1(index_t s0)
construct a one dimension shape, stride will equal s0
Definition: tensor.h:188
TBlob(const mshadow::Tensor< Device, dim, DType > &src)
constructor from tensor
Definition: tensor_blob.h:148
A dynamic sized array data structure that is optimized for storing small number of elements with same...
Definition: tuple.h:54
TBlob(DType *dptr, const mxnet::TShape &shape, int dev_mask, int dev_id=-1)
constructor that construct TBlob from contiguous memory
Definition: tensor_blob.h:90
int ndim
Number of dimensions.
Definition: dlpack.h:135
TBlob & operator=(const TBlob &src)
assignment from TBlob (copy assignment)
Definition: tensor_blob.h:179
static const int kDevMask
device flag number, identifies this device
Definition: tensor.h:25
const DLTensor & dltensor() const
return the corresponding DLTensor
Definition: tensor_blob.h:272
A Shape class that is used to represent shape of each tensor.
Definition: tuple.h:413
mshadow::Tensor< Device, 3, DType > FlatTo3D(int axis_begin, int axis_end, mshadow::Stream< Device > *stream=NULL) const
flatten the tensor to 3 dimension, collapse the dimension: [0, axis_begin), [axis_begin, axis_end], (axis_end, ndim).
Definition: tensor_blob.h:340
mshadow::Tensor< Device, dim, DType > FlatToKD(mshadow::Stream< Device > *stream=NULL) const
flatten the tensor to specified number of dimensions, collapse the highest dimensions or pad with hig...
Definition: tensor_blob.h:356
mshadow::Tensor< Device, 2, DType > FlatTo2D(mshadow::Stream< Device > *stream=NULL) const
flatten the tensor to 2 dimension, collapse the higher dimensions together
Definition: tensor_blob.h:211
virtual void Check(void *head) const
Definition: tensor_blob.h:457
int ndim() const
Definition: tuple.h:193
bool CheckContiguous(void) const
Definition: tensor_blob.h:189
namespace for mshadow
Definition: base.h:282
mshadow::Tensor< Device, 3, DType > FlatTo3D(int axis, mshadow::Stream< Device > *stream=NULL) const
flatten the tensor to 3 dimension, collapse the dimension before and after specified axis...
Definition: tensor_blob.h:324
mshadow::index_t index_t
index type usually use unsigned
Definition: base.h:95
constexpr const int kCPU
Definition: tensor_blob.h:43
DLContext ctx
The device context of the tensor.
Definition: dlpack.h:133
CUDA GPU device.
Definition: dlpack.h:42
The data type the tensor can hold.
Definition: dlpack.h:93
general tensor
Definition: tensor.h:402
uint16_t lanes
Number of lanes in the type, used for vector types.
Definition: dlpack.h:105
Plain C Tensor object, does not manage memory.
Definition: dlpack.h:111
FieldEntry()
Definition: tensor_blob.h:453
ndarray interface
Definition: ndarray.h:82
uint64_t byte_offset
The offset in bytes to the beginning pointer to data.
Definition: dlpack.h:146
int dev_mask() const
device mask of the corresponding device
Definition: tensor_blob.h:261
TBlob(const TBlob &src)
constructor from TBlob (copy constructor)
Definition: tensor_blob.h:155
tensor blob class that can be used to hold tensor of any dimension, any device and any data type...
Definition: tensor_blob.h:66
int dev_id() const
device index of the corresponding device
Definition: tensor_blob.h:265
computaion stream structure, used for asynchronous computations
Definition: tensor.h:365
size_t Size(void) const
total number of elements in the tensor
Definition: tensor_blob.h:249