mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
dd3a77bc96
Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/105928 Approved by: https://github.com/albanD
37 lines
1.4 KiB
Python
37 lines
1.4 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._get_graph_executor_optimize(True)
|
|
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._get_graph_executor_optimize(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._get_graph_executor_optimize(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._get_graph_executor_optimize(True)
|
|
elif executor_name == "simple":
|
|
torch._C._get_graph_executor_optimize(False)
|
|
elif executor_name == "legacy":
|
|
torch._C._jit_set_profiling_executor(False)
|
|
torch._C._get_graph_executor_optimize(True)
|
|
elif executor_name == "default":
|
|
pass
|