| @@ -186,6 +186,27 @@ class TestNpuGraphEx(TestCase): | |||
| 186 | in4 = torch.randn(1000, 1000, dtype=torch.float16, device="npu") | 186 | in4 = torch.randn(1000, 1000, dtype=torch.float16, device="npu") |
| 187 | res = compiled_model(in1, in2, in3, in4) | 187 | res = compiled_model(in1, in2, in3, in4) |
| 188 | 188 | ||
| 189 | + def test_limit_core_num_with_stream(self): | ||
| 190 | + class Model(torch.nn.Module): | ||
| 191 | + def __init__(self): | ||
| 192 | + super().__init__() | ||
| 193 | + | ||
| 194 | + def forward(self, in1, in2, in3, in4): | ||
| 195 | + stream1 = torch.npu.Stream() | ||
| 196 | + with torch.npu.npugraph_ex.scope.limit_core_num(4, 5, stream=stream1): | ||
| 197 | + mm_result = torch.mm(in3, in4) | ||
| 198 | + with torch.npu.stream(stream1): | ||
| 199 | + add_result = torch.add(in1, in2) | ||
| 200 | + mm1_result = torch.mm(in3, in4) | ||
| 201 | + return add_result, mm_result, mm1_result | ||
| 202 | + | ||
| 203 | + model = Model().npu() | ||
| 204 | + compiled_model = torch.compile(model, backend="npugraph_ex", fullgraph=True, dynamic=False) | ||
| 205 | + in1 = torch.randn(1000, 1000, dtype=torch.float16, device="npu") | ||
| 206 | + in2 = torch.randn(1000, 1000, dtype=torch.float16, device="npu") | ||
| 207 | + in3 = torch.randn(1000, 1000, dtype=torch.float16, device="npu") | ||
| 208 | + in4 = torch.randn(1000, 1000, dtype=torch.float16, device="npu") | ||
| 209 | + res = compiled_model(in1, in2, in3, in4) | ||
| 189 | 210 | ||
| 190 | if __name__ == "__main__": | 211 | if __name__ == "__main__": |
| 191 | from torch._dynamo.test_case import run_tests | 212 | from torch._dynamo.test_case import run_tests |
| @@ -2307,7 +2307,7 @@ | |||
| 2307 | "signature": "(cache_bin, print_output=True, file=None)" | 2307 | "signature": "(cache_bin, print_output=True, file=None)" |
| 2308 | }, | 2308 | }, |
| 2309 | "torch_npu.npu.npugraph_ex.scope.limit_core_num": { | 2309 | "torch_npu.npu.npugraph_ex.scope.limit_core_num": { |
| 2310 | - "signature": "(op_aicore_num: int, op_vectorcore_num: int)" | 2310 | + "signature": "(op_aicore_num: int, op_vectorcore_num: int, stream=None)" |
| 2311 | }, | 2311 | }, |
 | |||
| 2312 | "torch_npu.npu.npugraph_ex.compile_fx": { | 2312 | "torch_npu.npu.npugraph_ex.compile_fx": { |
| 2313 | "signature": "(gm, example_inputs=None, options=None)" | 2313 | "signature": "(gm, example_inputs=None, options=None)" |
| @@ -1,6 +1,6 @@ | |||||||||||||||
| 1 | __all__ = ["limit_core_num"] | 1 | __all__ = ["limit_core_num"] | ||||||||||||
| 2 | 2 | ||||||||||||||
| 3 | 3 | ||||||||||||||
| 4 | -def limit_core_num(op_aicore_num: int, op_vectorcore_num: int): | 4 | +def limit_core_num(op_aicore_num: int, op_vectorcore_num: int, stream=None): | ||||||||||||
🔵 Low Priority 第4行： 虽然这不是功能缺陷，但缺失类型注解会降低代码可读性，且 IDE/类型检查器无法对 建议：为 stream 参数添加类型注解，例如   | |||||||||||||||
| 5 | from torch_npu.dynamo.npugraph_ex import scope | 5 | from torch_npu.dynamo.npugraph_ex import scope | ||||||||||||
| 6 | - return scope.limit_core_num(op_aicore_num, op_vectorcore_num) | 6 | + return scope.limit_core_num(op_aicore_num, op_vectorcore_num, stream=stream) | ||||||||||||
🟡 Medium Priority 变更后的 wrapper 受影响的合约/行为：旧版 失效模式：若 torchair 包未同步升级，所有 对 触发条件：torchair 版本未包含对应的 建议：在 wrapper 中判断 stream 是否为 None，仅在非 None 时才以关键字参数传递，保持对旧版 torchair 的向后兼容。 改动建议
| |||||||||||||||
🔵 Low Priority
第2310行：签名字符串正确反映了新增的
stream=None参数，JSON key"torch_npu.npu.npugraph_ex.scope.limit_core_num"与实际模块路径一致，未发现拼写错误或格式问题。该文件是本次变更的配套文档更新，与 Python 代码变更保持一致。
建议：无需修改，签名正确。