mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
9f4bcdf075
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/11254 Previously we use DeviceType in caffe2.proto directly, but it's an `enum` and have implicit conversion to int, which does not have type safety, e.g. we have to explicitly check for a device type is valid in event.h: ``` template <int d> struct EventCreateFunctionRegisterer { explicit EventCreateFunctionRegisterer(EventCreateFunction f) { static_assert(d < MaxDeviceTypes, ""); Event::event_creator_[d] = f; } }; ``` at::DeviceType is an `enum class`, and it does not have implicit conversion to int, and provides better type safety guarantees. In this diff we have done the following refactor(taking CPU as an example): 1. caffe2::DeviceType → caffe2::DeviceTypeProto 2. caffe2::CPU → caffe2::PROTO_CPU 3. caffe2::DeviceType = at::DeviceType 4. caffe2::CPU = at::DeviceType::CPU codemod -d caffe2/caffe2 --extensions h,cc,cpp 'device_type\(\), ' 'device_type(), PROTO_' + some manual changes In short, after this diff, in c++, caffe2::CPU refers to the at::DeviceType::CPU and the old proto caffe2::CPU will be caffe2::PROTO_CPU. In python side, we have a temporary workaround that alias `caffe2_pb2.CPU = caffe2_pb2.PROOT_CPU` to make the change easier to review and this will be removed later. Reviewed By: ezyang Differential Revision: D9545704 fbshipit-source-id: 461a28a4ca74e616d3ee183a607078a717fd38a7
94 lines
2.9 KiB
C++
94 lines
2.9 KiB
C++
#define NO_IMPORT_ARRAY
|
|
|
|
#include "pybind_state.h"
|
|
|
|
#include <pybind11/pybind11.h>
|
|
#include <pybind11/stl.h>
|
|
|
|
#include "caffe2/core/hip/common_miopen.h"
|
|
#include "caffe2/core/hip/context_hip.h"
|
|
#include "caffe2/operators/hip/operator_fallback_hip.h"
|
|
#include "caffe2/python/pybind_state_registry.h"
|
|
|
|
namespace caffe2 {
|
|
namespace python {
|
|
|
|
REGISTER_HIP_OPERATOR(Python, GPUFallbackOp<PythonOp<CPUContext, false>>);
|
|
REGISTER_HIP_OPERATOR(
|
|
PythonGradient,
|
|
GPUFallbackOp<PythonGradientOp<CPUContext, false>>);
|
|
|
|
REGISTER_HIP_OPERATOR(PythonDLPack, PythonOp<HIPContext, true>);
|
|
REGISTER_HIP_OPERATOR(PythonDLPackGradient, PythonGradientOp<HIPContext, true>);
|
|
|
|
REGISTER_BLOB_FEEDER(HIP, TensorFeeder<HIPContext>);
|
|
|
|
namespace py = pybind11;
|
|
|
|
void addHIPGlobalMethods(py::module& m) {
|
|
m.def("num_hip_devices", &NumHipDevices);
|
|
m.def("set_default_gpu_id", &SetDefaultGPUID);
|
|
m.def("get_default_gpu_id", &GetDefaultGPUID);
|
|
m.def("get_hip_version", &HipVersion);
|
|
m.def("get_miopen_version", &miopenCompiledVersion);
|
|
m.def("get_hip_peer_access_pattern", []() {
|
|
std::vector<std::vector<bool>> pattern;
|
|
CAFFE_ENFORCE(caffe2::GetHipPeerAccessPattern(&pattern));
|
|
return pattern;
|
|
});
|
|
m.def("get_device_properties", [](int deviceid) {
|
|
auto& prop = GetDeviceProperty(deviceid);
|
|
std::map<std::string, py::object> obj;
|
|
obj["name"] = py::cast(prop.name);
|
|
obj["major"] = py::cast(prop.major);
|
|
obj["minor"] = py::cast(prop.minor);
|
|
return obj;
|
|
});
|
|
};
|
|
|
|
void addHIPObjectMethods(py::module& m) {
|
|
py::class_<DLPackWrapper<HIPContext>>(m, "DLPackTensorHIP")
|
|
.def_property_readonly(
|
|
"data",
|
|
[](DLPackWrapper<HIPContext>* t) -> py::object {
|
|
CAFFE_ENFORCE_EQ(
|
|
t->device_option.device_type(),
|
|
PROTO_HIP,
|
|
"Expected HIP device option for HIP tensor");
|
|
|
|
return t->data();
|
|
},
|
|
"Return DLPack tensor with tensor's data.")
|
|
.def(
|
|
"feed",
|
|
[](DLPackWrapper<HIPContext>* t, py::object obj) {
|
|
CAFFE_ENFORCE_EQ(
|
|
t->device_option.device_type(),
|
|
PROTO_HIP,
|
|
"Expected HIP device option for HIP tensor");
|
|
t->feed(obj);
|
|
},
|
|
"Copy data from given DLPack tensor into this tensor.")
|
|
.def_property_readonly(
|
|
"_shape",
|
|
[](const DLPackWrapper<HIPContext>& t) { return t.tensor->dims(); })
|
|
.def(
|
|
"_reshape",
|
|
[](DLPackWrapper<HIPContext>* t, std::vector<TIndex> dims) {
|
|
t->tensor->Resize(dims);
|
|
});
|
|
}
|
|
|
|
PYBIND11_MODULE(caffe2_pybind11_state_hip, m) {
|
|
m.doc() = "pybind11 stateful interface to Caffe2 workspaces - GPU edition";
|
|
|
|
addGlobalMethods(m);
|
|
addHIPGlobalMethods(m);
|
|
addObjectMethods(m);
|
|
addHIPObjectMethods(m);
|
|
for (const auto& addition : PybindAdditionRegistry()->Keys()) {
|
|
PybindAdditionRegistry()->Create(addition, m);
|
|
}
|
|
}
|
|
} // namespace python
|
|
} // namespace caffe2
|