intel-extension-for-pytorch

1
import torch
2
import intel_extension_for_pytorch as ipex
3
from functools import wraps
4

5

6
class AutoMixPrecision(object):
7
    def __init__(self, enable_or_not=False, train=False):
8
        self.old_value = ipex.get_auto_mix_precision()
9
        self.train_old_value = ipex.get_train()
10
        self.enable_or_not = enable_or_not
11
        self.train = train
12

13
    def __enter__(self):
14
        if self.enable_or_not:
15
            ipex.enable_auto_mixed_precision(
16
                mixed_dtype=torch.bfloat16, train=self.train
17
            )
18
        else:
19
            ipex.enable_auto_mixed_precision(mixed_dtype=None)
20

21
    def __exit__(self, *args, **kwargs):
22
        if self.old_value:
23
            ipex.enable_auto_mixed_precision(
24
                mixed_dtype=torch.bfloat16, train=self.train_old_value
25
            )
26
        else:
27
            ipex.enable_auto_mixed_precision(mixed_dtype=None)
28

29

30
class AutoDNNL(object):
31
    def __init__(self, enable_or_not=False):
32
        self.old_value = ipex._get_auto_optimization()
33
        self.enable_or_not = enable_or_not
34

35
    def __enter__(self):
36
        if self.enable_or_not:
37
            ipex.core.enable_auto_dnnl()
38
        else:
39
            ipex.core.disable_auto_dnnl()
40

41
    def __exit__(self, *args, **kwargs):
42
        if self.old_value:
43
            ipex.core.enable_auto_dnnl()
44
        else:
45
            ipex.core.disable_auto_dnnl()
46

47

48
def runtime_thread_affinity_test_env(func):
49
    @wraps(func)
50
    def wrapTheFunction(*args):
51
        # In some cases, the affinity of main thread may be changed: MultiStreamModule of stream 1
52
        # Ensure, we restore the affinity of main thread
53
        previous_cpu_pool = ipex._C.get_current_cpu_pool()
54
        func(*args)
55
        ipex._C.set_cpu_pool(previous_cpu_pool)
56

57
    return wrapTheFunction
58