frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 13:44:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

Reland "add an API for external backends to register custom device names (#86992)" (#87453)

Browse Source 
Re-land of https://github.com/pytorch/pytorch/pull/86992

This reverts commit a895af92506f206889610251624590798d0deabd.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/87453
Approved by: https://github.com/ezyang, https://github.com/albanD
This commit is contained in:
Brian Hirsh
2022-10-21 06:21:41 -07:00
committed by PyTorch MergeBot
parent 70c46d32e2
commit ce0c6e828e
8 changed files with 122 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
torch/csrc/Module.cpp
View File

@ -441,6 +441,20 @@ PyObject* THModule_getCppBacktrace(PyObject* _unused, PyObject* args) {
c10::get_backtrace(frames_to_skip, maximum_number_of_frames, true));
END_HANDLE_TH_ERRORS
}
static PyObject* THModule_rename_privateuse1_backend(
PyObject* _unused,
PyObject* arg) {
HANDLE_TH_ERRORS
THPUtils_assert(
THPUtils_checkString(arg),
"_rename_privateuse1_backend expects a str, "
"but got %s",
THPUtils_typename(arg));
const std::string backend_name = THPUtils_unpackString(arg);
c10::register_privateuse1_backend(backend_name);
Py_RETURN_NONE;
END_HANDLE_TH_ERRORS
}
PyObject* THPModule_setAllowTF32CuDNN(PyObject* _unused, PyObject* arg) {
THPUtils_assert(
@ -990,6 +1004,10 @@ static PyMethodDef TorchMethods[] = {
{"_to_dlpack", THPModule_toDLPack, METH_O, nullptr},
{"_from_dlpack", THPModule_fromDLPack, METH_O, nullptr},
{"_get_cpp_backtrace", THModule_getCppBacktrace, METH_VARARGS, nullptr},
{"_rename_privateuse1_backend",
THModule_rename_privateuse1_backend,
METH_O,
nullptr},
{"set_flush_denormal", THPModule_setFlushDenormal, METH_O, nullptr},
{"get_default_dtype", THPModule_getDefaultDtype, METH_NOARGS, nullptr},
{"_get_default_device", THPModule_getDefaultDevice, METH_NOARGS, nullptr},