已合并
【bugfix】pugraph_ex.scope.limit_core_num增加默认参数stream #38934
mihudan创建于 6月22日
【bugfix】pugraph_ex.scope.limit_core_num增加默认参数stream #38934
已合并
共 3 个文件变更+24-3
| @@ -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): |
| 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) |