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

Allow single node fusion for nvfuser (#70000)

Browse Source 
Summary:
Setting `PYTORCH_NVFUSER_ONE_OP_FUSION=1` will take all nodes nvFuser support, instead of waiting for fusion opportunity.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/70000

Reviewed By: samdow

Differential Revision: D33292195

Pulled By: davidberard98

fbshipit-source-id: 8ed5ce5e82fbb6737e8ab5ce4223b038eaf47756
This commit is contained in:
jjsjann123
2021-12-23 17:06:28 -08:00
committed by Facebook GitHub Bot
parent 5ccf28d066
commit e429a68478
5 changed files with 106 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -687,6 +687,18 @@ void initJITBindings(PyObject* module) {
checkAliasAnnotation(g, std::move(stack), unqualified_op_name);
})
.def("_jit_set_nvfuser_enabled", &RegisterCudaFuseGraph::registerPass)
.def(
"_jit_set_nvfuser_single_node_mode",
[](bool flag) { return fuser::cuda::setSingletonFusion(flag); })
.def(
"_jit_nvfuser_single_node_mode",
[]() { return fuser::cuda::getSingletonFusion(); })
.def(
"_jit_set_nvfuser_horizontal_mode",
[](bool flag) { return fuser::cuda::setHorizontalFusion(flag); })
.def(
"_jit_nvfuser_horizontal_mode",
[]() { return fuser::cuda::getHorizontalFusion(); })
.def(
"_jit_set_nvfuser_guard_mode",
[](bool profiling_flag) {