forked from mooncake-track/Mooncake
106 lines
3.6 KiB
C++
106 lines
3.6 KiB
C++
#pragma once
|
|
|
|
#include <pybind11/numpy.h>
|
|
#include <pybind11/pybind11.h>
|
|
#include <pybind11/stl.h>
|
|
|
|
#include <array>
|
|
#include <functional>
|
|
|
|
namespace py = pybind11;
|
|
|
|
namespace mooncake {
|
|
|
|
// Avoid global py::module_ objects
|
|
inline py::module_ torch_module() { return py::module_::import("torch"); }
|
|
|
|
enum class TensorDtype : int32_t {
|
|
FLOAT32 = 0,
|
|
FLOAT64 = 1,
|
|
INT8 = 2,
|
|
UINT8 = 3,
|
|
INT16 = 4,
|
|
UINT16 = 5,
|
|
INT32 = 6,
|
|
UINT32 = 7,
|
|
INT64 = 8,
|
|
UINT64 = 9,
|
|
BOOL = 10,
|
|
FLOAT16 = 11,
|
|
BFLOAT16 = 12,
|
|
FLOAT8_E4M3 = 13,
|
|
FLOAT8_E5M2 = 14,
|
|
NR_DTYPES = 15,
|
|
UNKNOWN = -1
|
|
};
|
|
|
|
template <typename T>
|
|
py::array create_typed_array(char *exported_data, size_t offset,
|
|
size_t total_length, bool take_ownership) {
|
|
if (take_ownership) {
|
|
py::capsule free_when_done(
|
|
exported_data, [](void *p) { delete[] static_cast<char *>(p); });
|
|
return py::array_t<T>({static_cast<ssize_t>(total_length / sizeof(T))},
|
|
(T *)(exported_data + offset), free_when_done);
|
|
}
|
|
|
|
return py::array_t<T>({static_cast<ssize_t>(total_length / sizeof(T))},
|
|
(T *)(exported_data + offset), py::none());
|
|
}
|
|
|
|
using ArrayCreatorFunc = std::function<py::array(char *, size_t, size_t, bool)>;
|
|
|
|
static const std::array<ArrayCreatorFunc, 15> array_creators = {{
|
|
create_typed_array<float>, // FLOAT32 = 0
|
|
create_typed_array<double>, // FLOAT64 = 1
|
|
create_typed_array<int8_t>, // INT8 = 2
|
|
create_typed_array<uint8_t>, // UINT8 = 3
|
|
create_typed_array<int16_t>, // INT16 = 4
|
|
create_typed_array<uint16_t>, // UINT16 = 5
|
|
create_typed_array<int32_t>, // INT32 = 6
|
|
create_typed_array<uint32_t>, // UINT32 = 7
|
|
create_typed_array<int64_t>, // INT64 = 8
|
|
create_typed_array<uint64_t>, // UINT64 = 9
|
|
create_typed_array<bool>, // BOOL = 10
|
|
create_typed_array<uint16_t>, // FLOAT16 = 11 (using uint16_t as storage)
|
|
create_typed_array<uint16_t>, // BFLOAT16 = 12 (using uint16_t as storage)
|
|
create_typed_array<uint8_t>, // FLOAT8_E4M3 = 13 (using uint8_t as storage)
|
|
create_typed_array<uint8_t>, // FLOAT8_E5M2 = 14 (using uint8_t as storage)
|
|
}};
|
|
|
|
inline TensorDtype get_tensor_dtype(py::object dtype_obj) {
|
|
if (dtype_obj.is_none()) {
|
|
return TensorDtype::UNKNOWN;
|
|
}
|
|
|
|
auto torch = torch_module();
|
|
|
|
if (dtype_obj.equal(torch.attr("float32"))) return TensorDtype::FLOAT32;
|
|
if (dtype_obj.equal(torch.attr("float64"))) return TensorDtype::FLOAT64;
|
|
if (dtype_obj.equal(torch.attr("int8"))) return TensorDtype::INT8;
|
|
if (dtype_obj.equal(torch.attr("uint8"))) return TensorDtype::UINT8;
|
|
if (dtype_obj.equal(torch.attr("int16"))) return TensorDtype::INT16;
|
|
if (dtype_obj.equal(torch.attr("uint16"))) return TensorDtype::UINT16;
|
|
if (dtype_obj.equal(torch.attr("int32"))) return TensorDtype::INT32;
|
|
if (dtype_obj.equal(torch.attr("uint32"))) return TensorDtype::UINT32;
|
|
if (dtype_obj.equal(torch.attr("int64"))) return TensorDtype::INT64;
|
|
if (dtype_obj.equal(torch.attr("uint64"))) return TensorDtype::UINT64;
|
|
if (dtype_obj.equal(torch.attr("bool"))) return TensorDtype::BOOL;
|
|
if (dtype_obj.equal(torch.attr("float16"))) return TensorDtype::FLOAT16;
|
|
if (dtype_obj.equal(torch.attr("bfloat16"))) return TensorDtype::BFLOAT16;
|
|
if (dtype_obj.equal(torch.attr("float8_e4m3fn")))
|
|
return TensorDtype::FLOAT8_E4M3;
|
|
if (dtype_obj.equal(torch.attr("float8_e5m2")))
|
|
return TensorDtype::FLOAT8_E5M2;
|
|
|
|
return TensorDtype::UNKNOWN;
|
|
}
|
|
|
|
struct TensorMetadata {
|
|
int32_t dtype;
|
|
int32_t ndim;
|
|
int64_t shape[4];
|
|
};
|
|
|
|
} // namespace mooncake
|