updaates

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

Author: anmyachev

benchmarks/triton_kernels_benchmark/benchmark_driver.py

@@ -6,14 +6,16 @@
 from triton.runtime.cache import get_cache_manager
 from triton.runtime.build import _build, quiet
 from triton._utils import parse_list_string
-from triton.backends.intel.driver import compile_module_from_src
+from triton.backends.intel.driver import compile_module_from_src, COMPILATION_HELPER
 
 import torch
 
 # ------------------------
 # Utils
 # ------------------------
 
+COMPILATION_HELPER.inject_pytorch_dep()
+
 
 class XPUUtils:
 

third_party/intel/backend/driver.py

@@ -68,6 +68,10 @@ def find_sycl(include_dir: list[str]) -> tuple[list[str], Optional[str]]:
 class CompilationHelper:
     _library_dir: list[str]
     _include_dir: list[str]
+    libraries: list[str]
+
+    # for benchmarks
+    _build_with_pytorch_dep: bool = False
 
     def __init__(self):
         self._library_dir = None
@@ -77,6 +81,12 @@ def __init__(self):
         if os.name != "nt":
             self.libraries += ["sycl"]
 
+    def inject_pytorch_dep(self):
+        # must be called before any cached properties (if pytorch is needed)
+        if self._build_with_pytorch_dep is False:
+            self._build_with_pytorch_dep = True
+            self.libraries += ['torch']
+
     @cached_property
     def _compute_compilation_options_lazy(self):
         ze_root = os.getenv("ZE_PATH", default="/usr/local")
@@ -91,9 +101,18 @@ def _compute_compilation_options_lazy(self):
 
         dirname = os.path.dirname(os.path.realpath(__file__))
         include_dir += [os.path.join(dirname, "include")]
-        # TODO: do we need this?
         library_dir += [os.path.join(dirname, "lib")]
 
+        if self._build_with_pytorch_dep:
+            import torch
+
+            torch_path = torch.utils.cmake_prefix_path
+            include_dir += [
+                os.path.join(torch_path, "../../include"),
+                os.path.join(torch_path, "../../include/torch/csrc/api/include"),
+            ]
+            library_dir += [os.path.join(torch_path, "../../lib")]
+
         self._library_dir = library_dir
         self._include_dir = include_dir
 
@@ -113,7 +132,7 @@ def libsycl_dir(self) -> Optional[str]:
         return self._libsycl_dir
 
 
-compilation_helper = CompilationHelper()
+COMPILATION_HELPER = CompilationHelper()
 
 
 def compile_module_from_src(src, name):
@@ -127,10 +146,10 @@ def compile_module_from_src(src, name):
             with open(src_path, "w") as f:
                 f.write(src)
             extra_compiler_args = []
-            if compilation_helper.libsycl_dir:
-                extra_compiler_args += ['-Wl,-rpath,' + compilation_helper.libsycl_dir]
-            so = _build(name, src_path, tmpdir, compilation_helper.library_dir, compilation_helper.include_dir,
-                        compilation_helper.libraries, extra_compile_args=extra_compiler_args)
+            if COMPILATION_HELPER.libsycl_dir:
+                extra_compiler_args += ['-Wl,-rpath,' + COMPILATION_HELPER.libsycl_dir]
+            so = _build(name, src_path, tmpdir, COMPILATION_HELPER.library_dir, COMPILATION_HELPER.include_dir,
+                        COMPILATION_HELPER.libraries, extra_compile_args=extra_compiler_args)
             with open(so, "rb") as f:
                 cache_path = cache.put(f.read(), file_name, binary=True)
     import importlib.util