InfiniCore/scripts/python_test.py at 5d3e3899e66a0e2a63ebfad455fc9f677a9cc4fa · InfiniTensor/InfiniCore · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import os
import subprocess
from set_env import set_env
import sys

PROJECT_DIR = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "..", "test", "infiniop")
)
os.chdir(PROJECT_DIR)


def run_tests(args):
    failed = []
    for test in [
        "add.py",
        "attention.py",
        "causal_softmax.py",
        "clip.py",
        "conv.py",
        #"dequantize_awq.py",
        "gelu.py",
        "gemm.py",
        #"layer_norm.py",
        "logsoftmax.py",
        #"lp_norm.py",
        "mul.py",
        "ones.py",
        "random_sample.py",
        "rearrange.py",
        "relu.py",
        "rms_norm.py",
        "rope.py",
        "sigmoid.py",
        #"softmax.py",
        "softplus.py",
        "sub.py",
        "swiglu.py",
        "tanh.py",
        "topkrouter.py",
        "topksoftmax.py",
        "zeros.py",
        "2dmrope.py",
        "3dmrope.py",
    ]:
        result = subprocess.run(
            f"python {test} {args} --debug", text=True, encoding="utf-8", shell=True
        )
        if result.returncode != 0:
            failed.append(test)

    return failed


if __name__ == "__main__":
    set_env()
    failed = run_tests(" ".join(sys.argv[1:]))
    if len(failed) == 0:
        print("\033[92mAll tests passed!\033[0m")
    else:
        print("\033[91mThe following tests failed:\033[0m")
        for test in failed:
            print(f"\033[91m - {test}\033[0m")
    exit(len(failed))