Skip to content

Commit

Permalink
Lazy compute compile options (#2029)
Browse files Browse the repository at this point in the history
Closes #2027

---------

Signed-off-by: Anatoly Myachev <anatoly.myachev@intel.com>
Co-authored-by: Pavel Chekin <pavel.chekin@intel.com>
  • Loading branch information
anmyachev and pbchekin authored Aug 28, 2024
1 parent ee594a9 commit 1ab4017
Showing 1 changed file with 37 additions and 9 deletions.
46 changes: 37 additions & 9 deletions third_party/intel/backend/driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,15 @@
import shutil
import tempfile
from pathlib import Path
from functools import cached_property

from triton.runtime.build import _build
from triton.runtime.cache import get_cache_manager
from triton.backends.compiler import GPUTarget
from triton.backends.driver import DriverBase
from packaging.version import Version
from packaging.specifiers import SpecifierSet

ze_root = os.getenv("ZE_PATH", default="/usr/local")
include_dir = [os.path.join(ze_root, "include")]


def find_sycl(include_dir: list[str]) -> tuple[list[str], list[str]]:
"""
Expand Down Expand Up @@ -65,13 +64,41 @@ def find_sycl(include_dir: list[str]) -> tuple[list[str], list[str]]:
return include_dir, library_dir


include_dir, library_dir = find_sycl(include_dir)
class CompilationHelper:
_library_dir: list[str]
_include_dir: list[str]

def __init__(self):
self._library_dir = None
self._include_dir = None
self.libraries = ['ze_loader', 'sycl']

@cached_property
def _compute_compilation_options_lazy(self):
ze_root = os.getenv("ZE_PATH", default="/usr/local")
include_dir = [os.path.join(ze_root, "include")]

include_dir, library_dir = find_sycl(include_dir)

dirname = os.path.dirname(os.path.realpath(__file__))
include_dir += [os.path.join(dirname, "include")]
library_dir += [os.path.join(dirname, "lib")]

self._library_dir = library_dir
self._include_dir = include_dir

@cached_property
def library_dir(self) -> list[str]:
self._compute_compilation_options_lazy
return self._library_dir

@cached_property
def include_dir(self) -> list[str]:
self._compute_compilation_options_lazy
return self._include_dir

dirname = os.path.dirname(os.path.realpath(__file__))
include_dir += [os.path.join(dirname, "include")]

library_dir += [os.path.join(dirname, "lib")]
libraries = ['ze_loader', 'sycl']
compilation_helper = CompilationHelper()


def compile_module_from_src(src, name):
Expand All @@ -83,7 +110,8 @@ def compile_module_from_src(src, name):
src_path = os.path.join(tmpdir, "main.cpp")
with open(src_path, "w") as f:
f.write(src)
so = _build(name, src_path, tmpdir, library_dir, include_dir, libraries)
so = _build(name, src_path, tmpdir, compilation_helper.library_dir, compilation_helper.include_dir,
compilation_helper.libraries)
with open(so, "rb") as f:
cache_path = cache.put(f.read(), f"{name}.so", binary=True)
import importlib.util
Expand Down

0 comments on commit 1ab4017

Please sign in to comment.