Skip to content

fix(torch): honor handle streams in generated operators - #880

Open
voltjia wants to merge 2 commits into
masterfrom
fix/torch-backend-stream
Open

fix(torch): honor handle streams in generated operators#880
voltjia wants to merge 2 commits into
masterfrom
fix/torch-backend-stream

Conversation

@voltjia

@voltjia voltjia commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Moves the existing shared C10 declaration and the NVIDIA, MetaX, Cambricon, and Moore specializations from src/linked/torch to src/torch, while updating the existing linked providers to include the new paths.
  • Guards generated ATen tensor conversion and operator execution with each existing platform's native PyTorch stream guard, passing both null and non-null handle streams through GetStreamFromExternal and adding only the build wiring and tests required by that behavior.

Motivation

Generated ATen backends received the InfiniOps handle stream through stream_, but they invoked PyTorch without installing that stream in C10. PyTorch work could therefore run on a different stream from the caller-selected InfiniOps stream and violate the caller's ordering assumptions.

The existing C10 adapters already map a null external stream to the platform's default stream. Using the same external-stream path for both cases keeps the call sites uniform and preserves the established C10 API instead of introducing another stream-selection abstraction.

No issue is linked.

Type of Change

  • feat - New feature, operator, or platform.
  • fix - Bug fix.
  • perf - Performance improvement without a behavioral change.
  • refactor - Code restructuring without a behavioral change.
  • test - Adding or fixing tests only.
  • docs - Documentation only.
  • build / ci - Build system or CI configuration.
  • chore - Tooling, formatting, or other non-code changes.
  • Breaking change.

Platforms Affected

  • CPU (WITH_CPU).
  • NVIDIA (WITH_NVIDIA).
  • Iluvatar (WITH_ILUVATAR).
  • MetaX (WITH_METAX).
  • Cambricon (WITH_CAMBRICON).
  • Moore (WITH_MOORE).
  • Ascend (WITH_ASCEND).
  • PyTorch C++ bindings (WITH_TORCH).
  • Build system / CMake / CI.
  • Python bindings / user-facing API.

Smoke Test Result

Current validation SHA: 3f5b8b777b991c1c1a99aaccb1e63bc26d339096.

Ruff 0.15.22 lint and format checks passed.
Clang Format 21.1.8 passed for every changed C++ file.
The generator test module passed with 19 tests.
The NVIDIA smoke build passed against the current InfiniRT master in accelerator-dev/nvidia:latest.
The focused NVIDIA handle-stream regression passed with 1 test and 18 deselected tests.
The NVIDIA smoke suite passed with 82 tests, 24 skipped tests, and 7792 deselected tests.
The main and shadow CI matrices passed for every configured accelerator backend.

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA Yes Local smoke plus main and shadow CI passed. The focused handle-stream ordering regression passed on an A100.
Iluvatar No N/A - Not affected. No C10 specialization was added.
MetaX Yes Main and shadow CI passed. The existing C10 specialization is reused.
Cambricon Yes Main and shadow CI passed. The existing C10 specialization is reused.
Moore Yes Main and shadow CI passed. The existing C10 specialization is reused.
Ascend No N/A - Not affected. No C10 specialization was added.
Full `pytest` output (optional)
Running 98 items in this shard.
82 passed, 24 skipped, 7792 deselected in 2.36s.

Benchmark / Performance Impact

N/A. This PR fixes stream ordering and moves existing C10 helpers without changing operator algorithms.

Notes for Reviewers

This PR intentionally contains only two logical changes. The first commit moves the C10 helpers that already existed for NVIDIA, MetaX, Cambricon, and Moore. The second commit uses those helpers around generated ATen calls and adds the build wiring and regression coverage required for that use.

No C10 specialization is added for CPU, Iluvatar, Hygon, or Ascend. Generated operators on those platforms retain the previous execution path through the discarded if constexpr branch, and additional platform adapters can be handled in separate PRs when needed.

The stream guard covers both ToAtenTensor conversion and the ATen call. A null stream_ is passed directly to the existing GetStreamFromExternal implementation, which selects the default C10 stream on the supported adapters, while a non-null value imports the caller-provided native stream.

The CMake changes are limited to dependencies that become direct requirements outside linked providers. MetaX links its C10 CUDA bridge and runtime adapter, Cambricon makes the existing torch_mlu setup available to generated Torch sources, and Moore reuses the existing torch_musa discovery while linking its stream implementation.

@voltjia
voltjia requested a review from a team August 4, 2026 11:36
@voltjia
voltjia force-pushed the fix/torch-backend-stream branch 14 times, most recently from 5d4f0dc to a074719 Compare August 11, 2026 05:41
@voltjia
voltjia force-pushed the fix/torch-backend-stream branch from a074719 to 232dda7 Compare August 11, 2026 10:05
@voltjia
voltjia force-pushed the fix/torch-backend-stream branch from 232dda7 to 3f5b8b7 Compare August 11, 2026 10:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant