pytorch

1
import torch
2

3

4
def set_fuser(fuser_name, executor_name):
5
    assert fuser_name in ["te", "old", "none", "default"]
6
    if fuser_name == "te":
7
        torch._C._jit_set_profiling_executor(True)
8
        torch._C._get_graph_executor_optimize(True)
9
        torch._C._jit_override_can_fuse_on_cpu(False)
10
        torch._C._jit_override_can_fuse_on_gpu(True)
11
        torch._C._jit_set_texpr_fuser_enabled(True)
12
    elif fuser_name == "old":
13
        torch._C._jit_set_profiling_executor(False)
14
        torch._C._get_graph_executor_optimize(False)
15
        torch._C._jit_override_can_fuse_on_gpu(True)
16
        torch._C._jit_set_texpr_fuser_enabled(False)
17
    elif fuser_name == "none":
18
        torch._C._jit_set_profiling_executor(False)
19
        torch._C._get_graph_executor_optimize(False)
20
        torch._C._jit_override_can_fuse_on_gpu(False)
21
        torch._C._jit_override_can_fuse_on_cpu(False)
22
        torch._C._jit_set_texpr_fuser_enabled(False)
23
    elif fuser_name == "default":
24
        pass
25

26
    # --executor overrides settings of --fuser
27
    if executor_name == "profiling":
28
        torch._C._jit_set_profiling_executor(True)
29
        torch._C._get_graph_executor_optimize(True)
30
    elif executor_name == "simple":
31
        torch._C._get_graph_executor_optimize(False)
32
    elif executor_name == "legacy":
33
        torch._C._jit_set_profiling_executor(False)
34
        torch._C._get_graph_executor_optimize(True)
35
    elif executor_name == "default":
36
        pass
37