This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 28308486 [Feature] [XPU] Enable Intel GPU Support (#888)
28308486 is described below
commit 283084866fdfcb849501dfe4003882e52f83a62b
Author: Polisetty V R K Jyothendra Varma <[email protected]>
AuthorDate: Wed Nov 4 21:27:28 2026 +0530
[Feature] [XPU] Enable Intel GPU Support (#888)
This PR enables Intel GPU Support
Signed-off-by: P V R K Jyothendra Varma
<[email protected]>
---
.gitignore | 1 +
docs/concepts/func_module.rst | 6 +++
docs/concepts/tensor.rst | 3 +-
include/tvm/ffi/container/tensor.h | 3 +-
include/tvm/ffi/device.h | 1 +
include/tvm/ffi/extra/c_env_api.h | 5 ++
python/tvm_ffi/cpp/extension.py | 100 ++++++++++++++++++++++++++++++-------
python/tvm_ffi/cython/device.pxi | 2 +
python/tvm_ffi/stream.py | 37 ++++++++++++--
tests/cpp/test_device.cc | 7 ++-
tests/python/test_device.py | 4 ++
tests/python/test_stream.py | 33 ++++++++++++
12 files changed, 177 insertions(+), 25 deletions(-)
diff --git a/.gitignore b/.gitignore
index fa91f214..448dd6f2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -15,6 +15,7 @@ __pycache__/
*.so
build/
+build_test/
*.ll
.npm
diff --git a/docs/concepts/func_module.rst b/docs/concepts/func_module.rst
index 508b4f7b..5f1b9fce 100644
--- a/docs/concepts/func_module.rst
+++ b/docs/concepts/func_module.rst
@@ -347,6 +347,12 @@ C++/CUDA source files and loads them as a module in one
step:
Essentially, :py:func:`tvm_ffi.cpp.load` is a convenience function that
JIT-compiles the source
files and loads the resulting library as a :py:class:`tvm_ffi.Module`.
+Device code is passed through ``cuda_files`` / ``cuda_sources`` and compiled
for the GPU backend given
+by ``backend``: ``"cuda"`` (``nvcc``), ``"hip"`` (``hipcc``), or ``"sycl"``
(``icpx -fsycl``, for Intel
+XPU). Without ``backend``, it is taken from the ``TVM_FFI_GPU_BACKEND``
environment variable or from
+the installed PyTorch build. SYCL is single-source C++, so with ``"sycl"``
every C++ source is
+compiled with ``icpx``.
+
System Library
~~~~~~~~~~~~~~
diff --git a/docs/concepts/tensor.rst b/docs/concepts/tensor.rst
index 26a806b4..8a68ab69 100644
--- a/docs/concepts/tensor.rst
+++ b/docs/concepts/tensor.rst
@@ -235,7 +235,8 @@ This is equivalent to the following PyTorch C++ code:
**Auto-Update Stream Context.** When converting framework tensors via
:py:func:`tvm_ffi.from_dlpack`,
TVM-FFI automatically updates the stream context to match the device of the
converted tensor.
For example, when converting a PyTorch tensor on ``torch.device('cuda:3')``,
TVM-FFI automatically
-captures the stream from :py:func:`torch.cuda.current_stream(device='cuda:3')`.
+captures the stream from
:py:func:`torch.cuda.current_stream(device='cuda:3')`. For ``torch.xpu``
tensors,
+the ``sycl::queue*`` of :py:func:`torch.xpu.current_stream` is captured the
same way.
**Set Stream Context.** Use :py:func:`tvm_ffi.use_torch_stream` or
:py:func:`tvm_ffi.use_raw_stream`
to manually set the stream context when automatic detection is insufficient.
diff --git a/include/tvm/ffi/container/tensor.h
b/include/tvm/ffi/container/tensor.h
index 2c25ba44..db687487 100644
--- a/include/tvm/ffi/container/tensor.h
+++ b/include/tvm/ffi/container/tensor.h
@@ -47,7 +47,8 @@ class Tensor;
*/
inline bool IsDirectAddressDevice(const DLDevice& device) {
return device.device_type <= kDLCUDAHost || device.device_type ==
kDLCUDAManaged ||
- device.device_type == kDLROCM || device.device_type == kDLROCMHost;
+ device.device_type == kDLROCM || device.device_type == kDLROCMHost ||
+ device.device_type == kDLOneAPI;
}
/*!
diff --git a/include/tvm/ffi/device.h b/include/tvm/ffi/device.h
index b2c7fd5f..3a34006c 100644
--- a/include/tvm/ffi/device.h
+++ b/include/tvm/ffi/device.h
@@ -47,6 +47,7 @@ TVM_FFI_INLINE static std::optional<DLDeviceType>
TryParseDLDeviceType(std::stri
if (name == "vpi") return kDLVPI;
if (name == "rocm") return kDLROCM;
if (name == "ext_dev") return kDLExtDev;
+ if (name == "oneapi" || name == "xpu") return kDLOneAPI;
if (name == "hexagon") return kDLHexagon;
if (name == "wgpu" || name == "webgpu") return kDLWebGPU;
if (name == "maia") return kDLMAIA;
diff --git a/include/tvm/ffi/extra/c_env_api.h
b/include/tvm/ffi/extra/c_env_api.h
index 6192db54..fba240e8 100644
--- a/include/tvm/ffi/extra/c_env_api.h
+++ b/include/tvm/ffi/extra/c_env_api.h
@@ -37,6 +37,11 @@ extern "C" {
// ----------------------------------------------------------------------------
/*!
* \brief The type of the stream handle.
+ *
+ * The meaning of the handle depends on the device type:
+ * - kDLCUDA: cudaStream_t
+ * - kDLROCM: hipStream_t
+ * - kDLOneAPI: non-owning sycl::queue*, valid for the duration of the call
*/
typedef void* TVMFFIStreamHandle;
diff --git a/python/tvm_ffi/cpp/extension.py b/python/tvm_ffi/cpp/extension.py
index 0bf06f07..f37f279e 100644
--- a/python/tvm_ffi/cpp/extension.py
+++ b/python/tvm_ffi/cpp/extension.py
@@ -23,6 +23,7 @@ import hashlib
import logging
import os
import platform
+import re
import shutil
import subprocess
import sys
@@ -36,7 +37,8 @@ from tvm_ffi.module import Module, load_module
from tvm_ffi.utils import FileLock
IS_WINDOWS = sys.platform == "win32"
-BACKEND_STR = Literal["cuda", "hip"]
+BACKEND_STR = Literal["cuda", "hip", "sycl"]
+_BACKENDS = ("cuda", "hip", "sycl")
logger = logging.getLogger(__name__)
@@ -57,18 +59,18 @@ def _find_compiler(name: str, *home_vars: str) -> str |
None:
@functools.lru_cache
def _detect_gpu_backend() -> BACKEND_STR:
- """Auto-detect whether to use CUDA or HIP (ROCm), defaulting to CUDA.
+ """Auto-detect whether to use CUDA, HIP (ROCm), or SYCL, defaulting to
CUDA.
Resolution order: the ``TVM_FFI_GPU_BACKEND`` override, then PyTorch build
- signals (``torch.version.hip`` / ``torch.version.cuda``), then a
discoverable
+ signals (``torch.version.hip`` / ``torch.version.cuda`` /
``torch.version.xpu``), then a discoverable
``hipcc`` when no ``nvcc`` can be found. Compilers are resolved from
``PATH``
and from the ``CUDA_HOME``/``CUDA_PATH`` and ``ROCM_HOME``/``ROCM_PATH``
toolkit homes, so an env-specified CUDA install is honored even when
``nvcc``
is off ``PATH``. The mere existence of ``/opt/rocm`` is not treated as
- evidence of HIP.
+ evidence of HIP. SYCL is never inferred from the presence of ``icpx``
alone.
"""
backend = os.environ.get("TVM_FFI_GPU_BACKEND", "").lower()
- if backend in ("cuda", "hip"):
+ if backend in _BACKENDS:
return backend # type: ignore[return-value]
try:
import torch # noqa: PLC0415
@@ -78,6 +80,8 @@ def _detect_gpu_backend() -> BACKEND_STR:
return "hip"
if getattr(version, "cuda", None):
return "cuda"
+ if getattr(version, "xpu", None):
+ return "sycl"
except Exception:
pass
if _find_compiler("nvcc", "CUDA_HOME", "CUDA_PATH") is None and (
@@ -89,13 +93,15 @@ def _detect_gpu_backend() -> BACKEND_STR:
def _resolve_gpu_backend(backend: str | None) -> BACKEND_STR:
if backend is not None:
- if backend in ("cuda", "hip"):
+ if backend in _BACKENDS:
return backend # type: ignore[return-value]
- raise ValueError(f"Invalid backend: {backend}. Supported backends are
'cuda' and 'hip'.")
+ raise ValueError(
+ f"Invalid backend: {backend}. Supported backends are 'cuda',
'hip', and 'sycl'."
+ )
return _detect_gpu_backend()
-def _hash_sources(
+def _hash_sources( # noqa: PLR0913
cpp_source: str | None,
cuda_source: str | None,
cpp_files: Sequence[str] | None,
@@ -106,8 +112,12 @@ def _hash_sources(
extra_ldflags: Sequence[str],
extra_include_paths: Sequence[str],
embed_cubin: Mapping[str, bytes] | None = None,
+ backend: str | None = None,
) -> str:
- """Generate a unique hash for the given sources and functions."""
+ """Generate a unique hash for the given sources and functions.
+
+ ``backend`` is only hashed when given, so existing cache keys are
unchanged.
+ """
m = hashlib.sha256()
def _hash(obj: Any) -> None:
@@ -146,6 +156,8 @@ def _hash_sources(
embed_cubin,
)
)
+ if backend is not None:
+ _hash(backend)
return m.hexdigest()[:16]
@@ -302,6 +314,32 @@ def _get_rocm_target() -> list[str]:
)
+def _find_sycl_compiler() -> str:
+ """Find the oneAPI DPC++ compiler (``icpx``)."""
+ icpx = _find_compiler("icpx", "CMPLR_ROOT")
+ if icpx is None and os.environ.get("ONEAPI_ROOT"):
+ icpx = shutil.which(
+ "icpx", path=str(Path(os.environ["ONEAPI_ROOT"]) / "compiler" /
"latest" / "bin")
+ )
+ if icpx is None:
+ raise RuntimeError(
+ "Could not find the SYCL compiler icpx. Please source the oneAPI
setvars.sh script "
+ "or set the CMPLR_ROOT environment variable."
+ )
+ return icpx
+
+
+def _get_sycl_targets() -> list[str]:
+ """Get the SYCL ahead-of-time target flags, or none to JIT compile from
SPIR-V."""
+ targets = os.environ.get("TVM_FFI_SYCL_TARGETS", "").replace(",", "
").split()
+ for target in targets: # e.g., "intel_gpu_pvc intel_gpu_bmg_g21"
+ if not re.fullmatch(r"[A-Za-z0-9_]+", target):
+ raise ValueError(f"Invalid SYCL target: {target}")
+ if not targets:
+ return []
+ return ["-fsycl-targets={}".format(",".join(targets))]
+
+
def _run_command_in_dev_prompt(
args: list[str],
cwd: str | os.PathLike[str],
@@ -402,6 +440,13 @@ def _generate_ninja_build( # noqa: PLR0915, PLR0912
object_mode = False
output_name = f"{name}{'.dll' if IS_WINDOWS else '.so'}"
has_cuda_sources = any(Path(s).suffix.lower() == ".cu" for s in sources)
+ with_sycl = backend == "sycl"
+ if with_sycl and has_cuda_sources:
+ raise ValueError(
+ "CUDA (.cu) sources cannot be built with the sycl backend, use C++
sources instead."
+ )
+ if with_sycl and IS_WINDOWS:
+ raise NotImplementedError("The sycl backend is not yet supported on
Windows")
with_hip = backend == "hip"
with_cuda = backend == "cuda"
with_backend = with_hip or with_cuda or has_cuda_sources
@@ -461,6 +506,11 @@ def _generate_ninja_build( # noqa: PLR0915, PLR0912
f"-L{Path(rocm_home) / 'lib'!s}",
"-lamdhip64",
]
+ if with_sycl:
+ # SYCL is single-source C++, so every C++ source is compiled and
linked with icpx.
+ sycl_flags = ["-fsycl", *_get_sycl_targets()]
+ default_cxxflags = ["-std=c++20", "-fPIC", "-O2", *sycl_flags]
+ default_ldflags += sycl_flags
if with_cuda:
default_cuda_cflags = ["-Xcompiler", "-fPIC", *default_cuda_cflags]
default_cuda_cflags += [_get_cuda_target()]
@@ -472,6 +522,9 @@ def _generate_ninja_build( # noqa: PLR0915, PLR0912
extra_cflags_list = [flag.strip() for flag in extra_cflags]
cflags = default_cflags + extra_cflags_list
cxxflags = default_cxxflags + extra_cflags_list
+ if with_sycl:
+ # SYCL device code is compiled by the C++ rule
+ cxxflags += [flag.strip() for flag in extra_cuda_cflags]
cuda_cflags = default_cuda_cflags + [flag.strip() for flag in
extra_cuda_cflags]
ldflags = default_ldflags + [flag.strip() for flag in extra_ldflags]
include_paths = default_include_paths + [
@@ -489,7 +542,10 @@ def _generate_ninja_build( # noqa: PLR0915, PLR0912
with_c = any(Path(s).suffix.lower() == ".c" for s in sources)
ninja: list[str] = []
ninja.append("ninja_required_version = 1.3")
- ninja.append("cxx = {}".format(os.environ.get("CXX", "cl" if IS_WINDOWS
else "c++")))
+ if with_sycl:
+ ninja.append(f"cxx = {_find_sycl_compiler()}")
+ else:
+ ninja.append("cxx = {}".format(os.environ.get("CXX", "cl" if
IS_WINDOWS else "c++")))
ninja.append("cxxflags = {}".format(" ".join(cxxflags)))
if with_c:
ninja.append("cc = {}".format(os.environ.get("CC", "cl" if IS_WINDOWS
else "cc")))
@@ -728,7 +784,7 @@ def _str_seq2list(seq: Sequence[str] | str | None) ->
list[str]:
return list(seq)
-def _build_impl( # noqa: PLR0913
+def _build_impl( # noqa: PLR0912, PLR0913
name: str,
sources: Sequence[str] | str | None,
extra_cflags: Sequence[str] | None,
@@ -747,7 +803,12 @@ def _build_impl( # noqa: PLR0913
assert source_path_list, "sources must be provided."
has_cuda = any(Path(p).suffix.lower() == ".cu" for p in source_path_list)
- resolved_backend = _resolve_gpu_backend(backend) if has_cuda else None
+ if has_cuda:
+ resolved_backend: str | None = _resolve_gpu_backend(backend)
+ elif backend == "sycl":
+ resolved_backend = "sycl"
+ else:
+ resolved_backend = None
extra_ldflags_list = list(extra_ldflags) if extra_ldflags is not None else
[]
extra_cflags_list = list(extra_cflags) if extra_cflags is not None else []
extra_cuda_cflags_list = list(extra_cuda_cflags) if extra_cuda_cflags is
not None else []
@@ -767,6 +828,7 @@ def _build_impl( # noqa: PLR0913
extra_ldflags_list,
extra_include_paths_list,
embed_cubin,
+ backend="sycl" if resolved_backend == "sycl" else None,
)
build_dir = Path(cache_dir).expanduser() / f"{name}_{source_hash}"
else:
@@ -916,7 +978,7 @@ def build_inline( # noqa: PLR0913
See the `examples/cubin_launcher` directory for examples how to use
cubin launcher to launch CUBIN kernels in TVM-FFI.
backend
- The GPU backend to use. It can be "cuda" or "hip".
+ The GPU backend to use. It can be "cuda", "hip", or "sycl".
If not specified, the backend will be automatically determined based
on the available GPU and the provided source code.
output
@@ -981,6 +1043,9 @@ def build_inline( # noqa: PLR0913
cuda_source = "\n".join(cuda_source_list)
with_backend = bool(cuda_source_list)
del cuda_source_list
+ with_sycl = with_backend and _resolve_gpu_backend(backend) == "sycl"
+ if with_sycl:
+ backend = "sycl"
extra_ldflags_list = list(extra_ldflags) if extra_ldflags is not None else
[]
extra_cflags_list = list(extra_cflags) if extra_cflags is not None else []
@@ -1018,6 +1083,7 @@ def build_inline( # noqa: PLR0913
extra_ldflags_list,
extra_include_paths_list,
embed_cubin,
+ backend="sycl" if with_sycl else None,
)
build_dir = Path(cache_dir).expanduser() / f"{name}_{source_hash}"
else:
@@ -1025,7 +1091,7 @@ def build_inline( # noqa: PLR0913
build_dir.mkdir(parents=True, exist_ok=True)
cpp_file = str((build_dir / "main.cpp").resolve())
- cuda_file = str((build_dir / "cuda.cu").resolve())
+ cuda_file = str((build_dir / ("sycl.cpp" if with_sycl else
"cuda.cu")).resolve())
with FileLock(str(build_dir / "lock")):
# write source files if they do not already exist
@@ -1144,7 +1210,7 @@ def load_inline( # noqa: PLR0913
for the duration of the process.
backend
- The GPU backend to use. It can be "cuda" or "hip".
+ The GPU backend to use. It can be "cuda", "hip", or "sycl".
If not specified, the backend will be automatically determined based
on the available GPU and the provided source code.
Returns
@@ -1309,7 +1375,7 @@ def build( # noqa: PLR0913
specify the cache directory.
backend
- The GPU backend to use. It can be "cuda" or "hip".
+ The GPU backend to use. It can be "cuda", "hip", or "sycl".
If not specified, the backend will be automatically determined based
on the available GPU and the provided source code.
output
@@ -1469,7 +1535,7 @@ def load( # noqa: PLR0913
for the duration of the process.
backend
- The GPU backend to use. It can be "cuda" or "hip".
+ The GPU backend to use. It can be "cuda", "hip", or "sycl".
If not specified, the backend will be automatically determined based
on the available GPU and the provided source code.
Returns
diff --git a/python/tvm_ffi/cython/device.pxi b/python/tvm_ffi/cython/device.pxi
index 0f3fa863..642bac68 100644
--- a/python/tvm_ffi/cython/device.pxi
+++ b/python/tvm_ffi/cython/device.pxi
@@ -124,6 +124,8 @@ cdef class Device:
"vpi": DLDeviceType.kDLVPI,
"rocm": DLDeviceType.kDLROCM,
"ext_dev": DLDeviceType.kDLExtDev,
+ "oneapi": DLDeviceType.kDLOneAPI,
+ "xpu": DLDeviceType.kDLOneAPI,
"hexagon": DLDeviceType.kDLHexagon,
"wgpu": DLDeviceType.kDLWebGPU,
"webgpu": DLDeviceType.kDLWebGPU,
diff --git a/python/tvm_ffi/stream.py b/python/tvm_ffi/stream.py
index a8053785..0cfc7a42 100644
--- a/python/tvm_ffi/stream.py
+++ b/python/tvm_ffi/stream.py
@@ -19,6 +19,7 @@
from __future__ import annotations
+import functools
from ctypes import c_void_p
from typing import Any
@@ -68,6 +69,20 @@ class StreamContext:
)
[email protected]
+def _xpu_dlpack_device_from_torch_device_index(torch_index: int) ->
core.Device:
+ """Return the FFI device of torch ``xpu:<torch_index>``.
+
+ The kDLOneAPI device_id is the global SYCL index that torch's DLPack
exporter reports,
+ which may differ from the torch XPU index, so the mapping comes from torch.
+ """
+ # TODO(tvm-ffi-team): remove this mapping once the DLPack, PyTorch, SYCL
and Level Zero
+ # device indices are aligned, and use device(str(current_stream.device))
like CUDA.
+ import torch # noqa: PLC0415
+
+ return device(*torch._C._torchDeviceToDLDevice(torch.device("xpu",
torch_index)))
+
+
try:
import torch
@@ -82,10 +97,17 @@ try:
"""Enter both Torch and FFI stream contexts."""
if self.torch_context:
self.torch_context.__enter__()
- current_stream = torch.cuda.current_stream()
- self.ffi_context = StreamContext(
- device(str(current_stream.device)), current_stream.cuda_stream
- )
+ current_stream = torch.get_device_module().current_stream()
+ if hasattr(current_stream, "cuda_stream"):
+ self.ffi_context = StreamContext(
+ device(str(current_stream.device)),
current_stream.cuda_stream
+ )
+ else:
+ # XPU
+ self.ffi_context = StreamContext(
+
_xpu_dlpack_device_from_torch_device_index(current_stream.device.index),
+ current_stream.sycl_queue,
+ )
self.ffi_context.__enter__()
return self
@@ -103,7 +125,8 @@ try:
Parameters
----------
context
- The wrapped torch stream or cuda graph.
+ The wrapped torch stream (``torch.cuda.stream`` /
``torch.xpu.stream``)
+ or cuda graph.
Returns
-------
@@ -122,6 +145,10 @@ try:
with tvm_ffi.use_torch_stream(torch.cuda.graph(g)):
...
+ s = torch.xpu.Stream()
+ with tvm_ffi.use_torch_stream(torch.xpu.stream(s)):
+ ...
+
Note
----
When working with a raw ``cudaStream_t`` handle, use
diff --git a/tests/cpp/test_device.cc b/tests/cpp/test_device.cc
index 3eb80859..a4051eb3 100644
--- a/tests/cpp/test_device.cc
+++ b/tests/cpp/test_device.cc
@@ -42,7 +42,7 @@ TEST(Device, AnyConversionWithString) {
{"cpu", kDLCPU}, {"cuda", kDLCUDA}, {"opencl", kDLOpenCL},
{"vulkan", kDLVulkan},
{"metal", kDLMetal}, {"mps", kDLMetal}, {"vpi", kDLVPI},
{"rocm", kDLROCM},
{"ext_dev", kDLExtDev}, {"hexagon", kDLHexagon}, {"wgpu", kDLWebGPU},
{"webgpu", kDLWebGPU},
- {"maia", kDLMAIA}, {"trn", kDLTrn},
+ {"maia", kDLMAIA}, {"trn", kDLTrn}, {"oneapi", kDLOneAPI},
{"xpu", kDLOneAPI},
};
for (const TestCase& test_case : test_cases) {
SCOPED_TRACE(test_case.name);
@@ -61,6 +61,11 @@ TEST(Device, AnyConversionWithString) {
EXPECT_EQ(opencl.device_type, kDLOpenCL);
EXPECT_EQ(opencl.device_id, 2);
+ AnyView xpu_view = "xpu:2";
+ DLDevice xpu = xpu_view.cast<DLDevice>();
+ EXPECT_EQ(xpu.device_type, kDLOneAPI);
+ EXPECT_EQ(xpu.device_id, 2);
+
AnyView max_index_view = "cuda:2147483647";
DLDevice max_index = max_index_view.cast<DLDevice>();
EXPECT_EQ(max_index.device_type, kDLCUDA);
diff --git a/tests/python/test_device.py b/tests/python/test_device.py
index 0b622c72..c770245c 100644
--- a/tests/python/test_device.py
+++ b/tests/python/test_device.py
@@ -50,6 +50,10 @@ def test_device_from_str() -> None:
("cuda:0", DLDeviceType.kDLCUDA, 0),
("cuda:3", DLDeviceType.kDLCUDA, 3),
("metal:2", DLDeviceType.kDLMetal, 2),
+ ("oneapi", DLDeviceType.kDLOneAPI, 0),
+ ("oneapi:1", DLDeviceType.kDLOneAPI, 1),
+ ("xpu", DLDeviceType.kDLOneAPI, 0),
+ ("xpu:2", DLDeviceType.kDLOneAPI, 2),
],
)
def test_device_dlpack_device_type(
diff --git a/tests/python/test_stream.py b/tests/python/test_stream.py
index ffdd122e..2d33e382 100644
--- a/tests/python/test_stream.py
+++ b/tests/python/test_stream.py
@@ -159,3 +159,36 @@ def test_torch_graph() -> None:
x = x + 1
run_with_gpu_lock(run_and_check)
+
+
[email protected](torch is None or not torch.xpu.is_available(),
reason="Requires torch and XPU")
+def test_torch_xpu_stream() -> None:
+ assert torch is not None
+ mod = tvm_ffi.cpp.load_inline(
+ name="check_xpu_stream",
+ cpp_sources="""
+ uint64_t get_tensor_stream(tvm::ffi::TensorView x) {
+ return reinterpret_cast<uint64_t>(
+ TVMFFIEnvGetStream(x.device().device_type,
x.device().device_id));
+ }
+ """,
+ functions=["get_tensor_stream"],
+ )
+
+ def run_and_check() -> None:
+ assert torch is not None
+ device_id = torch.xpu.current_device()
+ x = torch.empty(4, device=f"xpu:{device_id}")
+ stream_1 = torch.xpu.Stream(device_id)
+ stream_2 = torch.xpu.Stream(device_id)
+ # torch.xpu tensors carry the current torch stream via the DLPack
exchange API
+ assert mod.get_tensor_stream(x) ==
torch.xpu.current_stream().sycl_queue
+ with torch.xpu.stream(stream_1):
+ assert mod.get_tensor_stream(x) == stream_1.sycl_queue
+ # use_torch_stream sets the FFI stream of the current torch stream
+ with tvm_ffi.use_torch_stream(torch.xpu.stream(stream_2)):
+ assert torch.xpu.current_stream() == stream_2
+ assert mod.get_tensor_stream(x) == stream_2.sycl_queue
+ assert tvm_ffi.get_raw_stream(tvm_ffi.from_dlpack(x).device) ==
stream_2.sycl_queue
+
+ run_with_gpu_lock(run_and_check)
