intel-extension-for-pytorch

1
import unittest
2
import os
3
import subprocess
4

5
import intel_extension_for_pytorch._C as core
6

7
supported_isa_set = [
8
    "default",
9
    "avx2",
10
    "avx2_vnni",
11
    "avx512",
12
    "avx512_vnni",
13
    "avx512_bf16",
14
    "amx",
15
    "avx512_fp16",
16
]
17

18

19
def get_isa_val(isa_name):
20
    if isa_name == "default":
21
        return 0
22
    elif isa_name == "avx2":
23
        return 1
24
    elif isa_name == "avx2_vnni":
25
        return 2
26
    elif isa_name == "avx512":
27
        return 3
28
    elif isa_name == "avx512_vnni":
29
        return 4
30
    elif isa_name == "avx512_bf16":
31
        return 5
32
    elif isa_name == "amx":
33
        return 6
34
    elif isa_name == "avx512_fp16":
35
        return 7
36
    else:
37
        return 100
38

39

40
def get_ipex_isa_env_setting():
41
    env_isa = os.getenv("ATEN_CPU_CAPABILITY")
42
    return env_isa
43

44

45
def get_currnet_isa_level():
46
    return core._get_current_isa_level().lower()
47

48

49
def get_highest_binary_support_isa_level():
50
    return core._get_highest_binary_support_isa_level().lower()
51

52

53
def get_highest_cpu_support_isa_level():
54
    return core._get_highest_cpu_support_isa_level().lower()
55

56

57
def check_not_sync_onednn_isa_level():
58
    return core._check_not_sync_onednn_isa_level()
59

60

61
class TestDynDisp(unittest.TestCase):
62
    def test_manual_select_kernel(self):
63
        env_isa = get_ipex_isa_env_setting()
64
        cur_isa = get_currnet_isa_level()
65
        max_bin_isa = get_highest_binary_support_isa_level()
66
        max_cpu_isa = get_highest_cpu_support_isa_level()
67

68
        expected_isa_val = min(get_isa_val(max_bin_isa), get_isa_val(max_cpu_isa))
69

70
        if env_isa is not None:
71
            expected_isa_val = min(get_isa_val(env_isa), expected_isa_val)
72

73
        actural_isa_val = get_isa_val(cur_isa)
74

75
        # Isa level and compiler version are not linear relationship.
76
        # gcc 9.4 can build avx512_vnni.
77
        # gcc 11.3 start to support avx2_vnni.
78
        self.assertTrue(actural_isa_val <= expected_isa_val)
79
        return
80

81
    def test_dyndisp_in_supported_set(self):
82
        env_isa = get_ipex_isa_env_setting()
83

84
        if env_isa is not None:
85
            return
86

87
        cur_isa = get_currnet_isa_level()
88
        expected_isa = cur_isa in supported_isa_set
89

90
        self.assertTrue(expected_isa)
91
        return
92

93
    @unittest.skipIf(
94
        check_not_sync_onednn_isa_level(), "skip this if not sync onednn isa level"
95
    )
96
    def test_ipex_set_onednn_isa_level(self):
97
        command = 'ATEN_CPU_CAPABILITY=avx2 python -c "import torch; import intel_extension_for_pytorch._C \
98
            as core; print(core._get_current_onednn_isa_level())" '
99
        with subprocess.Popen(
100
            command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
101
        ) as p:
102
            out = p.stdout.readlines()
103
            onednn_isa_level = str(out[-1], "utf-8").strip()
104
            self.assertTrue(onednn_isa_level == "AVX2")
105

106
    @unittest.skipIf(
107
        check_not_sync_onednn_isa_level(), "skip this if not sync onednn isa level"
108
    )
109
    def test_onednn_do_not_set_isa_level(self):
110
        command = 'ONEDNN_MAX_CPU_ISA=avx2 python -c "import torch; import intel_extension_for_pytorch._C \
111
            as core; print(core._get_current_isa_level().lower())" '
112
        cur_ipex_isa = get_currnet_isa_level()
113
        with subprocess.Popen(
114
            command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
115
        ) as p:
116
            out = p.stdout.readlines()
117
            cur_ipex_isa_1 = str(out[-1], "utf-8").strip()
118
            self.assertTrue(cur_ipex_isa == cur_ipex_isa_1)
119

120

121
if __name__ == "__main__":
122
    unittest.main()
123