Enable tools/test_aot.py for Windows+icpx (#3629)

I plan to add msvc support separately.

---------

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

python/test/unit/tools/test_aot.py

@@ -3,6 +3,8 @@
 import subprocess
 import sys
 import tempfile
+import shutil
+import sysconfig
 
 import numpy as np
 
@@ -99,19 +101,34 @@ def kernel(C, A, B, M, N, K,
 }"""
 
 
+def select_compiler():
+    gxx = shutil.which("g++")
+    icpx = shutil.which("icpx")
+    cl = shutil.which("cl")
+    cxx = (icpx or cl) if os.name == "nt" else (icpx or gxx)
+    if cxx is None:
+        raise RuntimeError("Failed to find C++ compiler. Please specify via CXX environment variable.")
+    return cxx
+
+
 def gen_kernel_library_xpu(dir, libname):
     cpp_files = glob.glob(os.path.join(dir, "*.cpp"))
-    subprocess.run(
-        ["g++"] + cpp_files + ["-I" + include_dir for include_dir in COMPILATION_HELPER.include_dir] + ["-c", "-fPIC"],
-        check=True,
-        cwd=dir,
-    )
+    cxx = select_compiler()
+    command = [cxx] + cpp_files + ["-I" + include_dir for include_dir in COMPILATION_HELPER.include_dir
+                                   ] + ["-c", "-fPIC" if os.name != "nt" else "-Wno-deprecated-declarations"]
+    subprocess.run(command, check=True, cwd=dir)
     o_files = glob.glob(os.path.join(dir, "*.o"))
 
-    subprocess.run(["g++"] + [*o_files, "-shared", "-o", libname] +
-                   ["-L" + library_dir for library_dir in COMPILATION_HELPER.library_dir] +
-                   ["-L" + dir
-                    for dir in COMPILATION_HELPER.libsycl_dir] + ["-lsycl", "-lze_loader"], check=True, cwd=dir)
+    extra_link_args = []
+    if "icpx" in cxx and os.name == "nt":
+        libname_without_ext = libname.split(".")[0]
+        extra_link_args = [f"/IMPLIB:{libname_without_ext}.lib"]
+
+    command = [cxx] + [*o_files, "-shared", "-o", libname] + [
+        "-L" + library_dir for library_dir in COMPILATION_HELPER.library_dir
+    ] + ["-L" + dir for dir in COMPILATION_HELPER.libsycl_dir
+         ] + ["-lsycl8" if os.name == "nt" else "-lsycl", "-lze_loader"] + extra_link_args
+    subprocess.run(command, check=True, cwd=dir)
 
 
 def gen_kernel_library(dir, libname):
@@ -133,6 +150,8 @@ def gen_kernel_library(dir, libname):
 
 
 def gen_test_bin(dir, M, N, K, exe="test", algo_id=0):
+    exe_extension = sysconfig.get_config_var("EXE")
+    exe = exe + exe_extension
     test_src = f"""
 int main(int argc, char **argv) {{
   int M = {M}, N = {N}, K = {K};
@@ -294,15 +313,18 @@ def gen_test_bin(dir, M, N, K, exe="test", algo_id=0):
         command.extend(["-l", "cuda", "-L", dir, "-l", "kernel", "-o", exe])
 
     if is_xpu():
-        command = ["g++", "test.cpp"]
+        cxx = select_compiler()
+        command = [cxx, "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(["-lsycl", "-lze_loader", "-L", dir, "-l", "kernel", "-o", exe])
+        if os.name == "nt":
+            command.extend(["-Wno-deprecated-declarations"])
+        command.extend(["-lsycl8" if os.name == "nt" else "-lsycl", "-lze_loader", "-L", dir, "-lkernel", "-o", exe])
     subprocess.run(command, check=True, cwd=dir)
 
 
@@ -415,7 +437,7 @@ def test_compile_link_matmul_no_specialization():
 
         # compile test case
         M, N, K = 16, 16, 16
-        gen_kernel_library(tmp_dir, "libkernel.so")
+        gen_kernel_library(tmp_dir, "libkernel.so" if os.name != "nt" else "kernel.dll")
         gen_test_bin(tmp_dir, M, N, K)
 
         # initialize test data
@@ -424,7 +446,7 @@ def test_compile_link_matmul_no_specialization():
         # run test case
         env = os.environ.copy()
         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)
+        subprocess.run([os.path.join(tmp_dir, "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)
@@ -445,7 +467,7 @@ def test_compile_link_matmul():
 
         # compile test case
         M, N, K = 16, 16, 16
-        gen_kernel_library(tmp_dir, "libkernel.so")
+        gen_kernel_library(tmp_dir, "libkernel.so" if os.name != "nt" else "kernel.dll")
         gen_test_bin(tmp_dir, M, N, K)
 
         # initialize test data
@@ -454,7 +476,7 @@ def test_compile_link_matmul():
         # run test case
         env = os.environ.copy()
         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)
+        subprocess.run([os.path.join(tmp_dir, "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)
@@ -476,7 +498,7 @@ def test_launcher_has_no_available_kernel():
 
         # compile test case
         M, N, K = 16, 16, 16
-        gen_kernel_library(tmp_dir, "libkernel.so")
+        gen_kernel_library(tmp_dir, "libkernel.so" if os.name != "nt" else "kernel.dll")
         gen_test_bin(tmp_dir, M, N, K)
 
         # initialize test data
@@ -486,15 +508,16 @@ def test_launcher_has_no_available_kernel():
         env = os.environ.copy()
         env["LD_LIBRARY_PATH"] = tmp_dir + ":" + env.get("LD_LIBRARY_PATH", "")
         result = subprocess.run(
-            ["./test", a_path, b_path, c_path],
+            [os.path.join(tmp_dir, "test"), a_path, b_path, c_path],
             env=env,
             cwd=tmp_dir,
             capture_output=True,
             text=True,
         )
 
         # It should fail since the launcher requires all the strides be 1 while they are not.
-        assert result.returncode == -6
+        # On windows: 3221226505 == 0xc0000409: STATUS_STACK_BUFFER_OVERRUN
+        assert result.returncode == -6 if os.name != "nt" else 0xc0000409
         assert "kernel launch failed" in result.stderr
 
 
@@ -519,7 +542,7 @@ def test_compile_link_autotune_matmul():
 
         link_aot_kernels(tmp_dir)
 
-        gen_kernel_library(tmp_dir, "libkernel.so")
+        gen_kernel_library(tmp_dir, "libkernel.so" if os.name != "nt" else "kernel.dll")
 
         # compile test case
         M, N, K = 64, 64, 64
@@ -535,7 +558,7 @@ def test_compile_link_autotune_matmul():
             env = os.environ.copy()
             env["LD_LIBRARY_PATH"] = tmp_dir + ":" + env.get("LD_LIBRARY_PATH", "")
             subprocess.run(
-                [f"./{test_name}", a_path, b_path, c_path],
+                [os.path.join(tmp_dir, test_name), a_path, b_path, c_path],
                 check=True,
                 cwd=tmp_dir,
                 env=env,

python/triton/tools/link.py

@@ -159,10 +159,10 @@ def make_global_decl(meta: KernelLinkerMeta) -> str:
         """
     if is_xpu():
         return f"""
-int32_t {meta.orig_kernel_name}_default(sycl::queue &stream, {gen_signature_with_full_args(meta)});
-int32_t {meta.orig_kernel_name}(sycl::queue &stream, {gen_signature_with_full_args(meta)}, int algo_id);
-void load_{meta.orig_kernel_name}();
-void unload_{meta.orig_kernel_name}();
+EXPORT_FUNC int32_t {meta.orig_kernel_name}_default(sycl::queue &stream, {gen_signature_with_full_args(meta)});
+EXPORT_FUNC int32_t {meta.orig_kernel_name}(sycl::queue &stream, {gen_signature_with_full_args(meta)}, int algo_id);
+EXPORT_FUNC void load_{meta.orig_kernel_name}();
+EXPORT_FUNC void unload_{meta.orig_kernel_name}();
         """
 
 
@@ -172,7 +172,7 @@ def make_default_algo_kernel(meta: KernelLinkerMeta) -> str:
         src = f"CUresult {meta.orig_kernel_name}_default(CUstream stream, {gen_signature_with_full_args(meta)}){{\n"
         src += f"  return {meta.orig_kernel_name}(stream, {', '.join(meta.arg_names)}, 0);\n"
     if is_xpu():
-        src = f"int32_t {meta.orig_kernel_name}_default(sycl::queue &stream, {gen_signature_with_full_args(meta)}){{\n"
+        src = f"EXPORT_FUNC int32_t {meta.orig_kernel_name}_default(sycl::queue &stream, {gen_signature_with_full_args(meta)}){{\n"
         src += f"  return {meta.orig_kernel_name}(stream, {', '.join(meta.arg_names)}, 0);\n"
     src += "}\n"
     return src
@@ -245,7 +245,7 @@ def make_kernel_meta_const_dispatcher(meta: KernelLinkerMeta) -> str:
     if is_cuda():
         src = f"CUresult {meta.orig_kernel_name}(CUstream stream, {gen_signature_with_full_args(meta)}, int algo_id){{\n"
     if is_xpu():
-        src = f"int32_t {meta.orig_kernel_name}(sycl::queue &stream, {gen_signature_with_full_args(meta)}, int algo_id){{\n"
+        src = f"EXPORT_FUNC int32_t {meta.orig_kernel_name}(sycl::queue &stream, {gen_signature_with_full_args(meta)}, int algo_id){{\n"
     src += f"  assert (algo_id < (int)sizeof({meta.orig_kernel_name}_kernels));\n"
     if is_cuda():
         src += f"  return {meta.orig_kernel_name}_kernels[algo_id](stream, {', '.join(meta.arg_names)});\n"
@@ -273,7 +273,7 @@ def make_func_pointers(names: str, meta: KernelLinkerMeta) -> str:
 def make_kernel_load_def(names: str, meta: KernelLinkerMeta) -> str:
     src = ""
     for mode in ["load", "unload"]:
-        src += f"void {mode}_{meta.orig_kernel_name}(void){{\n"
+        src += f"EXPORT_FUNC void {mode}_{meta.orig_kernel_name}(void){{\n"
         for name in names:
             src += f"  {mode}_{name}();\n"
         src += "}\n\n"
@@ -344,6 +344,11 @@ def make_get_num_algos_def(meta: KernelLinkerMeta) -> str:
             out += "#include <sycl/sycl.hpp>\n"
             out += "#include <stdint.h>\n"
             out += "#include <stdio.h>\n"
+            out += "#if defined(_WIN32)\n"
+            out += "#define EXPORT_FUNC __declspec(dllexport)\n"
+            out += "#else\n"
+            out += "#define EXPORT_FUNC\n"
+            out += "#endif\n"
         out += "\n".join(algo_decls)
         out += "\n"
         out += get_num_algos_decl
@@ -386,6 +391,11 @@ def make_get_num_algos_def(meta: KernelLinkerMeta) -> str:
             out += "#include <stdint.h>\n"
             out += "#include <assert.h>\n"
             out += "#include <cstdint>\n"
+            out += "#if defined(_WIN32)\n"
+            out += "#define EXPORT_FUNC __declspec(dllexport)\n"
+            out += "#else\n"
+            out += "#define EXPORT_FUNC\n"
+            out += "#endif\n"
             out += "\n"
             out += "\n".join(defs)
             out += "\n"