mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
4dfaca6fb1
I want to start using `TensorMetadata` elsewhere in profiler so we have a common representation of Tensor. The main changes in this PR are: 1) Replace raw pointers with strong typedefs and create a custom type caster to handle moving them to Python. 2) Adding a `device()` method to handle reassembling type and index. Differential Revision: [D39563965](https://our.internmc.facebook.com/intern/diff/D39563965/) Pull Request resolved: https://github.com/pytorch/pytorch/pull/85161 Approved by: https://github.com/chaekit
34 lines
915 B
C++
34 lines
915 B
C++
#pragma once
|
|
|
|
#include <pybind11/pybind11.h>
|
|
|
|
#include <c10/util/strong_type.h>
|
|
#include <torch/csrc/utils/pybind.h>
|
|
#include <torch/csrc/utils/python_numbers.h>
|
|
|
|
namespace pybind11 {
|
|
namespace detail {
|
|
// Strong typedefs don't make much sense in Python since everything is duck
|
|
// typed. So instead we simply cast them to ints, return them, and let the
|
|
// caller handle correctness.
|
|
template <typename T>
|
|
struct strong_pointer_type_caster {
|
|
template <typename T_>
|
|
static handle cast(
|
|
T_&& src,
|
|
return_value_policy /*policy*/,
|
|
handle /*parent*/) {
|
|
const auto* ptr = reinterpret_cast<const void*>(src.value_of());
|
|
return ptr ? handle(THPUtils_packUInt64(reinterpret_cast<intptr_t>(ptr)))
|
|
: none();
|
|
}
|
|
|
|
bool load(handle /*src*/, bool /*convert*/) {
|
|
return false;
|
|
}
|
|
|
|
PYBIND11_TYPE_CASTER(T, _("strong_pointer"));
|
|
};
|
|
} // namespace detail
|
|
} // namespace pybind11
|