From f2392c4a9753bf16e78a39352939552cd1414f26 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 11 Aug 2026 15:39:10 +0800 Subject: [PATCH 1/2] refactor(torch): move C10 helpers out of linked --- src/linked/torch/c10.h | 13 ------------- src/linked/torch/cambricon/ops/rms_norm/apex.h | 2 +- src/linked/torch/cambricon/ops/silu_and_mul/apex.h | 2 +- src/linked/torch/metax/ops/rms_norm/vllm.h | 2 +- src/linked/torch/metax/ops/silu_and_mul/vllm.h | 2 +- src/linked/torch/moore/ops/rms_norm/vllm.h | 2 +- src/linked/torch/moore/ops/silu_and_mul/vllm.h | 2 +- .../torch/nvidia/ops/awq_marlin_repack/vllm.h | 2 +- .../nvidia/ops/flash_attn_varlen_func/flash_attn.h | 2 +- .../nvidia/ops/flash_attn_with_kvcache/flash_attn.h | 2 +- .../torch/nvidia/ops/get_cutlass_moe_mm_data/vllm.h | 2 +- .../torch/nvidia/ops/gptq_marlin_repack/vllm.h | 2 +- src/linked/torch/nvidia/ops/grouped_topk/vllm.h | 2 +- src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.h | 2 +- .../torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.h | 2 +- src/linked/torch/nvidia/ops/topk_sigmoid/vllm.h | 2 +- src/linked/torch/nvidia/ops/topk_softmax/vllm.h | 2 +- src/torch/c10.h | 13 +++++++++++++ src/{linked => }/torch/cambricon/c10.h | 12 ++++++------ src/{linked => }/torch/metax/c10.h | 12 ++++++------ src/{linked => }/torch/moore/c10.h | 12 ++++++------ src/{linked => }/torch/nvidia/c10.h | 12 ++++++------ 22 files changed, 53 insertions(+), 53 deletions(-) delete mode 100644 src/linked/torch/c10.h create mode 100644 src/torch/c10.h rename src/{linked => }/torch/cambricon/c10.h (64%) rename src/{linked => }/torch/metax/c10.h (69%) rename src/{linked => }/torch/moore/c10.h (70%) rename src/{linked => }/torch/nvidia/c10.h (69%) diff --git a/src/linked/torch/c10.h b/src/linked/torch/c10.h deleted file mode 100644 index 3a205ce00..000000000 --- a/src/linked/torch/c10.h +++ /dev/null @@ -1,13 +0,0 @@ -#ifndef INFINI_OPS_LINKED_TORCH_C10_H_ -#define INFINI_OPS_LINKED_TORCH_C10_H_ - -#include "device.h" - -namespace infini::ops::linked::torch { - -template -struct C10; - -} // namespace infini::ops::linked::torch - -#endif // INFINI_OPS_LINKED_TORCH_C10_H_ diff --git a/src/linked/torch/cambricon/ops/rms_norm/apex.h b/src/linked/torch/cambricon/ops/rms_norm/apex.h index f89d614ab..d689c7d03 100644 --- a/src/linked/torch/cambricon/ops/rms_norm/apex.h +++ b/src/linked/torch/cambricon/ops/rms_norm/apex.h @@ -1,8 +1,8 @@ #ifndef INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_RMS_NORM_APEX_H_ #define INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_RMS_NORM_APEX_H_ -#include "linked/torch/cambricon/c10.h" #include "linked/torch/ops/rms_norm.h" +#include "torch/cambricon/c10.h" namespace infini::ops::linked::torch::cambricon { diff --git a/src/linked/torch/cambricon/ops/silu_and_mul/apex.h b/src/linked/torch/cambricon/ops/silu_and_mul/apex.h index 27161fbb9..5c9439b75 100644 --- a/src/linked/torch/cambricon/ops/silu_and_mul/apex.h +++ b/src/linked/torch/cambricon/ops/silu_and_mul/apex.h @@ -1,8 +1,8 @@ #ifndef INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_SILU_AND_MUL_APEX_H_ #define INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_SILU_AND_MUL_APEX_H_ -#include "linked/torch/cambricon/c10.h" #include "linked/torch/ops/silu_and_mul.h" +#include "torch/cambricon/c10.h" namespace infini::ops::linked::torch::cambricon { diff --git a/src/linked/torch/metax/ops/rms_norm/vllm.h b/src/linked/torch/metax/ops/rms_norm/vllm.h index e4923cd51..3b81e51a7 100644 --- a/src/linked/torch/metax/ops/rms_norm/vllm.h +++ b/src/linked/torch/metax/ops/rms_norm/vllm.h @@ -1,8 +1,8 @@ #ifndef INFINI_OPS_LINKED_TORCH_METAX_OPS_RMS_NORM_VLLM_H_ #define INFINI_OPS_LINKED_TORCH_METAX_OPS_RMS_NORM_VLLM_H_ -#include "linked/torch/metax/c10.h" #include "linked/torch/ops/rms_norm.h" +#include "torch/metax/c10.h" namespace infini::ops::linked::torch::metax { diff --git a/src/linked/torch/metax/ops/silu_and_mul/vllm.h b/src/linked/torch/metax/ops/silu_and_mul/vllm.h index 94e0503a5..b7fb695b5 100644 --- a/src/linked/torch/metax/ops/silu_and_mul/vllm.h +++ b/src/linked/torch/metax/ops/silu_and_mul/vllm.h @@ -1,8 +1,8 @@ #ifndef INFINI_OPS_LINKED_TORCH_METAX_OPS_SILU_AND_MUL_VLLM_H_ #define INFINI_OPS_LINKED_TORCH_METAX_OPS_SILU_AND_MUL_VLLM_H_ -#include "linked/torch/metax/c10.h" #include "linked/torch/ops/silu_and_mul.h" +#include "torch/metax/c10.h" namespace infini::ops::linked::torch::metax { diff --git a/src/linked/torch/moore/ops/rms_norm/vllm.h b/src/linked/torch/moore/ops/rms_norm/vllm.h index f7a4ef3f7..c721b1c2e 100644 --- a/src/linked/torch/moore/ops/rms_norm/vllm.h +++ b/src/linked/torch/moore/ops/rms_norm/vllm.h @@ -1,8 +1,8 @@ #ifndef INFINI_OPS_LINKED_TORCH_MOORE_OPS_RMS_NORM_VLLM_H_ #define INFINI_OPS_LINKED_TORCH_MOORE_OPS_RMS_NORM_VLLM_H_ -#include "linked/torch/moore/c10.h" #include "linked/torch/ops/rms_norm.h" +#include "torch/moore/c10.h" namespace infini::ops::linked::torch::moore { diff --git a/src/linked/torch/moore/ops/silu_and_mul/vllm.h b/src/linked/torch/moore/ops/silu_and_mul/vllm.h index c236434ae..15cbef1d9 100644 --- a/src/linked/torch/moore/ops/silu_and_mul/vllm.h +++ b/src/linked/torch/moore/ops/silu_and_mul/vllm.h @@ -1,8 +1,8 @@ #ifndef INFINI_OPS_LINKED_TORCH_MOORE_OPS_SILU_AND_MUL_VLLM_H_ #define INFINI_OPS_LINKED_TORCH_MOORE_OPS_SILU_AND_MUL_VLLM_H_ -#include "linked/torch/moore/c10.h" #include "linked/torch/ops/silu_and_mul.h" +#include "torch/moore/c10.h" namespace infini::ops::linked::torch::moore { diff --git a/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.h b/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.h index a350622dc..f40444c83 100644 --- a/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.h +++ b/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.h @@ -1,8 +1,8 @@ #ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_AWQ_MARLIN_REPACK_VLLM_H_ #define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_AWQ_MARLIN_REPACK_VLLM_H_ -#include "linked/torch/nvidia/c10.h" #include "linked/torch/ops/awq_marlin_repack.h" +#include "torch/nvidia/c10.h" namespace infini::ops::linked::torch::nvidia { diff --git a/src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.h b/src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.h index 9a02c7179..9d6ddfc4b 100644 --- a/src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.h +++ b/src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.h @@ -3,8 +3,8 @@ #include -#include "linked/torch/nvidia/c10.h" #include "linked/torch/ops/flash_attn_varlen_func.h" +#include "torch/nvidia/c10.h" namespace infini::ops::linked::torch::nvidia { diff --git a/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.h b/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.h index fddb2f73e..db52f9a49 100644 --- a/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.h +++ b/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.h @@ -1,8 +1,8 @@ #ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_ #define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_ -#include "linked/torch/nvidia/c10.h" #include "linked/torch/ops/flash_attn_with_kvcache.h" +#include "torch/nvidia/c10.h" namespace infini::ops::linked::torch::nvidia { diff --git a/src/linked/torch/nvidia/ops/get_cutlass_moe_mm_data/vllm.h b/src/linked/torch/nvidia/ops/get_cutlass_moe_mm_data/vllm.h index 52c9abba7..aa9a1acf1 100644 --- a/src/linked/torch/nvidia/ops/get_cutlass_moe_mm_data/vllm.h +++ b/src/linked/torch/nvidia/ops/get_cutlass_moe_mm_data/vllm.h @@ -1,8 +1,8 @@ #ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GET_CUTLASS_MOE_MM_DATA_VLLM_H_ #define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GET_CUTLASS_MOE_MM_DATA_VLLM_H_ -#include "linked/torch/nvidia/c10.h" #include "linked/torch/ops/get_cutlass_moe_mm_data.h" +#include "torch/nvidia/c10.h" namespace infini::ops::linked::torch::nvidia { diff --git a/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.h b/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.h index 5a92d836f..0323cccdc 100644 --- a/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.h +++ b/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.h @@ -1,8 +1,8 @@ #ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GPTQ_MARLIN_REPACK_VLLM_H_ #define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GPTQ_MARLIN_REPACK_VLLM_H_ -#include "linked/torch/nvidia/c10.h" #include "linked/torch/ops/gptq_marlin_repack.h" +#include "torch/nvidia/c10.h" namespace infini::ops::linked::torch::nvidia { diff --git a/src/linked/torch/nvidia/ops/grouped_topk/vllm.h b/src/linked/torch/nvidia/ops/grouped_topk/vllm.h index da6a1037d..5eef6bf75 100644 --- a/src/linked/torch/nvidia/ops/grouped_topk/vllm.h +++ b/src/linked/torch/nvidia/ops/grouped_topk/vllm.h @@ -3,8 +3,8 @@ #include -#include "linked/torch/nvidia/c10.h" #include "linked/torch/ops/grouped_topk.h" +#include "torch/nvidia/c10.h" namespace infini::ops::linked::torch::nvidia { diff --git a/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.h b/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.h index de361ad18..06da2b865 100644 --- a/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.h +++ b/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.h @@ -3,8 +3,8 @@ #include -#include "linked/torch/nvidia/c10.h" #include "linked/torch/ops/moe_wna16_gemm.h" +#include "torch/nvidia/c10.h" namespace infini::ops::linked::torch::nvidia { diff --git a/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.h b/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.h index f64e1ef4a..5e0bfe1b8 100644 --- a/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.h +++ b/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.h @@ -8,8 +8,8 @@ #include #include -#include "linked/torch/nvidia/c10.h" #include "linked/torch/ops/moe_wna16_marlin_gemm.h" +#include "torch/nvidia/c10.h" namespace infini::ops::linked::torch::nvidia { diff --git a/src/linked/torch/nvidia/ops/topk_sigmoid/vllm.h b/src/linked/torch/nvidia/ops/topk_sigmoid/vllm.h index c825f7b1c..7043eff9a 100644 --- a/src/linked/torch/nvidia/ops/topk_sigmoid/vllm.h +++ b/src/linked/torch/nvidia/ops/topk_sigmoid/vllm.h @@ -3,8 +3,8 @@ #include -#include "linked/torch/nvidia/c10.h" #include "linked/torch/ops/topk_sigmoid.h" +#include "torch/nvidia/c10.h" namespace infini::ops::linked::torch::nvidia { diff --git a/src/linked/torch/nvidia/ops/topk_softmax/vllm.h b/src/linked/torch/nvidia/ops/topk_softmax/vllm.h index 1463dd5a2..3e2813e46 100644 --- a/src/linked/torch/nvidia/ops/topk_softmax/vllm.h +++ b/src/linked/torch/nvidia/ops/topk_softmax/vllm.h @@ -3,8 +3,8 @@ #include -#include "linked/torch/nvidia/c10.h" #include "linked/torch/ops/topk_softmax.h" +#include "torch/nvidia/c10.h" namespace infini::ops::linked::torch::nvidia { diff --git a/src/torch/c10.h b/src/torch/c10.h new file mode 100644 index 000000000..46c684eb6 --- /dev/null +++ b/src/torch/c10.h @@ -0,0 +1,13 @@ +#ifndef INFINI_OPS_TORCH_C10_H_ +#define INFINI_OPS_TORCH_C10_H_ + +#include "device.h" + +namespace infini::ops { + +template +struct C10; + +} // namespace infini::ops + +#endif // INFINI_OPS_TORCH_C10_H_ diff --git a/src/linked/torch/cambricon/c10.h b/src/torch/cambricon/c10.h similarity index 64% rename from src/linked/torch/cambricon/c10.h rename to src/torch/cambricon/c10.h index 3d62c8d6d..ed7daaeb0 100644 --- a/src/linked/torch/cambricon/c10.h +++ b/src/torch/cambricon/c10.h @@ -1,13 +1,13 @@ -#ifndef INFINI_OPS_LINKED_TORCH_CAMBRICON_C10_H_ -#define INFINI_OPS_LINKED_TORCH_CAMBRICON_C10_H_ +#ifndef INFINI_OPS_TORCH_CAMBRICON_C10_H_ +#define INFINI_OPS_TORCH_CAMBRICON_C10_H_ #include #include #include -#include "linked/torch/c10.h" +#include "torch/c10.h" -namespace infini::ops::linked::torch { +namespace infini::ops { template <> struct C10 { @@ -22,6 +22,6 @@ struct C10 { } }; -} // namespace infini::ops::linked::torch +} // namespace infini::ops -#endif // INFINI_OPS_LINKED_TORCH_CAMBRICON_C10_H_ +#endif // INFINI_OPS_TORCH_CAMBRICON_C10_H_ diff --git a/src/linked/torch/metax/c10.h b/src/torch/metax/c10.h similarity index 69% rename from src/linked/torch/metax/c10.h rename to src/torch/metax/c10.h index a0e5daf56..c0d9d1f95 100644 --- a/src/linked/torch/metax/c10.h +++ b/src/torch/metax/c10.h @@ -1,13 +1,13 @@ -#ifndef INFINI_OPS_LINKED_TORCH_METAX_C10_H_ -#define INFINI_OPS_LINKED_TORCH_METAX_C10_H_ +#ifndef INFINI_OPS_TORCH_METAX_C10_H_ +#define INFINI_OPS_TORCH_METAX_C10_H_ #include #include #include -#include "linked/torch/c10.h" +#include "torch/c10.h" -namespace infini::ops::linked::torch { +namespace infini::ops { template <> struct C10 { @@ -23,6 +23,6 @@ struct C10 { } }; -} // namespace infini::ops::linked::torch +} // namespace infini::ops -#endif // INFINI_OPS_LINKED_TORCH_METAX_C10_H_ +#endif // INFINI_OPS_TORCH_METAX_C10_H_ diff --git a/src/linked/torch/moore/c10.h b/src/torch/moore/c10.h similarity index 70% rename from src/linked/torch/moore/c10.h rename to src/torch/moore/c10.h index 3d69221bf..f410277de 100644 --- a/src/linked/torch/moore/c10.h +++ b/src/torch/moore/c10.h @@ -1,14 +1,14 @@ -#ifndef INFINI_OPS_LINKED_TORCH_MOORE_C10_H_ -#define INFINI_OPS_LINKED_TORCH_MOORE_C10_H_ +#ifndef INFINI_OPS_TORCH_MOORE_C10_H_ +#define INFINI_OPS_TORCH_MOORE_C10_H_ #include #include #include #include -#include "linked/torch/c10.h" +#include "torch/c10.h" -namespace infini::ops::linked::torch { +namespace infini::ops { template <> struct C10 { @@ -24,6 +24,6 @@ struct C10 { } }; -} // namespace infini::ops::linked::torch +} // namespace infini::ops -#endif // INFINI_OPS_LINKED_TORCH_MOORE_C10_H_ +#endif // INFINI_OPS_TORCH_MOORE_C10_H_ diff --git a/src/linked/torch/nvidia/c10.h b/src/torch/nvidia/c10.h similarity index 69% rename from src/linked/torch/nvidia/c10.h rename to src/torch/nvidia/c10.h index a24833303..b5c4c5a91 100644 --- a/src/linked/torch/nvidia/c10.h +++ b/src/torch/nvidia/c10.h @@ -1,13 +1,13 @@ -#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_C10_H_ -#define INFINI_OPS_LINKED_TORCH_NVIDIA_C10_H_ +#ifndef INFINI_OPS_TORCH_NVIDIA_C10_H_ +#define INFINI_OPS_TORCH_NVIDIA_C10_H_ #include #include #include -#include "linked/torch/c10.h" +#include "torch/c10.h" -namespace infini::ops::linked::torch { +namespace infini::ops { template <> struct C10 { @@ -23,6 +23,6 @@ struct C10 { } }; -} // namespace infini::ops::linked::torch +} // namespace infini::ops -#endif // INFINI_OPS_LINKED_TORCH_NVIDIA_C10_H_ +#endif // INFINI_OPS_TORCH_NVIDIA_C10_H_ From 3f5b8b777b991c1c1a99aaccb1e63bc26d339096 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 11 Aug 2026 18:03:34 +0800 Subject: [PATCH 2/2] fix(torch): honor handle streams in generated operators --- CMakeLists.txt | 8 +++++ scripts/generate_torch_ops.py | 35 ++++++++++++++++++++-- src/CMakeLists.txt | 51 ++++++++++++++++++-------------- tests/test_abs.py | 45 ++++++++++++++++++++++++++++ tests/test_generate_torch_ops.py | 17 +++++++++++ 5 files changed, 132 insertions(+), 24 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index e2ada3897..e5b91fed3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -276,6 +276,14 @@ if(WITH_TORCH OR WITH_LINKED) -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) + list(APPEND TORCH_LIBRARIES + ${C10_CUDA_LIB} ${MACA_TORCH_RUNTIME_LIB}) + 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 diff --git a/scripts/generate_torch_ops.py b/scripts/generate_torch_ops.py index ad6e085f2..0c90cedad 100644 --- a/scripts/generate_torch_ops.py +++ b/scripts/generate_torch_ops.py @@ -74,6 +74,13 @@ "kHygon", ) +_C10_DEVICE_TYPES = ( + "kNvidia", + "kCambricon", + "kMetax", + "kMoore", +) + # YAML scalar-type tokens → C++ types. Reference types (e.g. `const Scalar&`) # are not used so the generated signatures match the existing hand-written # ones, which pass by value to keep pybind11 binding generation simple. @@ -1270,7 +1277,6 @@ def _generate_torch_method_source(name: str, op: Op) -> str: op_type = _op_cpp_type(name) 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};") def _optional_aten_type(param: Param) -> str: return _NULLOPT_BY_TYPE[param.aten_type].removesuffix("{}") @@ -1446,6 +1452,10 @@ def _render_arg(schema_index, p): op_type=op_type, op_call_signature=_format_signature(op), tensor_conversions="\n".join(conversion_lines), + out_device_index=out_device_index, + c10_condition=" ||\n ".join( + f"kDev == Device::Type::{dev}" for dev in _C10_DEVICE_TYPES + ), # The generated call expression resolves the right kernel via C++ # overload resolution from the argument types we pass. aten_call=aten_call, @@ -1456,6 +1466,12 @@ def _render_arg(schema_index, p): def _generate_torch_source(name: str, ops: list[Op]) -> str: op_type = _op_cpp_type(name) methods = "\n\n".join(_generate_torch_method_source(name, op) for op in ops) + c10_includes = "\n".join( + f"#ifdef WITH_{dev.removeprefix('k').upper()}\n" + f'#include "torch/{dev.removeprefix("k").lower()}/c10.h"\n' + f"#endif" + for dev in _C10_DEVICE_TYPES + ) # Guard each explicit instantiation by the matching `WITH_` macro # so a build that only enables a subset of devices does not pay the # ATen template-instantiation cost (and memory pressure) for the @@ -1471,6 +1487,7 @@ def _generate_torch_source(name: str, ops: list[Op]) -> str: return _TORCH_SOURCE_TEMPLATE.format( name=name, methods=methods, + c10_includes=c10_includes, instantiations=instantiations, ) @@ -1524,9 +1541,21 @@ class Operator<{op_type}, kDev, {slot}> : public {op_type} {{ _TORCH_METHOD_TEMPLATE = """\ template void Operator<{op_type}, kDev, {slot}>::operator()({op_call_signature}) const {{ + const auto device_index = {out_device_index}; + const auto run = [&] {{ {tensor_conversions} - {aten_call}; + {aten_call}; + }}; + + if constexpr ({c10_condition}) {{ + const typename C10::StreamGuard stream_guard{{ + C10::GetStreamFromExternal(stream_, device_index)}}; + run(); + return; + }} + + run(); }} """ @@ -1534,6 +1563,8 @@ class Operator<{op_type}, kDev, {slot}> : public {op_type} {{ _TORCH_SOURCE_TEMPLATE = """\ #include "torch/{name}/{name}.h" +#include "torch/c10.h" +{c10_includes} #include "torch/tensor_.h" namespace infini::ops {{ diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 579f7a3af..d63dc3e8b 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -693,7 +693,7 @@ if(_infini_ops_linked_uses_torch) list(APPEND TORCH_SOURCES ${INFINI_OPS_LINKED_SOURCES}) endif() -if(WITH_CAMBRICON AND _infini_ops_linked_uses_torch) +if(WITH_CAMBRICON AND TORCH_SOURCES) execute_process( COMMAND "${_TORCH_PYTHON}" -c "import importlib.metadata, pathlib; root = pathlib.Path(importlib.metadata.distribution('torch-mlu').locate_file('')).resolve(); include = root / 'torch_mlu/csrc'; lib = include / 'lib'; print(include, lib, sep=';')" @@ -808,32 +808,39 @@ if(TORCH_SOURCES) set(_torch_vendor_include_flags "") if(WITH_METAX) - list(APPEND _torch_vendor_include_flags "-I${MACA_PATH}/include") + list(APPEND _torch_vendor_include_flags + "-I${MACA_PATH}/include" + "-I${MACA_PATH}/include/mcr" + "-I${MACA_PATH}/tools/cu-bridge/include") endif() if(WITH_MOORE) list(APPEND _torch_vendor_include_flags "-I${MUSA_ROOT}/include") - if(WITH_LINKED AND _infini_ops_linked_uses_torch) - execute_process( - COMMAND "${_TORCH_PYTHON}" -c - "import importlib.metadata, pathlib; root = pathlib.Path(importlib.metadata.distribution('torch-musa').locate_file('')).resolve(); compat = root / 'torch_musa/share/generated_cuda_compatible/include'; print(root, compat, sep=';')" - OUTPUT_VARIABLE _torch_musa_include_dirs - OUTPUT_STRIP_TRAILING_WHITESPACE - RESULT_VARIABLE _torch_musa_include_result - ) - if(NOT _torch_musa_include_result EQUAL 0 OR - NOT _torch_musa_include_dirs) + execute_process( + COMMAND "${_TORCH_PYTHON}" -c + "import importlib.metadata, pathlib; root = pathlib.Path(importlib.metadata.distribution('torch-musa').locate_file('')).resolve(); compat = root / 'torch_musa/share/generated_cuda_compatible/include'; lib = root / 'torch_musa/lib'; print(root, compat, lib, sep=';')" + OUTPUT_VARIABLE _torch_musa_paths + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE _torch_musa_paths_result + ) + if(NOT _torch_musa_paths_result EQUAL 0 OR NOT _torch_musa_paths) + message(FATAL_ERROR + "Could not determine `torch_musa` include and library directories.") + endif() + list(GET _torch_musa_paths 2 _torch_musa_library_dir) + list(REMOVE_AT _torch_musa_paths 2) + foreach(_dir IN LISTS _torch_musa_paths) + if(NOT IS_DIRECTORY "${_dir}") message(FATAL_ERROR - "Could not determine `torch_musa` include directories.") + "Could not find `torch_musa` include directory `${_dir}`.") endif() - foreach(_dir IN LISTS _torch_musa_include_dirs) - if(NOT IS_DIRECTORY "${_dir}") - message(FATAL_ERROR - "Could not find `torch_musa` include directory `${_dir}`.") - endif() - list(APPEND _torch_include_flags - "-isystem" "${_dir}") - endforeach() - endif() + list(APPEND _torch_include_flags "-isystem" "${_dir}") + endforeach() + + find_library(TORCH_MUSA_LIB musa_python + HINTS "${_torch_musa_library_dir}" + NO_DEFAULT_PATH REQUIRED) + target_link_libraries(infiniops PUBLIC "${TORCH_MUSA_LIB}") + list(APPEND TORCH_RUNTIME_DIRS "${_torch_musa_library_dir}") endif() # Vendor-specific defines required by forked `torch` headers. diff --git a/tests/test_abs.py b/tests/test_abs.py index 27781cc48..f7082db13 100644 --- a/tests/test_abs.py +++ b/tests/test_abs.py @@ -63,3 +63,48 @@ def _torch_abs(input, out): out.copy_(torch.abs(input)) return out + + +@pytest.mark.smoke +def test_abs_torch_backend_uses_handle_stream(device): + pytorch_slot = 8 + + if device != "cuda": + pytest.skip("The stream regression requires CUDA") + if not infini.ops.Abs.active_implementation_indices("nvidia"): + pytest.skip("The stream regression requires the NVIDIA backend") + if pytorch_slot not in infini.ops.Abs.active_implementation_indices(device): + pytest.skip("The PyTorch backend is not active") + + input = torch.full((4096,), -1.0, device=device) + out = torch.full_like(input, torch.nan) + stream = torch.cuda.Stream() + + def call_abs(): + infini.ops.abs( + input, + out, + stream=stream.cuda_stream, + implementation_index=pytorch_slot, + ) + + try: + call_abs() + stream.synchronize() + out.fill_(torch.nan) + torch.cuda.synchronize() + + with torch.cuda.stream(stream): + torch.cuda._sleep(50_000_000) + call_abs() + + default_stream = torch.cuda.default_stream() + with torch.cuda.stream(default_stream): + snapshot = out.clone() + default_stream.synchronize() + assert torch.isnan(snapshot).all() + + stream.synchronize() + torch.testing.assert_close(out, input.abs()) + finally: + torch.cuda.synchronize() diff --git a/tests/test_generate_torch_ops.py b/tests/test_generate_torch_ops.py index c7debc32a..43dffc4bc 100644 --- a/tests/test_generate_torch_ops.py +++ b/tests/test_generate_torch_ops.py @@ -96,6 +96,23 @@ def test_schema_self_param_renders_as_input_in_public_cpp_api(): assert "at::_softmax_out(at_out, at_self" in source +def test_torch_source_uses_existing_c10_stream_guards(): + 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 "stream_ == nullptr" not in method + assert "C10::GetStreamFromExternal(stream_, device_index)" in method + + for device in ("nvidia", "cambricon", "metax", "moore"): + assert f'#include "torch/{device}/c10.h"' in source + + for device in ("cpu", "ascend", "iluvatar", "hygon"): + assert f'#include "torch/{device}/c10.h"' not in source + + def test_optional_tensor_params_are_exposed_and_forwarded_to_aten(): module = _load_generator_module() op = module._parse_func(