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 75404888d..c9d76cf5e 100644 --- a/docs/aten-operators.md +++ b/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::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. 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/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 new file mode 100644 index 000000000..f1614f4e1 --- /dev/null +++ b/src/torch/stream_.h @@ -0,0 +1,89 @@ +#ifndef INFINI_OPS_TORCH_STREAM__H_ +#define INFINI_OPS_TORCH_STREAM__H_ + +#include + +#include + +#if defined(WITH_NVIDIA) || defined(WITH_METAX) +#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 { + +template +class TorchStreamGuard { + public: + 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 + : public CudaTorchStreamGuard { + public: + using CudaTorchStreamGuard::CudaTorchStreamGuard; +}; +#endif + +#ifdef WITH_METAX +template <> +class TorchStreamGuard + : public CudaTorchStreamGuard { + public: + using CudaTorchStreamGuard::CudaTorchStreamGuard; +}; +#endif + +#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 27781cc48..70cf9fa41 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,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()) 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(