[Feature]: 开放NPUStream::stream() 外部流场景下支持清空 default stream repo-pytorch-AtomGit

已关闭

[Feature]: 开放NPUStream::stream() 外部流场景下支持清空 default stream repo #3775

pwc10490创建于 8月6日关闭于 8月10日

pwc10490成员

8月6日创建

欢迎加入社区，感谢您对社区的贡献 🎉!

💻 需求背景、当前现状、期望实现的功能内容、具体的设计方案、以及测试方案

关联特性提案：#2163（支持 C++ 外部流接口 getStreamFromExternal）

💻 需求背景、当前现状、期望实现的功能内容、具体的设计方案、以及测试方案

1：背景

1.1 关联特性

#2163 引入了外部流（External Stream）能力，其功能支持矩阵（4.2 节）
将 NPUStream::stream() 列为「✅ 支持功能 —— 获取底层指针」。

1.2 现状

NPUStream::stream()（torch_npu/csrc/core/npu/NPUStream.cpp:417）
当前对外部流采用早退返回逻辑：

aclrtStream NPUStream::stream() const
{
    auto cur_ptr = NPUStream_internals(*this);
    AT_ASSERT(cur_ptr, PTA_ERROR(ErrCode::PTR));
    if (isExternalStreamId(this->id())) {
        return cur_ptr->stream;          // ← 外部流直接返回，跳过 repo 清空
    }
    if (c10_npu::option::OptionsManager::GetPerStreamQueue()) {
        // PerStreamQueue 分支：清空当前流自身 repo
        ...
        return cur_ptr->stream;
    } else {
        // 默认分支：清空 default_stream 的 repo
        auto ptr = NPUStream_internals(getDefaultNPUStream());
        ...
    }
    return cur_ptr->stream;
}

对比普通流与外部流的行为：

场景	普通流	外部流（当前）	问题
PerStreamQueue 关闭	清空 default repo 后返回	跳过清空，直接返回	default stream 上已入队任务未下发，与外部流算子存在执行顺序风险
PerStreamQueue 开启	清空自身 repo 后返回	早退返回（规避 null repo）	若早退被移除，`cur_ptr->repo` 为 nullptr，`CheckInit()` null-deref 崩溃

根因：外部流在 getStreamFromExternal（NPUStream.cpp:509）中执行
internals->repo.reset()，repo 为 nullptr，无独立 Task Queue。
因此外部流依赖 default stream 的 repo 进行任务下发保障，但早退逻辑
跳过了对 default repo 的 MakeSureQueueEmpty。

1.3 期望行为

PerStreamQueue 关闭 + 外部流：清空 default_stream 的 repo 后，
返回外部流的 aclrtStream（与普通流行为对齐，保障任务下发顺序）。
PerStreamQueue 开启 + 外部流：抛 TORCH_CHECK 异常拒绝
（与 enCurrentNPUStream 已有校验一致，避免 null repo 崩溃）。

2：设计方案

2.1 修改点

仅修改 torch_npu/csrc/core/npu/NPUStream.cpp 中
NPUStream::stream()（line 417），两处变动：

2.1.1 移除外部流早退逻辑

删除：

    if (isExternalStreamId(this->id())) {
        return cur_ptr->stream;
    }

使外部流进入与普通流一致的分支判断。

2.1.2 PerStreamQueue 分支增加外部流校验

