mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
89 lines
2.3 KiB
C++
89 lines
2.3 KiB
C++
#ifndef THP_COPY_UTILS_H
|
|
#define THP_COPY_UTILS_H
|
|
|
|
#include <functional>
|
|
#include <vector>
|
|
#include "Types.h"
|
|
|
|
|
|
typedef std::function<void(PyObject*, PyObject*)> THPCopyFunction;
|
|
struct THPCopyInfo {
|
|
PyTypeObject* srcType; // Python type of src tensor/storage
|
|
THPCopyFunction copy; // copy function
|
|
bool async; // true if copy implements an 'async' copy
|
|
};
|
|
typedef std::vector<THPCopyInfo> THPCopyList;
|
|
|
|
inline bool tryTHPCopy(const THPCopyList& v, PyObject* dst, PyObject* src, bool async)
|
|
{
|
|
for (auto it = v.begin(); it != v.end(); ++it) {
|
|
if (it->async == async && PyType_IsSubtype(Py_TYPE(src), it->srcType)) {
|
|
(it->copy)(dst, src);
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
inline bool THPCopy(const THPCopyList& v, PyObject* dst, PyObject* src, bool async)
|
|
{
|
|
if (tryTHPCopy(v, dst, src, async)) {
|
|
return true;
|
|
} else if (async && tryTHPCopy(v, dst, src, false)) {
|
|
return true;
|
|
}
|
|
THPUtils_setError("copy from %s to %s isn't implemented",
|
|
THPUtils_typename(src), THPUtils_typename(dst));
|
|
return false;
|
|
}
|
|
|
|
inline PyObject * THPCopyMethod(const THPCopyList& v, PyObject *self, PyObject *args, PyObject *kwargs)
|
|
{
|
|
PyObject *src;
|
|
PyObject *async = Py_False;
|
|
static char *kwlist[] = {"source", "async", NULL};
|
|
if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|O:copy_", kwlist, &src, &async)) {
|
|
return NULL;
|
|
}
|
|
if (!PyBool_Check(async)) {
|
|
return PyErr_Format(PyExc_TypeError, "copy_() expected bool for argument async (got '%s')",
|
|
Py_TYPE(async)->tp_name);
|
|
}
|
|
|
|
if (!THPCopy(v, self, src, async == Py_True)) {
|
|
return NULL;
|
|
}
|
|
|
|
Py_INCREF(self);
|
|
return self;
|
|
}
|
|
|
|
template <typename TensorDst, typename TensorSrc>
|
|
void THPInsertCopyFunction(
|
|
THPCopyList& copyList,
|
|
void (*copyFunc)(LIBRARY_STATE_TYPE TensorDst* x, TensorSrc* z),
|
|
bool async=false)
|
|
{
|
|
auto wrapper = [copyFunc](PyObject* dst_, PyObject* src_) {
|
|
TensorDst* dst = THPTypeInfo<TensorDst>::cdata(dst_);
|
|
TensorSrc* src = THPTypeInfo<TensorSrc>::cdata(src_);
|
|
|
|
PyThreadState *_save = NULL;
|
|
try {
|
|
Py_UNBLOCK_THREADS;
|
|
copyFunc(LIBRARY_STATE dst, src);
|
|
Py_BLOCK_THREADS;
|
|
} catch (...) {
|
|
if (_save) {
|
|
Py_BLOCK_THREADS;
|
|
}
|
|
throw;
|
|
}
|
|
};
|
|
|
|
PyTypeObject* srcType = THPTypeInfo<TensorSrc>::pyType();
|
|
copyList.push_back({ srcType, wrapper, async });
|
|
}
|
|
|
|
#endif
|