mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
d59714e3b1
- Reduce setup.py diff. - Expunge WITH_TOFFEE from codebase. - Elaborate on a comment. - Move gen_toffee.sh to tools - Delete densenet test. - Use 'using' to inherit a constructor. - Delete outdated comment. - Comment about why primspecs can return fewer outputs. - Remove dead, commented out includes. Signed-off-by: Edward Z. Yang <ezyang@fb.com>
70 lines
2.1 KiB
C++
70 lines
2.1 KiB
C++
#include <Python.h>
|
|
#include <pybind11/pybind11.h>
|
|
|
|
namespace py = pybind11;
|
|
|
|
#include "THP.h"
|
|
#include "torch/csrc/jit/ir.h"
|
|
#include "torch/csrc/jit/graph_fuser.h"
|
|
#include "torch/csrc/jit/init_pass.h"
|
|
#include "torch/csrc/jit/dead_code_elimination.h"
|
|
#include "torch/csrc/jit/python_tracer.h"
|
|
#include "torch/csrc/utils/python_strings.h"
|
|
#include "torch/csrc/DynamicTypes.h"
|
|
#include "torch/csrc/toffee/export.h"
|
|
|
|
PyObject * THPJIT_initExtension(PyObject *_unused)
|
|
{
|
|
// Leaving this code here, because it will likely be useful at some point
|
|
//PyObject *jit_module = PyImport_ImportModule("torch.jit");
|
|
//THPUtils_assert(jit_module, "class loader couldn't access "
|
|
//"torch.jit module");
|
|
//PyObject *jit_dict = PyModule_GetDict(jit_module);
|
|
|
|
Py_RETURN_TRUE;
|
|
}
|
|
|
|
// stub to run all C++ only tests for the JIT
|
|
// the stuff in test_jit.cpp is kept separate from the rest of PyTorch
|
|
// so we can build and iterate on it faster.
|
|
// from test_jit.cpp
|
|
namespace torch { namespace jit { extern void runJITCPPTests(); } };
|
|
|
|
namespace {
|
|
|
|
using namespace torch::jit;
|
|
|
|
using pass_type = void (std::shared_ptr<Graph>&);
|
|
|
|
template<pass_type pass>
|
|
PyObject * wrap_pass(PyObject *_unused, PyObject *py_state) {
|
|
HANDLE_TH_ERRORS
|
|
auto trace = py::handle(py_state).cast<tracer::TracingState*>();
|
|
pass(trace->graph);
|
|
Py_RETURN_NONE;
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
PyObject * run_cpp_tests(PyObject *_unused, PyObject *_unused2) {
|
|
HANDLE_TH_ERRORS
|
|
runJITCPPTests();
|
|
Py_RETURN_NONE;
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
struct PyMethodDef _THPJIT_methods[] = {
|
|
{"_jit_init", (PyCFunction)THPJIT_initExtension, METH_NOARGS, NULL},
|
|
{"_jit_pass_init", (PyCFunction)wrap_pass<MatchJITOps>, METH_O, "init"},
|
|
{"_jit_pass_fuse", (PyCFunction)wrap_pass<FuseGraph>, METH_O, "fuse"},
|
|
{"_jit_pass_dco", (PyCFunction)wrap_pass<EliminateDeadCode>, METH_O, "dco"},
|
|
{"_jit_pass_lint", (PyCFunction)wrap_pass<LintGraph>, METH_O, "lint"},
|
|
{"_jit_run_cpp_tests",(PyCFunction)run_cpp_tests, METH_NOARGS, NULL},
|
|
{NULL}
|
|
};
|
|
|
|
} // anonymous namespace
|
|
|
|
PyMethodDef* THPJIT_methods() {
|
|
return _THPJIT_methods;
|
|
}
|