mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
33d51a9b32
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/43967 Test Plan: Imported from OSS Reviewed By: asuhan Differential Revision: D23469048 Pulled By: bertmaher fbshipit-source-id: 1005a7ae08974059ff9d467492caa3a388070eeb
35 lines
1.4 KiB
Python
35 lines
1.4 KiB
Python
import torch
|
|
|
|
def set_fuser(fuser_name, executor_name):
|
|
assert fuser_name in ['te', 'old', 'none']
|
|
if fuser_name == 'te':
|
|
torch._C._jit_set_profiling_executor(True)
|
|
torch._C._jit_set_profiling_mode(True)
|
|
torch._C._jit_set_bailout_depth(20)
|
|
torch._C._jit_override_can_fuse_on_cpu(False)
|
|
torch._C._jit_override_can_fuse_on_gpu(True)
|
|
torch._C._jit_set_texpr_fuser_enabled(True)
|
|
elif fuser_name == 'old':
|
|
torch._C._jit_set_profiling_executor(False)
|
|
torch._C._jit_set_profiling_mode(False)
|
|
torch._C._jit_override_can_fuse_on_gpu(True)
|
|
torch._C._jit_set_texpr_fuser_enabled(False)
|
|
elif fuser_name == 'none':
|
|
torch._C._jit_set_profiling_executor(False)
|
|
torch._C._jit_set_profiling_mode(False)
|
|
torch._C._jit_override_can_fuse_on_gpu(False)
|
|
torch._C._jit_override_can_fuse_on_cpu(False)
|
|
torch._C._jit_set_texpr_fuser_enabled(False)
|
|
|
|
# --executor overrides settings of --fuser
|
|
if executor_name == 'profiling':
|
|
torch._C._jit_set_profiling_executor(True)
|
|
torch._C._jit_set_profiling_mode(True)
|
|
torch._C._jit_set_bailout_depth(20)
|
|
elif executor_name == 'simple':
|
|
torch._C._jit_set_profiling_executor(True)
|
|
torch._C._jit_set_profiling_mode(False)
|
|
elif executor_name == 'legacy':
|
|
torch._C._jit_set_profiling_executor(False)
|
|
torch._C._jit_set_profiling_mode(False)
|