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

Reland: D27652485: [nnc] Enable CPU fusion only when num_threads == 1" (#56120)

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

This reverts commit ad17fadbfc786dc1ccb42e822208ff03c2a2b72c (D27786457).

The big annoyance here is that depending on the threading mode you may not be
able to toggle num_threads at will, so the fusion tests won't fail.

I hate this solution, but I'm adding a secondary override for the TE fuser.
Now you need to both turn on fusion (_jit_override_can_fuse_on_cpu), and you're
OK if you're running with 1 thread, or you can add
`_jit_set_texpr_parallel_cpu_enabled` to enable it anyways.

This is (a) mainly for tests, since a real user probably won't fiddle aimlessly
with the thread count, and (b) will go away once NNC's threading support is
fully baked.

Test Plan: Imported from OSS

Reviewed By: Krovatkin

Differential Revision: D27788199

Pulled By: bertmaher

fbshipit-source-id: 070d04474f15e9689dbdf8cc1fde43050c6506b1
This commit is contained in:
Bert Maher
2021-04-15 15:47:03 -07:00
committed by Facebook GitHub Bot
parent e1752ffa04
commit 8e82e932f3
7 changed files with 39 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -631,6 +631,8 @@ void initJITBindings(PyObject* module) {
.def("_jit_texpr_set_fallback_allowed", &tensorexpr::setFallbackAllowed)
.def("_jit_set_texpr_reductions_enabled", &setTexprReductionsEnabled)
.def("_jit_texpr_reductions_enabled", &texprReductionsEnabled)
.def("_jit_set_texpr_parallel_cpu_enabled", &setTexprParallelCPUEnabled)
.def("_jit_texpr_parallel_cpu_enabled", &texprParallelCPUEnabled)
.def(
"_jit_set_te_generate_block_code",
[](bool gen_block_code) {