fix(torch): honor handle streams in generated operators

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-<hash>.so`, `libopenblasp-<hash>.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.

docs/aten-operators.md

@@ -48,6 +48,18 @@ generated ATen wrappers. Hand-written ATen backends may use another explicit
 implementation index, but must avoid colliding with the operator's existing
 implementations.
 
+Generated wrappers obtain each backend's native stream type from InfiniRT's
+`Runtime<kDev>::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
 
 1. Make sure the target environment has PyTorch and `torchgen` installed.

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<kDev> 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 {{

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")

src/torch/stream_.h

@@ -0,0 +1,89 @@
+#ifndef INFINI_OPS_TORCH_STREAM__H_
+#define INFINI_OPS_TORCH_STREAM__H_
+
+#include <c10/core/StreamGuard.h>
+
+#include <optional>
+
+#if defined(WITH_NVIDIA) || defined(WITH_METAX)
+#include <c10/cuda/CUDAStream.h>
+#endif
+
+#ifdef WITH_ASCEND
+#include <torch_npu/csrc/core/npu/NPUStream.h>
+
+#if __has_include(<torch_npu/csrc/core/npu/NPUStreamUtils.h>)
+#define INFINI_OPS_HAS_TORCH_NPU_EXTERNAL_STREAM 1
+#endif
+#endif
+
+#include "device.h"
+#include "runtime.h"
+
+namespace infini::ops::detail {
+
+template <Device::Type kDev>
+class TorchStreamGuard {
+ public:
+  TorchStreamGuard(void*, int) {}
+};
+
+#if defined(WITH_NVIDIA) || defined(WITH_METAX)
+template <Device::Type kDev>
+class CudaTorchStreamGuard {
+ public:
+  CudaTorchStreamGuard(void* stream, int device_index) {
+    if (stream == nullptr) return;
+
+    stream_guard_.emplace(c10::cuda::getStreamFromExternal(
+        reinterpret_cast<typename Runtime<kDev>::Stream>(stream),
+        static_cast<c10::DeviceIndex>(device_index)));
+  }
+
+ private:
+  std::optional<c10::StreamGuard> stream_guard_;
+};
+#endif
+
+#ifdef WITH_NVIDIA
+template <>
+class TorchStreamGuard<Device::Type::kNvidia>
+    : public CudaTorchStreamGuard<Device::Type::kNvidia> {
+ public:
+  using CudaTorchStreamGuard<Device::Type::kNvidia>::CudaTorchStreamGuard;
+};
+#endif
+
+#ifdef WITH_METAX
+template <>
+class TorchStreamGuard<Device::Type::kMetax>
+    : public CudaTorchStreamGuard<Device::Type::kMetax> {
+ public:
+  using CudaTorchStreamGuard<Device::Type::kMetax>::CudaTorchStreamGuard;
+};
+#endif
+
+#ifdef INFINI_OPS_HAS_TORCH_NPU_EXTERNAL_STREAM
+template <>
+class TorchStreamGuard<Device::Type::kAscend> {
+ public:
+  TorchStreamGuard(void* stream, int device_index) {
+    if (stream == nullptr) return;
+
+    stream_guard_.emplace(c10_npu::getStreamFromExternal(
+        reinterpret_cast<Runtime<Device::Type::kAscend>::Stream>(stream),
+        static_cast<c10::DeviceIndex>(device_index)));
+  }
+
+ private:
+  std::optional<c10::StreamGuard> 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

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,57 @@ def _torch_abs(input, out):
     out.copy_(torch.abs(input))
 
     return out
+
+
+@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(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)
+
+    def call_abs(**kwargs):
+        infini.ops.abs(
+            input,
+            out,
+            implementation_index=_PYTORCH_SLOT,
+            **kwargs,
+        )
+
+    call_abs()
+    out.fill_(torch.nan)
+    stream = accelerator.Stream()
+    accelerator.synchronize()
+
+    with accelerator.stream(stream):
+        accelerator._sleep(1_000_000_000)
+        if stream_source == "current":
+            call_abs()
+
+    if stream_source == "handle":
+        call_abs(stream=getattr(stream, stream_attr))
+
+    accelerator.default_stream().synchronize()
+    snapshot = out.clone()
+    assert torch.isnan(snapshot).all()
+
+    stream.synchronize()
+    torch.testing.assert_close(out, input.abs())

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<kDev> 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(