mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
bc5710f2f7
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/45349 Test Plan: Imported from OSS Reviewed By: Krovatkin Differential Revision: D23935518 Pulled By: ZolotukhinM fbshipit-source-id: 5a7c508c6fc84eafbc23399f095d732b903510dc
41 lines
1.6 KiB
Python
41 lines
1.6 KiB
Python
import torch
|
|
|
|
def set_fuser(fuser_name, executor_name):
|
|
assert fuser_name in ['te', 'old', 'none', 'default']
|
|
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_set_num_profiled_runs(1)
|
|
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)
|
|
elif fuser_name == 'default':
|
|
pass
|
|
|
|
# --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)
|
|
torch._C._jit_set_num_profiled_runs(2)
|
|
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)
|
|
elif executor_name == 'default':
|
|
pass
|