Port and run tests in python/test/unit/tools (#2953)

Author: yudongsi

python/test/unit/tools/test_aot.py

@@ -7,8 +7,10 @@
 import numpy as np
 
 import triton
+from triton._internal_testing import is_cuda, is_xpu
 from triton.backends.compiler import GPUTarget
 from triton.backends.nvidia.driver import include_dir, library_dirs
+from triton.backends.intel.driver import COMPILATION_HELPER
 
 kernel_utils_src = """
 import triton
@@ -97,21 +99,42 @@ def kernel(C, A, B, M, N, K,
 }"""
 
 
-def gen_kernel_library(dir, libname):
-    c_files = glob.glob(os.path.join(dir, "*.c"))
+def gen_kernel_library_xpu(dir, libname):
+    cpp_files = glob.glob(os.path.join(dir, "*.cpp"))
     subprocess.run(
-        ["gcc"] + c_files + ["-I", include_dir[0], "-c", "-fPIC"],
+        ["icpx"] + cpp_files + ["-I", COMPILATION_HELPER.include_dir[0], "-c", "-fsycl", "-fPIC"],
         check=True,
         cwd=dir,
     )
     o_files = glob.glob(os.path.join(dir, "*.o"))
 
-    command = ["gcc", *o_files, "-shared", "-o", libname]
-    for lib_dir in library_dirs():
+    command = ["icpx", "-fsycl", "-lze_loader", *o_files, "-shared", "-o", libname]
+    for lib_dir in COMPILATION_HELPER.library_dir:
         command.extend(["-L", lib_dir])
+    if COMPILATION_HELPER.libsycl_dir:
+        for lib_dir in COMPILATION_HELPER.libsycl_dir:
+            command.extend(["-L", lib_dir])
     subprocess.run(command, check=True, cwd=dir)
 
 
+def gen_kernel_library(dir, libname):
+    if is_xpu():
+        gen_kernel_library_xpu(dir, libname)
+    else:
+        c_files = glob.glob(os.path.join(dir, "*.c"))
+        subprocess.run(
+            ["gcc"] + c_files + ["-I", include_dir[0], "-c", "-fPIC"],
+            check=True,
+            cwd=dir,
+        )
+        o_files = glob.glob(os.path.join(dir, "*.o"))
+
+        command = ["gcc", *o_files, "-shared", "-o", libname]
+        for lib_dir in library_dirs():
+            command.extend(["-L", lib_dir])
+        subprocess.run(command, check=True, cwd=dir)
+
+
 def gen_test_bin(dir, M, N, K, exe="test", algo_id=0):
     test_src = f"""
 int main(int argc, char **argv) {{
@@ -171,15 +194,118 @@ def gen_test_bin(dir, M, N, K, exe="test", algo_id=0):
 }}
 """
     src = test_utils_src + test_src
-    with open(os.path.join(dir, "test.c"), "w") as file:
+    if is_xpu():
+        src = f"""
+#include "kernel.h"
+#include <assert.h>
+#include <cmath>
+#include <cstddef>
+#include <level_zero/ze_api.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <string.h>
+#include <sycl/sycl.hpp>
+
+static void write_buffer_to_csv(char *filename, int32_t *buffer, int size) {{
+    FILE *file = fopen(filename, "w");
+    if (file == NULL) {{
+        printf("Could not open file %s\\n", filename);
+        return;
+    }}
+    for (int i = 0; i < size; i++) {{
+        fprintf(file, "%d", buffer[i]);
+        if (i < size - 1) {{
+            fprintf(file, ",");
+        }}
+    }}
+    fclose(file);
+}}
+
+static void read_csv_to_buffer(char *filename, int16_t *buffer, int size) {{
+    FILE *file = fopen(filename, "r");
+    if (file == NULL) {{
+        printf("Could not open file %s\\n", filename);
+        return;
+    }}
+    int index = 0;
+    while (fscanf(file, "%hd,", &buffer[index]) != EOF && index < size) {{
+        index++;
+    }}
+    fclose(file);
+}}
+int main(int argc, char ** argv) {{
+    int M = {M}, N = {N}, K = {K};
+
+    // initialize sycl handles
+    sycl::queue q{{sycl::gpu_selector_v}};
+    sycl::ext::intel::device_ptr<sycl::float16> A =
+        sycl::malloc_device<sycl::float16>(M * K * 2, q);
+    sycl::ext::intel::device_ptr<sycl::float16> B =
+        sycl::malloc_device<sycl::float16>(K * N * 2, q);
+    sycl::ext::intel::device_ptr<sycl::float16> C =
+        sycl::malloc_device<sycl::float16>(M * N * 4, q);
+
+    // initialize input data
+    int16_t hA[M * K];
+    int16_t hB[K * N];
+    memset(hA, 0, M * K * 2);
+    memset(hB, 0, K * N * 2);
+    read_csv_to_buffer(argv[1], hA, M * K);
+    read_csv_to_buffer(argv[2], hB, K * N);
+    q.memcpy(A, hA, M * K * 2).wait();
+    q.memcpy(B, hB, K * N * 2).wait();
+
+    // launch kernel
+    load_matmul_fp16();
+    int32_t ret;
+    int algo_id = {algo_id};
+    if (algo_id == 0) {{
+        ret = matmul_fp16_default(q, C, A, B, M, N, K, N, 1, K, 1, N, 1);
+    }} else {{
+        ret = matmul_fp16(q, C, A, B, M, N, K, N, 1, K, 1, N, 1, {algo_id});
+    }}
+    if (ret != 0) fprintf(stderr, "kernel launch failed\\n");
+    assert(ret == 0);
+
+    q.wait();
+
+    // read data
+    int32_t hC[M * N];
+    memset(hC, 0, M * N * 4);
+    q.memcpy(hC, C, M * N * 4).wait();
+    write_buffer_to_csv(argv[3], hC, M * N);
+
+    // free sycl resources
+    unload_matmul_fp16();
+    sycl::free(A, q);
+    sycl::free(B, q);
+    sycl::free(C, q);
+}}
+"""
+    src_name = "test.c"
+    if is_xpu():
+        src_name = "test.cpp"
+    with open(os.path.join(dir, src_name), "w") as file:
         file.write(src)
 
-    command = ["gcc", "test.c"]
-    for inc_dir in include_dir:
-        command.extend(["-I", inc_dir])
-    for lib_dir in library_dirs():
-        command.extend(["-L", lib_dir])
-    command.extend(["-l", "cuda", "-L", dir, "-l", "kernel", "-o", exe])
+    if is_cuda():
+        command = ["gcc", "test.c"]
+        for inc_dir in include_dir:
+            command.extend(["-I", inc_dir])
+        for lib_dir in library_dirs():
+            command.extend(["-L", lib_dir])
+        command.extend(["-l", "cuda", "-L", dir, "-l", "kernel", "-o", exe])
+
+    if is_xpu():
+        command = ["icpx", "test.cpp"]
+        for inc_dir in COMPILATION_HELPER.include_dir:
+            command.extend(["-I", inc_dir])
+        for lib_dir in COMPILATION_HELPER.library_dir:
+            command.extend(["-L", lib_dir])
+        if COMPILATION_HELPER.libsycl_dir:
+            for lib_dir in COMPILATION_HELPER.libsycl_dir:
+                command.extend(["-L", lib_dir])
+        command.extend(["-fsycl", "-lze_loader", "-L", dir, "-l", "kernel", "-o", exe])
     subprocess.run(command, check=True, cwd=dir)
 
 
@@ -283,6 +409,7 @@ def test_compile_link_matmul_no_specialization():
 
     with tempfile.TemporaryDirectory() as tmp_dir:
         dtype = "fp16"
+
         BM, BN, BK = 16, 16, 16
 
         kernel_path = write_triton_kernels(tmp_dir, kernel_src, kernel_utils_src)
@@ -299,9 +426,8 @@ def test_compile_link_matmul_no_specialization():
 
         # run test case
         env = os.environ.copy()
-        env["LD_LIBRARY_PATH"] = tmp_dir
+        env["LD_LIBRARY_PATH"] = tmp_dir + ":" + env.get("LD_LIBRARY_PATH", "")
         subprocess.run(["./test", a_path, b_path, c_path], env=env, check=True, cwd=tmp_dir)
-
         # read data and compare against reference
         c = np.genfromtxt(c_path, delimiter=",", dtype=np.int32)
         c_tri = c.reshape((M, N)).view(np.float32)
@@ -330,7 +456,7 @@ def test_compile_link_matmul():
 
         # run test case
         env = os.environ.copy()
-        env["LD_LIBRARY_PATH"] = tmp_dir
+        env["LD_LIBRARY_PATH"] = tmp_dir + ":" + env.get("LD_LIBRARY_PATH", "")
         subprocess.run(["./test", a_path, b_path, c_path], env=env, check=True, cwd=tmp_dir)
 
         # read data and compare against reference
@@ -361,7 +487,7 @@ def test_launcher_has_no_available_kernel():
 
         # run test case
         env = os.environ.copy()
-        env["LD_LIBRARY_PATH"] = tmp_dir
+        env["LD_LIBRARY_PATH"] = tmp_dir + ":" + env.get("LD_LIBRARY_PATH", "")
         result = subprocess.run(
             ["./test", a_path, b_path, c_path],
             env=env,
@@ -410,7 +536,7 @@ def test_compile_link_autotune_matmul():
             gen_test_bin(tmp_dir, M, N, K, exe=test_name, algo_id=algo_id)
 
             env = os.environ.copy()
-            env["LD_LIBRARY_PATH"] = tmp_dir
+            env["LD_LIBRARY_PATH"] = tmp_dir + ":" + env.get("LD_LIBRARY_PATH", "")
             subprocess.run(
                 [f"./{test_name}", a_path, b_path, c_path],
                 check=True,
@@ -440,3 +566,21 @@ def test_ttgir_to_ptx():
         ptx = k.asm["ptx"]
         assert ".target sm_80" in ptx
         assert ".address_size 64" in ptx
+
+
+def test_ttgir_to_spv():
+    src = """
+module attributes {"ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32, "ttg.num-ctas" = 1 : i32} {
+  tt.func public @sum_kernel_0d1d(%arg0: !tt.ptr<i32>, %arg1: !tt.ptr<i32>) {
+    tt.return
+  }
+}
+"""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        kernel_path = os.path.join(tmp_dir, "empty_kernel.ttgir")
+        with open(kernel_path, "w") as fp:
+            fp.write(src)
+        k = triton.compile(kernel_path, target=triton.runtime.driver.active.get_current_target())
+        spv = k.asm['spvdis']
+        assert "OpCapability KernelAttributesINTEL" in spv
+        assert "SubgroupSize 32" in spv

python/test/unit/tools/test_disasm.py

@@ -19,3 +19,19 @@ def kernel(X, i: tl.constexpr):
     sass = h.asm["sass"]
     # check that the sass has a store instruction.
     assert "STG.E" in sass
+
+
+def test_disam_spvbin():
+    if not triton.runtime.driver.active.get_current_target().backend == "xpu":
+        pytest.skip("Test requires XPU.")
+
+    @triton.jit
+    def kernel(X, i: tl.constexpr):
+        tl.store(X, i)
+
+    x = torch.empty(1, dtype=torch.int32, device='xpu')
+    h = kernel[(1, )](x, i=12)
+    assert x[0] == 12
+    dis = h.asm["spvdis"]
+    # check that the spvdis has a store instruction.
+    assert "OpStore" in dis

python/triton/compiler/compiler.py

@@ -8,7 +8,7 @@
 from ..runtime.autotuner import OutOfResources
 from ..runtime.cache import get_cache_manager, get_dump_manager, get_override_manager
 from ..runtime.driver import driver
-from ..tools.disasm import get_sass
+from ..tools.disasm import get_sass, get_spvdis
 # TODO: this shouldn't be here
 from .code_generator import ast_to_ttir
 from pathlib import Path
@@ -175,6 +175,8 @@ def parse(full_name, ext, context):
         return Path(full_name).read_text()
     if ext == "cubin" or ext == "hsaco":
         return Path(full_name).read_bytes()
+    if ext == "spv":
+        return Path(full_name).read_bytes()
 
 
 def filter_traceback(e: BaseException):
@@ -339,6 +341,8 @@ def __missing__(self, key):
 
         if key == "sass":
             value = get_sass(self["cubin"])
+        if key == "spvdis":
+            value = get_spvdis(self["spv"])
         else:
             raise KeyError("Unknown key: '%s'" % key)