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

JIT pass to insert XNNPACK ops (#34048)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/34048

Rewrites the graph to insert xnnpack prepack and packed run ops for
conv2d and linear.

Test Plan:
python test/test_xnnpack_integration.py

Imported from OSS

Differential Revision: D20185658

fbshipit-source-id: c4c073c912ad33e822e7beb4ed86c9f895129d55
This commit is contained in:
Kimish Patel
2020-03-14 12:48:24 -07:00
committed by Facebook GitHub Bot
parent 4c30fc7238
commit 7dd5da2026
9 changed files with 377 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
torch/csrc/jit/python/init.cpp
View File

@ -46,6 +46,7 @@
#include <torch/csrc/jit/passes/tensorexpr_fuser.h>
#include <torch/csrc/jit/passes/utils/check_alias_annotation.h>
#include <torch/csrc/jit/passes/freeze_module.h>
#include <torch/csrc/jit/passes/xnnpack_rewrite.h>
#include <torch/csrc/jit/runtime/print_handler.h>
#include <torch/csrc/jit/python/pybind_utils.h>
#include <torch/csrc/jit/python/python_arg_flatten.h>
@ -409,6 +410,16 @@ void initJITBindings(PyObject* module) {
[](Graph& g, std::vector<at::Tensor> inps) {
return debugGetFusedKernelCode(g, inps);
})
.def(
"_jit_pass_insert_xnnpack_ops",
[](std::shared_ptr<Graph>& graph) {
return insertXNNPACKOps(graph);
})
.def(
"_jit_pass_insert_xnnpack_ops",
[](script::Module& module) {
return insertXNNPACKOps(module);
})
.def(
"_jit_pass_onnx_unpack_quantized_weights",
[](std::shared_ptr<Graph>& graph,