在 PerStreamQueue 分支入口增加 TORCH_CHECK，参照
enCurrentNPUStream（NPUStream.cpp:753-755）已有校验：

    if (c10_npu::option::OptionsManager::GetPerStreamQueue()) {
        TORCH_CHECK(!isExternalStreamId(this->id()),
            "External NPU stream is not supported by task queue enqueue.",
            PTA_ERROR(ErrCode::NOT_SUPPORT));
        if (!this->isSyncLaunchStream() && cur_ptr->repo->CheckInit()) {
            ...

2.2 修改后完整逻辑

aclrtStream NPUStream::stream() const
{
    auto cur_ptr = NPUStream_internals(*this);
    AT_ASSERT(cur_ptr, PTA_ERROR(ErrCode::PTR));
    if (c10_npu::option::OptionsManager::GetPerStreamQueue()) {
        TORCH_CHECK(!isExternalStreamId(this->id()),
            "External NPU stream is not supported by task queue enqueue.",
            PTA_ERROR(ErrCode::NOT_SUPPORT));
        if (!this->isSyncLaunchStream() && cur_ptr->repo->CheckInit()) {
            NPUStatus ret = cur_ptr->repo->MakeSureQueueEmpty();
            if (ret != NPU_STATUS_SUCCESS) {
                ASCEND_LOGE("MakeSureQueueEmpty fail, ret: %s", ret.c_str());
                return nullptr;
            }
        }
        return cur_ptr->stream;
    } else {
        auto ptr = NPUStream_internals(getDefaultNPUStream());
        AT_ASSERT(ptr, PTA_ERROR(ErrCode::PTR));
        if (!this->isSyncLaunchStream() && ptr->repo->CheckInit()) {
            NPUStatus ret = ptr->repo->MakeSureQueueEmpty();
            if (ret != NPU_STATUS_SUCCESS) {
                ASCEND_LOGE("MakeSureQueueEmpty fail, ret: %s", ret.c_str());
                return nullptr;
            }
        }
    }
    return cur_ptr->stream;
}

2.3 行为对照

场景	修改前	修改后
PerStreamQueue 关闭 + 普通流	清空 default repo，返回 stream	不变
PerStreamQueue 关闭 + 外部流	跳过清空，直接返回	清空 default repo，返回外部 stream ✅
PerStreamQueue 开启 + 普通流	清空自身 repo，返回 stream	不变
PerStreamQueue 开启 + 外部流	早退返回（隐患）	TORCH_CHECK 抛异常（与 enCurrentNPUStream 一致） ✅

2.4 一致性说明

PerStreamQueue 模式下外部流的拒绝，与现有 enCurrentNPUStream
（NPUStream.cpp:753-760）对外部流的校验完全一致：

// enCurrentNPUStream 已有校验
TORCH_CHECK(!isExternalStreamId(task_stream->id()),
    "External NPU stream is not supported by task queue enqueue.",
    PTA_ERROR(ErrCode::NOT_SUPPORT));
TORCH_CHECK(!isExternalStreamId(static_cast<c10::StreamId>(ptr->stream_id)),
    "External NPU stream is not supported by task queue enqueue.",
    PTA_ERROR(ErrCode::NOT_SUPPORT));

外部流在 PerStreamQueue 模式下本就不支持任务入队，stream() 的校验
补全了获取底层指针环节的同一约束，避免后续 null-deref。

3：测试方案

3.1 PerStreamQueue 关闭场景

构造外部流（getStreamFromExternal）
在 default stream 上入队若干任务
调用 NPUStream::stream() 获取外部流指针
验证 default stream repo 被清空（任务已下发执行）
验证返回的 aclrtStream 与外部流指针一致

3.2 PerStreamQueue 开启场景

开启 PerStreamQueue（OptionsManager::GetPerStreamQueue() 返回 true）
构造外部流
调用 NPUStream::stream()
验证抛出 TORCH_CHECK 异常，错误信息包含
"External NPU stream is not supported by task queue enqueue."

3.3 回归

普通流场景行为不变（PerStreamQueue 开/关两分支）

4：影响面

改动文件：torch_npu/csrc/core/npu/NPUStream.cpp（1 文件，+3/-3 行）
影响版本：v2.9.0（后续可同步至 v2.7.1 / v2.10.0 / v2.11.0 / v2.12.0 / master）
兼容性：PerStreamQueue 关闭时外部流行为变化为"更正确"（多清空一次 default repo），
不破坏现有用法；PerStreamQueue 开启时外部流本就不被 enCurrentNPUStream 支持，
stream() 的拒绝是同一约束的补全。

5：关联 PR

fork 仓提交：pwc10490/pytorch_stream_patch v2.9.0 分支
commit：cbb7ae836 [fix] clear default stream repo for external stream in NPUStream::stream()
PR 链接：（待创建后回填）

欢迎加入社区，感谢您对社区的贡献 🎉!