mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
cf19c3697d
### Description
Run clang-format in CI. Formatted all c/c++, objective-c/c++ files.
Excluded
```
'onnxruntime/core/mlas/**',
'onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/**',
```
because they contain assembly or is data heavy
### Motivation and Context
Coding style consistency
265 lines
9.1 KiB
C++
265 lines
9.1 KiB
C++
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
#include "core/dlpack/dlpack_converter.h"
|
|
|
|
namespace onnxruntime {
|
|
namespace dlpack {
|
|
|
|
namespace {
|
|
|
|
DLDataType GetDlpackDataType(const OrtValue& ort_value) {
|
|
ORT_ENFORCE(ort_value.IsTensor(), "Only tensor-type OrtValues are supported");
|
|
DLDataType dtype;
|
|
dtype.lanes = 1;
|
|
const Tensor& tensor = ort_value.Get<Tensor>();
|
|
switch (tensor.GetElementType()) {
|
|
case ONNX_NAMESPACE::TensorProto_DataType_DOUBLE:
|
|
dtype.code = DLDataTypeCode::kDLFloat;
|
|
dtype.bits = sizeof(double);
|
|
break;
|
|
case ONNX_NAMESPACE::TensorProto_DataType_FLOAT:
|
|
dtype.code = DLDataTypeCode::kDLFloat;
|
|
dtype.bits = sizeof(float);
|
|
break;
|
|
case ONNX_NAMESPACE::TensorProto_DataType_INT8:
|
|
dtype.code = DLDataTypeCode::kDLInt;
|
|
dtype.bits = sizeof(int8_t);
|
|
break;
|
|
case ONNX_NAMESPACE::TensorProto_DataType_INT16:
|
|
dtype.code = DLDataTypeCode::kDLInt;
|
|
dtype.bits = sizeof(int16_t);
|
|
break;
|
|
case ONNX_NAMESPACE::TensorProto_DataType_INT32:
|
|
dtype.code = DLDataTypeCode::kDLInt;
|
|
dtype.bits = sizeof(int);
|
|
break;
|
|
case ONNX_NAMESPACE::TensorProto_DataType_INT64:
|
|
dtype.code = DLDataTypeCode::kDLInt;
|
|
dtype.bits = sizeof(int64_t);
|
|
break;
|
|
case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16:
|
|
dtype.code = DLDataTypeCode::kDLFloat;
|
|
dtype.bits = sizeof(MLFloat16);
|
|
break;
|
|
// Currently bool is same as uint8 on both code and bits.
|
|
// PyTorch's to_dlpack also does this, but in from_dlpack,
|
|
// a torch.uint8 tensor is generated. This limitation from
|
|
// PyTorch means we cannot create a torch.bool tensor
|
|
// from a DLPack tensor.
|
|
case ONNX_NAMESPACE::TensorProto_DataType_BOOL:
|
|
dtype.code = DLDataTypeCode::kDLUInt;
|
|
dtype.bits = sizeof(bool);
|
|
break;
|
|
case ONNX_NAMESPACE::TensorProto_DataType_UINT8:
|
|
dtype.code = DLDataTypeCode::kDLUInt;
|
|
dtype.bits = sizeof(uint8_t);
|
|
break;
|
|
case ONNX_NAMESPACE::TensorProto_DataType_UINT16:
|
|
dtype.code = DLDataTypeCode::kDLUInt;
|
|
dtype.bits = sizeof(uint16_t);
|
|
break;
|
|
case ONNX_NAMESPACE::TensorProto_DataType_UINT32:
|
|
dtype.code = DLDataTypeCode::kDLUInt;
|
|
dtype.bits = sizeof(uint32_t);
|
|
break;
|
|
case ONNX_NAMESPACE::TensorProto_DataType_UINT64:
|
|
dtype.code = DLDataTypeCode::kDLUInt;
|
|
dtype.bits = sizeof(uint64_t);
|
|
break;
|
|
case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16:
|
|
dtype.code = DLDataTypeCode::kDLBfloat;
|
|
dtype.bits = sizeof(BFloat16);
|
|
break;
|
|
default:
|
|
ORT_THROW("Unexpected data type of ", tensor.GetElementType());
|
|
}
|
|
|
|
dtype.bits *= 8; // bits.
|
|
return dtype;
|
|
}
|
|
|
|
OrtDevice GetOrtDevice(const DLDevice& device) {
|
|
switch (device.device_type) {
|
|
case DLDeviceType::kDLCPU:
|
|
return OrtDevice();
|
|
case DLDeviceType::kDLCUDA:
|
|
case DLDeviceType::kDLROCM:
|
|
return OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, static_cast<OrtDevice::DeviceId>(device.device_id));
|
|
default:
|
|
ORT_THROW("Unsupported device type");
|
|
}
|
|
}
|
|
|
|
MLDataType GetOrtValueDataType(const DLDataType& dtype, bool is_bool_tensor) {
|
|
if (dtype.lanes != 1) ORT_THROW("ORT does not support lanes != 1");
|
|
switch (dtype.code) {
|
|
case DLDataTypeCode::kDLUInt:
|
|
switch (dtype.bits) {
|
|
case 8:
|
|
return is_bool_tensor ? DataTypeImpl::GetType<bool>() : DataTypeImpl::GetType<uint8_t>();
|
|
case 16:
|
|
return DataTypeImpl::GetType<uint16_t>();
|
|
case 32:
|
|
return DataTypeImpl::GetType<uint32_t>();
|
|
case 64:
|
|
return DataTypeImpl::GetType<uint64_t>();
|
|
default:
|
|
ORT_THROW("Unsupported kUInt bits " + std::to_string(dtype.bits));
|
|
}
|
|
case DLDataTypeCode::kDLInt:
|
|
switch (dtype.bits) {
|
|
case 8:
|
|
return DataTypeImpl::GetType<int8_t>();
|
|
case 16:
|
|
return DataTypeImpl::GetType<int16_t>();
|
|
case 32:
|
|
return DataTypeImpl::GetType<int32_t>();
|
|
case 64:
|
|
return DataTypeImpl::GetType<int64_t>();
|
|
default:
|
|
ORT_THROW("Unsupported kInt bits " + std::to_string(dtype.bits));
|
|
}
|
|
case DLDataTypeCode::kDLFloat:
|
|
switch (dtype.bits) {
|
|
case 16:
|
|
return DataTypeImpl::GetType<MLFloat16>();
|
|
case 32:
|
|
return DataTypeImpl::GetType<float>();
|
|
case 64:
|
|
return DataTypeImpl::GetType<double>();
|
|
default:
|
|
ORT_THROW("Unsupported kFloat bits " + std::to_string(dtype.bits));
|
|
}
|
|
case DLDataTypeCode::kDLBfloat:
|
|
switch (dtype.bits) {
|
|
case 16:
|
|
return DataTypeImpl::GetType<BFloat16>();
|
|
default:
|
|
ORT_THROW("Unsupported kBFloat bits " + std::to_string(dtype.bits));
|
|
}
|
|
default:
|
|
ORT_THROW("Unsupported code " + std::to_string(dtype.code));
|
|
}
|
|
}
|
|
|
|
const char* GetOrtDeviceName(const OrtDevice& device) {
|
|
switch (device.Type()) {
|
|
case OrtDevice::CPU:
|
|
return CPU;
|
|
case OrtDevice::GPU:
|
|
return CUDA;
|
|
default:
|
|
ORT_THROW("Unknown device type: ", device.Type());
|
|
}
|
|
}
|
|
|
|
bool IsContiguousTensor(const DLTensor& tensor) {
|
|
if (!tensor.strides) {
|
|
return true;
|
|
}
|
|
|
|
int64_t running_size = 1;
|
|
for (int i = tensor.ndim - 1; i >= 0; i--) {
|
|
if (tensor.shape[i] == 0) {
|
|
return true;
|
|
}
|
|
|
|
if (tensor.shape[i] != 1 && tensor.strides[i] != running_size) {
|
|
return false;
|
|
}
|
|
|
|
running_size *= tensor.shape[i];
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
DLDevice GetDlpackDevice(const OrtValue& ort_value, const int64_t& device_id) {
|
|
ORT_ENFORCE(ort_value.IsTensor(), "Only OrtValues that are Tensors are currently supported");
|
|
DLDevice device;
|
|
device.device_id = static_cast<int>(device_id);
|
|
const Tensor& tensor = ort_value.Get<Tensor>();
|
|
const auto& location = tensor.Location();
|
|
switch (location.device.Type()) {
|
|
case OrtDevice::CPU:
|
|
device.device_type = DLDeviceType::kDLCPU;
|
|
break;
|
|
case OrtDevice::GPU:
|
|
#ifdef USE_ROCM
|
|
device.device_type = DLDeviceType::kDLROCM;
|
|
#else
|
|
device.device_type = DLDeviceType::kDLCUDA;
|
|
#endif
|
|
break;
|
|
default:
|
|
ORT_THROW("Cannot pack tensors on this device.");
|
|
}
|
|
|
|
return device;
|
|
}
|
|
|
|
struct OrtDLManagedTensor {
|
|
OrtValue handle;
|
|
DLManagedTensor tensor;
|
|
};
|
|
|
|
static void DlpackDeleter(DLManagedTensor* arg) { delete static_cast<OrtDLManagedTensor*>(arg->manager_ctx); }
|
|
|
|
// This function should use smart pointers inside
|
|
#if defined(_MSC_VER) && !defined(__clang__)
|
|
#pragma warning(push)
|
|
#pragma warning(disable : 26409)
|
|
#pragma warning(disable : 26400)
|
|
#endif
|
|
// This function returns a pointer to DLManagedTensor constructed from an OrtValue
|
|
// The OrtValue inside OrtDLManagedTensor will increase its own buffer's ref count by one
|
|
// When the consumer of DLManagedTensor is done with the tensor, it should invoke the deleter.
|
|
DLManagedTensor* OrtValueToDlpack(OrtValue& ort_value) {
|
|
ORT_ENFORCE(ort_value.IsTensor(), "Only tensor type OrtValues are supported");
|
|
OrtDLManagedTensor* ort_dlmanaged_tensor(new OrtDLManagedTensor);
|
|
|
|
Tensor& tensor = *ort_value.GetMutable<Tensor>();
|
|
ort_dlmanaged_tensor->handle = ort_value;
|
|
ort_dlmanaged_tensor->tensor.manager_ctx = ort_dlmanaged_tensor;
|
|
ort_dlmanaged_tensor->tensor.deleter = &DlpackDeleter;
|
|
ort_dlmanaged_tensor->tensor.dl_tensor.data = (tensor.MutableDataRaw());
|
|
ort_dlmanaged_tensor->tensor.dl_tensor.device = GetDlpackDevice(ort_value, tensor.Location().device.Id());
|
|
ort_dlmanaged_tensor->tensor.dl_tensor.ndim = static_cast<int>(tensor.Shape().NumDimensions());
|
|
ort_dlmanaged_tensor->tensor.dl_tensor.dtype = GetDlpackDataType(ort_value);
|
|
ort_dlmanaged_tensor->tensor.dl_tensor.shape =
|
|
tensor.Shape().NumDimensions() > 0 ? &const_cast<TensorShape&>(tensor.Shape())[0] : nullptr;
|
|
ort_dlmanaged_tensor->tensor.dl_tensor.strides = nullptr;
|
|
ort_dlmanaged_tensor->tensor.dl_tensor.byte_offset = 0;
|
|
return &(ort_dlmanaged_tensor->tensor);
|
|
}
|
|
#if defined(_MSC_VER) && !defined(__clang__)
|
|
#pragma warning(pop)
|
|
#endif
|
|
OrtValue DlpackToOrtValue(DLManagedTensor* dlpack, bool is_bool_tensor) {
|
|
// ORT only supports contiguous tensor for now.
|
|
ORT_ENFORCE(IsContiguousTensor(dlpack->dl_tensor), "ORT only supports contiguous tensor for now.");
|
|
OrtDevice device = GetOrtDevice(dlpack->dl_tensor.device);
|
|
MLDataType data_type = GetOrtValueDataType(dlpack->dl_tensor.dtype, is_bool_tensor);
|
|
OrtMemoryInfo info(GetOrtDeviceName(device), OrtDeviceAllocator, device, device.Id());
|
|
std::unique_ptr<Tensor> p_tensor = std::make_unique<Tensor>(
|
|
data_type, TensorShape(dlpack->dl_tensor.shape, static_cast<size_t>(dlpack->dl_tensor.ndim)),
|
|
dlpack->dl_tensor.data, info);
|
|
|
|
OrtValue ort_value;
|
|
std::function<void(void*)> deleter = [dlpack](void* p) {
|
|
ORT_ENFORCE(dlpack->deleter != NULL, "A dlpack structure must have a deleter.");
|
|
dlpack->deleter(dlpack);
|
|
auto deleter = DataTypeImpl::GetType<Tensor>()->GetDeleteFunc();
|
|
if (deleter != NULL)
|
|
deleter(p);
|
|
};
|
|
|
|
ort_value.Init(p_tensor.release(), DataTypeImpl::GetType<Tensor>(), deleter);
|
|
return ort_value;
|
|
}
|
|
|
|
} // namespace dlpack
|
|
} // namespace onnxruntime
|