From e9f2781dcf320fc57b2325f91394a4936f387bf0 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 4 Aug 2026 19:29:51 +0800 Subject: [PATCH 1/2] fix(torch): honor handle streams in generated operators --- docs/aten-operators.md | 5 ++++ scripts/generate_torch_ops.py | 4 ++++ src/torch/stream_.h | 39 ++++++++++++++++++++++++++++++++ tests/test_abs.py | 38 +++++++++++++++++++++++++++++++ tests/test_generate_torch_ops.py | 13 +++++++++++ 5 files changed, 99 insertions(+) create mode 100644 src/torch/stream_.h diff --git a/docs/aten-operators.md b/docs/aten-operators.md index 75404888d..f69c32ba5 100644 --- a/docs/aten-operators.md +++ b/docs/aten-operators.md @@ -48,6 +48,11 @@ generated ATen wrappers. Hand-written ATen backends may use another explicit implementation index, but must avoid colliding with the operator's existing implementations. +On NVIDIA, generated wrappers temporarily install the stream stored in `Handle` +as the current ATen CUDA stream and restore the previous device and stream when +the call returns. Other PyTorch device backends continue to use their +backend-specific current stream until an external-stream bridge is provided. + ## Add a generated ATen operator 1. Make sure the target environment has PyTorch and `torchgen` installed. diff --git a/scripts/generate_torch_ops.py b/scripts/generate_torch_ops.py index ad6e085f2..96460445e 100644 --- a/scripts/generate_torch_ops.py +++ b/scripts/generate_torch_ops.py @@ -1271,6 +1271,9 @@ def _generate_torch_method_source(name: str, op: Op) -> str: conversion_lines = [] out_device_index = f"{op.out_params[0].api_name}.device().index()" conversion_lines.append(f" const auto device_index = {out_device_index};") + conversion_lines.append( + " const detail::TorchStreamGuard stream_guard{stream_, device_index};" + ) def _optional_aten_type(param: Param) -> str: return _NULLOPT_BY_TYPE[param.aten_type].removesuffix("{}") @@ -1534,6 +1537,7 @@ class Operator<{op_type}, kDev, {slot}> : public {op_type} {{ _TORCH_SOURCE_TEMPLATE = """\ #include "torch/{name}/{name}.h" +#include "torch/stream_.h" #include "torch/tensor_.h" namespace infini::ops {{ diff --git a/src/torch/stream_.h b/src/torch/stream_.h new file mode 100644 index 000000000..e79567d3b --- /dev/null +++ b/src/torch/stream_.h @@ -0,0 +1,39 @@ +#ifndef INFINI_OPS_TORCH_STREAM__H_ +#define INFINI_OPS_TORCH_STREAM__H_ + +#ifdef WITH_NVIDIA +#include +#include +#include +#endif + +#include "device.h" + +namespace infini::ops::detail { + +template +class TorchStreamGuard { + public: + TorchStreamGuard(void*, int) {} +}; + +#ifdef WITH_NVIDIA +template <> +class TorchStreamGuard { + public: + TorchStreamGuard(void* stream, int device_index) + : device_guard_{static_cast(device_index)}, + stream_guard_{c10::cuda::getStreamFromExternal( + reinterpret_cast(stream), + static_cast(device_index))} {} + + private: + c10::cuda::CUDAGuard device_guard_; + + c10::cuda::CUDAStreamGuard stream_guard_; +}; +#endif + +} // namespace infini::ops::detail + +#endif diff --git a/tests/test_abs.py b/tests/test_abs.py index 27781cc48..1b6cc378f 100644 --- a/tests/test_abs.py +++ b/tests/test_abs.py @@ -21,6 +21,8 @@ (torch.bfloat16, 1e-2, 5e-3), ) +_PYTORCH_SLOT = 8 + @pytest.mark.auto_act_and_assert @pytest.mark.parametrize("shape, input_strides, out_strides", _SHAPE_CASES) @@ -63,3 +65,39 @@ def _torch_abs(input, out): out.copy_(torch.abs(input)) return out + + +def test_abs_torch_backend_uses_handle_stream(device): + if device != "cuda": + pytest.skip("CUDA stream coverage requires the NVIDIA backend") + + if _PYTORCH_SLOT not in infini.ops.Abs.active_implementation_indices(device): + pytest.skip("PyTorch backend slot 8 is not active on CUDA") + + input = torch.full((4096,), -1.0, device=device) + out = torch.full_like(input, torch.nan) + infini.ops.abs( + input, + out, + implementation_index=_PYTORCH_SLOT, + ) + out.fill_(torch.nan) + stream = torch.cuda.Stream() + torch.cuda.synchronize() + + with torch.cuda.stream(stream): + torch.cuda._sleep(1_000_000_000) + + infini.ops.abs( + input, + out, + stream=stream.cuda_stream, + implementation_index=_PYTORCH_SLOT, + ) + + torch.cuda.default_stream().synchronize() + snapshot = out.clone() + assert torch.isnan(snapshot).all() + + stream.synchronize() + torch.testing.assert_close(out, input.abs()) diff --git a/tests/test_generate_torch_ops.py b/tests/test_generate_torch_ops.py index c7debc32a..b353531a9 100644 --- a/tests/test_generate_torch_ops.py +++ b/tests/test_generate_torch_ops.py @@ -96,6 +96,19 @@ def test_schema_self_param_renders_as_input_in_public_cpp_api(): assert "at::_softmax_out(at_out, at_self" in source +def test_generated_torch_source_installs_handle_stream_guard(): + module = _load_generator_module() + op = module._parse_func("abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)") + + method = module._generate_torch_method_source("abs", op) + source = module._generate_torch_source("abs", [op]) + + assert ( + "detail::TorchStreamGuard stream_guard{stream_, device_index};" in method + ) + assert '#include "torch/stream_.h"' in source + + def test_optional_tensor_params_are_exposed_and_forwarded_to_aten(): module = _load_generator_module() op = module._parse_func( From f99e11e925cbf42f7d19c6a8cd6e8343bf7190a0 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 5 Aug 2026 16:09:47 +0800 Subject: [PATCH 2/2] fix(torch): extend handle stream bridge across backends --- CMakeLists.txt | 47 ++++++++++++++++++++++++++- docs/aten-operators.md | 15 ++++++--- src/CMakeLists.txt | 6 ++-- src/torch/stream_.h | 74 +++++++++++++++++++++++++++++++++++------- tests/test_abs.py | 60 ++++++++++++++++++++++------------ 5 files changed, 162 insertions(+), 40 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 293b15dce..2aea0cffa 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -274,6 +274,47 @@ if(WITH_TORCH) -Wl,--no-as-needed ${TORCH_CUDA_LIB} ${C10_CUDA_LIB} -Wl,--as-needed) endif() + if(WITH_METAX) + find_library(C10_CUDA_LIB c10_cuda HINTS ${_torch_lib_dirs} REQUIRED) + find_library(MACA_TORCH_RUNTIME_LIB runtime_cu + HINTS "$ENV{MACA_PATH}/lib" REQUIRED) + # The C10 external-stream bridge and its CUDA-to-MACA adapter both + # provide symbols referenced directly by generated Torch sources. + list(APPEND TORCH_LIBRARIES + ${C10_CUDA_LIB} ${MACA_TORCH_RUNTIME_LIB}) + endif() + + if(WITH_ASCEND) + execute_process( + COMMAND ${_TORCH_PYTHON} -c + "import importlib.util, os; s = importlib.util.find_spec('torch_npu'); print(os.path.dirname(s.origin) if s and s.origin else '')" + OUTPUT_VARIABLE TORCH_NPU_PACKAGE_DIR + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE _torch_npu_result + ) + + if(NOT _torch_npu_result EQUAL 0 OR NOT TORCH_NPU_PACKAGE_DIR) + message(FATAL_ERROR + "`WITH_ASCEND` and `WITH_TORCH` require `torch_npu`.") + endif() + + set(TORCH_NPU_INCLUDE_DIR "${TORCH_NPU_PACKAGE_DIR}/include") + set(TORCH_NPU_LIBRARY_DIR "${TORCH_NPU_PACKAGE_DIR}/lib") + list(APPEND TORCH_NPU_INCLUDE_DIRS + "${TORCH_NPU_INCLUDE_DIR}" + "${TORCH_NPU_INCLUDE_DIR}/third_party/acl/inc" + "${TORCH_NPU_INCLUDE_DIR}/third_party/hccl/inc") + list(APPEND TORCH_INCLUDE_DIRS ${TORCH_NPU_INCLUDE_DIRS}) + list(APPEND TORCH_RUNTIME_DIRS "${TORCH_NPU_LIBRARY_DIR}") + + find_library(TORCH_NPU_LIB torch_npu + HINTS "${TORCH_NPU_LIBRARY_DIR}" REQUIRED) + # Retain torch_npu so its PrivateUse1 kernels and stream APIs remain + # available even when the generated wrappers only call ATen symbols. + list(APPEND TORCH_LIBRARIES + -Wl,--no-as-needed ${TORCH_NPU_LIB} -Wl,--as-needed) + endif() + # `auditwheel`-repaired `torch` wheels bundle transitive dependencies # (e.g. `libgfortran-.so`, `libopenblasp-.so`) in a sibling # `torch.libs/` directory that `library_paths()` does not return. When @@ -429,10 +470,14 @@ if(WITH_METAX) # Normally can be found at: `/opt/maca/`. set(MACA_PATH $ENV{MACA_PATH}) + set(MACA_INCLUDE_DIRS + "${MACA_PATH}/include" + "${MACA_PATH}/include/mcr" + "${MACA_PATH}/tools/cu-bridge/include") set(CMAKE_C_COMPILER ${CMAKE_CURRENT_SOURCE_DIR}/scripts/mxcc_wrapper.sh) set(CMAKE_CXX_COMPILER ${CMAKE_CURRENT_SOURCE_DIR}/scripts/mxcc_wrapper.sh) - include_directories("${MACA_PATH}/include") + include_directories(${MACA_INCLUDE_DIRS}) link_directories("${MACA_PATH}/lib") # Libraries: mcruntime / mcdnn / mcblas. diff --git a/docs/aten-operators.md b/docs/aten-operators.md index f69c32ba5..c9d76cf5e 100644 --- a/docs/aten-operators.md +++ b/docs/aten-operators.md @@ -48,10 +48,17 @@ generated ATen wrappers. Hand-written ATen backends may use another explicit implementation index, but must avoid colliding with the operator's existing implementations. -On NVIDIA, generated wrappers temporarily install the stream stored in `Handle` -as the current ATen CUDA stream and restore the previous device and stream when -the call returns. Other PyTorch device backends continue to use their -backend-specific current stream until an external-stream bridge is provided. +Generated wrappers obtain each backend's native stream type from InfiniRT's +`Runtime::Stream`. On NVIDIA and MetaX, they temporarily install the +stream stored in `Handle` as the current ATen CUDA stream. On Ascend, they use +the equivalent NPU stream guard when the installed `torch_npu` provides its +external-stream API. Each guard restores the previous device and stream when +the call returns. When `Handle` carries no stream, generated calls leave +PyTorch's current stream unchanged. + +Older `torch_npu` releases without external-stream support continue to use the +current NPU stream because they cannot represent an arbitrary `aclrtStream`. +CPU and PyTorch device backends without a bridge also keep their current stream. ## Add a generated ATen operator diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 1edeebebd..2b5c0ae67 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -254,7 +254,7 @@ if(WITH_METAX) target_compile_options(infiniops PRIVATE "-x" "maca") target_sources(infiniops PRIVATE ${METAX_SOURCES}) - target_include_directories(infiniops PUBLIC "${MACA_PATH}/include") + target_include_directories(infiniops PUBLIC ${MACA_INCLUDE_DIRS}) target_link_libraries(infiniops PUBLIC ${MACA_RUNTIME_LIB} ${MACA_DNN_LIB} @@ -721,7 +721,9 @@ if(WITH_TORCH) set(_torch_vendor_include_flags "") if(WITH_METAX) - list(APPEND _torch_vendor_include_flags "-I${MACA_PATH}/include") + foreach(_dir ${MACA_INCLUDE_DIRS}) + list(APPEND _torch_vendor_include_flags "-I${_dir}") + endforeach() endif() if(WITH_MOORE) list(APPEND _torch_vendor_include_flags "-I${MUSA_ROOT}/include") diff --git a/src/torch/stream_.h b/src/torch/stream_.h index e79567d3b..f1614f4e1 100644 --- a/src/torch/stream_.h +++ b/src/torch/stream_.h @@ -1,13 +1,24 @@ #ifndef INFINI_OPS_TORCH_STREAM__H_ #define INFINI_OPS_TORCH_STREAM__H_ -#ifdef WITH_NVIDIA -#include +#include + +#include + +#if defined(WITH_NVIDIA) || defined(WITH_METAX) #include -#include +#endif + +#ifdef WITH_ASCEND +#include + +#if __has_include() +#define INFINI_OPS_HAS_TORCH_NPU_EXTERNAL_STREAM 1 +#endif #endif #include "device.h" +#include "runtime.h" namespace infini::ops::detail { @@ -17,23 +28,62 @@ class TorchStreamGuard { TorchStreamGuard(void*, int) {} }; +#if defined(WITH_NVIDIA) || defined(WITH_METAX) +template +class CudaTorchStreamGuard { + public: + CudaTorchStreamGuard(void* stream, int device_index) { + if (stream == nullptr) return; + + stream_guard_.emplace(c10::cuda::getStreamFromExternal( + reinterpret_cast::Stream>(stream), + static_cast(device_index))); + } + + private: + std::optional stream_guard_; +}; +#endif + #ifdef WITH_NVIDIA template <> -class TorchStreamGuard { +class TorchStreamGuard + : public CudaTorchStreamGuard { public: - TorchStreamGuard(void* stream, int device_index) - : device_guard_{static_cast(device_index)}, - stream_guard_{c10::cuda::getStreamFromExternal( - reinterpret_cast(stream), - static_cast(device_index))} {} + using CudaTorchStreamGuard::CudaTorchStreamGuard; +}; +#endif - private: - c10::cuda::CUDAGuard device_guard_; +#ifdef WITH_METAX +template <> +class TorchStreamGuard + : public CudaTorchStreamGuard { + public: + using CudaTorchStreamGuard::CudaTorchStreamGuard; +}; +#endif - c10::cuda::CUDAStreamGuard stream_guard_; +#ifdef INFINI_OPS_HAS_TORCH_NPU_EXTERNAL_STREAM +template <> +class TorchStreamGuard { + public: + TorchStreamGuard(void* stream, int device_index) { + if (stream == nullptr) return; + + stream_guard_.emplace(c10_npu::getStreamFromExternal( + reinterpret_cast::Stream>(stream), + static_cast(device_index))); + } + + private: + std::optional stream_guard_; }; #endif } // namespace infini::ops::detail +#ifdef INFINI_OPS_HAS_TORCH_NPU_EXTERNAL_STREAM +#undef INFINI_OPS_HAS_TORCH_NPU_EXTERNAL_STREAM +#endif + #endif diff --git a/tests/test_abs.py b/tests/test_abs.py index 1b6cc378f..70cf9fa41 100644 --- a/tests/test_abs.py +++ b/tests/test_abs.py @@ -67,35 +67,53 @@ def _torch_abs(input, out): return out -def test_abs_torch_backend_uses_handle_stream(device): - if device != "cuda": - pytest.skip("CUDA stream coverage requires the NVIDIA backend") - +@pytest.mark.smoke +@pytest.mark.parametrize("stream_source", ("current", "handle")) +def test_abs_torch_backend_uses_selected_stream(device, stream_source): + stream_apis = { + "cuda": (torch.cuda, "cuda_stream"), + "npu": (getattr(torch, "npu", None), "npu_stream"), + } + accelerator, stream_attr = stream_apis.get(device, (None, None)) + + if ( + stream_source == "handle" + and device == "npu" + and accelerator is not None + and not hasattr(accelerator, "ExternalStream") + ): + pytest.skip("The installed torch_npu does not support external streams") + + if accelerator is None or not hasattr(accelerator, "_sleep"): + pytest.skip("The device does not expose the required stream test APIs") if _PYTORCH_SLOT not in infini.ops.Abs.active_implementation_indices(device): - pytest.skip("PyTorch backend slot 8 is not active on CUDA") + pytest.skip(f"PyTorch backend slot 8 is not active on {device}") input = torch.full((4096,), -1.0, device=device) out = torch.full_like(input, torch.nan) - infini.ops.abs( - input, - out, - implementation_index=_PYTORCH_SLOT, - ) + + def call_abs(**kwargs): + infini.ops.abs( + input, + out, + implementation_index=_PYTORCH_SLOT, + **kwargs, + ) + + call_abs() out.fill_(torch.nan) - stream = torch.cuda.Stream() - torch.cuda.synchronize() + stream = accelerator.Stream() + accelerator.synchronize() - with torch.cuda.stream(stream): - torch.cuda._sleep(1_000_000_000) + with accelerator.stream(stream): + accelerator._sleep(1_000_000_000) + if stream_source == "current": + call_abs() - infini.ops.abs( - input, - out, - stream=stream.cuda_stream, - implementation_index=_PYTORCH_SLOT, - ) + if stream_source == "handle": + call_abs(stream=getattr(stream, stream_attr)) - torch.cuda.default_stream().synchronize() + accelerator.default_stream().synchronize() snapshot = out.clone() assert torch.isnan(snapshot).all()