mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
a7ac591d3b
Summary: Adding support for DLPack tensors to Python op Reviewed By: Yangqing Differential Revision: D6577702 fbshipit-source-id: e14ef213fcdb2930ffe164667971a92aa8db503c
90 lines
3.0 KiB
C++
90 lines
3.0 KiB
C++
/**
|
|
* Copyright (c) 2016-present, Facebook, Inc.
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#include "pybind_state_dlpack.h"
|
|
|
|
namespace caffe2 {
|
|
namespace python {
|
|
|
|
namespace py = pybind11;
|
|
|
|
const DLDeviceType* CaffeToDLDeviceType(int device_type) {
|
|
static std::map<int, DLDeviceType> dl_device_type_map{
|
|
{CPU, kCPU},
|
|
{CUDA, kGPU},
|
|
};
|
|
const auto it = dl_device_type_map.find(device_type);
|
|
return it == dl_device_type_map.end() ? nullptr : &it->second;
|
|
}
|
|
|
|
const DLDataType* CaffeToDLType(const TypeMeta& meta) {
|
|
static std::map<CaffeTypeId, DLDataType> dl_type_map{
|
|
{TypeMeta::Id<int8_t>(), DLDataType{0, 8, 1}},
|
|
{TypeMeta::Id<int16_t>(), DLDataType{0, 16, 1}},
|
|
{TypeMeta::Id<int32_t>(), DLDataType{0, 32, 1}},
|
|
{TypeMeta::Id<int64_t>(), DLDataType{0, 64, 1}},
|
|
{TypeMeta::Id<uint8_t>(), DLDataType{1, 8, 1}},
|
|
{TypeMeta::Id<uint16_t>(), DLDataType{1, 16, 1}},
|
|
{TypeMeta::Id<float16>(), DLDataType{2, 16, 1}},
|
|
{TypeMeta::Id<float>(), DLDataType{2, 32, 1}},
|
|
{TypeMeta::Id<double>(), DLDataType{2, 64, 1}},
|
|
};
|
|
const auto it = dl_type_map.find(meta.id());
|
|
return it == dl_type_map.end() ? nullptr : &it->second;
|
|
}
|
|
|
|
const TypeMeta& DLTypeToCaffe(const DLDataType& dl_type) {
|
|
try {
|
|
if (dl_type.lanes != 1) {
|
|
throw std::invalid_argument("invalid type");
|
|
}
|
|
static std::map<int, std::map<int, TypeMeta>> dl_caffe_type_map{
|
|
{0,
|
|
std::map<int, TypeMeta>{
|
|
{8, TypeMeta::Make<int8_t>()},
|
|
{16, TypeMeta::Make<int16_t>()},
|
|
{32, TypeMeta::Make<int32_t>()},
|
|
{64, TypeMeta::Make<int64_t>()},
|
|
}},
|
|
{1,
|
|
std::map<int, TypeMeta>{
|
|
{8, TypeMeta::Make<uint8_t>()},
|
|
{16, TypeMeta::Make<uint16_t>()},
|
|
}},
|
|
{2,
|
|
std::map<int, TypeMeta>{
|
|
{16, TypeMeta::Make<float16>()},
|
|
{32, TypeMeta::Make<float>()},
|
|
{64, TypeMeta::Make<double>()},
|
|
}},
|
|
};
|
|
if (!dl_caffe_type_map.count(dl_type.code)) {
|
|
throw std::invalid_argument("invalid type");
|
|
}
|
|
const auto& bits_map = dl_caffe_type_map.at(dl_type.code);
|
|
if (!bits_map.count(dl_type.bits)) {
|
|
throw std::invalid_argument("invalid type");
|
|
}
|
|
return bits_map.at(dl_type.bits);
|
|
} catch (const std::invalid_argument& e) {
|
|
CAFFE_THROW(
|
|
"Unsupported DLDataType: ", dl_type.code, dl_type.bits, dl_type.lanes);
|
|
}
|
|
}
|
|
|
|
} // namespace python
|
|
} // namespace caffe2
|