已合并
[fix] torch2.12 dynamo原生支持record_stream #36504
fu_hao创建于 5月22日
[fix] torch2.12 dynamo原生支持record_stream #36504
已合并
共 2 个文件变更+6-37
| @@ -11,8 +11,8 @@ class TraceStreamEventTests(unittest.TestCase): | |||
| 11 | node_names = (node.name for node in gm.graph.nodes) | 11 | node_names = (node.name for node in gm.graph.nodes) |
| 12 | self.assertIn("current_stream", node_names) | 12 | self.assertIn("current_stream", node_names) |
| 13 | self.assertIn("set_stream", node_names) | 13 | self.assertIn("set_stream", node_names) |
| 14 | - self.assertIn("fake_record_stream", node_names) | 14 | + self.assertIn("record_stream", node_names) |
| 15 | - return gm | 15 | + return gm |
| 16 | 16 | ||
| 17 | 17 | ||
| 18 | def test_stream_in_graph(a): | 18 | def test_stream_in_graph(a): |
| @@ -24,14 +24,12 @@ class TraceStreamEventTests(unittest.TestCase): | |||
| 24 | event.wait() | 24 | event.wait() |
| 25 | r = torch.add(r, 1) | 25 | r = torch.add(r, 1) |
| 26 | r.record_stream(s) | 26 | r.record_stream(s) |
| 27 | - r = torch.add(r, 1) | 27 | + r = torch.add(r, 1) |
| 28 | return r | 28 | return r |
| 29 | - | 29 | + |
| 30 | i = torch.randn([3, 3], device="npu:0") | 30 | i = torch.randn([3, 3], device="npu:0") |
| 31 | - with self.assertRaises(RuntimeError) as context: | 31 | + r = test_stream_in_graph(i) |
| 32 | - r = test_stream_in_graph(i) | 32 | + return r |
| 33 | - return r | ||
| 34 | - self.assertIn("tensor.record_stream is not supported on torch.compile", str(context.exception)) | ||
| 35 | 33 | ||
| 36 | 34 | ||
| 37 | if __name__ == '__main__': | 35 | if __name__ == '__main__': |
| @@ -356,34 +356,6 @@ def patch_npu_stream_context(): | |||
| 356 | ] = _handle_npu_device_interface_stream | 356 | ] = _handle_npu_device_interface_stream |
| 357 | 357 | ||
| 358 | 358 | ||
| 359 | -def fake_record_stream(self, s): | ||
| 360 | - """ | ||
| 361 | - let dynamo trace Tensor.record_stream as this empty function, | ||
| 362 | - and you can replace it later in your compile backend to an actual function | ||
| 363 | - """ | ||
| 364 | - if isinstance(self, torch._subclasses.fake_tensor.FakeTensor): | ||
| 365 | - return | ||
| 366 | - raise RuntimeError( | ||
| 367 | - "tensor.record_stream is not supported on torch.compile! " | ||
| 368 | - "You should write a pass to replace torch.npu.fake_record_stream to an actual function in FX graph " | ||
| 369 | - "before aot_autograd." | ||
| 370 | - ) | ||
| 371 | - | ||
| 372 | - | ||
| 373 | -def patch_record_stream(): | ||
| 374 | - torch.npu.fake_record_stream = fake_record_stream | ||
| 375 | - | ||
| 376 | - def method_record_stream(self, s): | ||
| 377 | - tx = torch._dynamo.symbolic_convert.InstructionTranslator.current_tx() | ||
| 378 | - return torch._dynamo.variables.TorchInGraphFunctionVariable( | ||
| 379 | - torch.npu.fake_record_stream | ||
| 380 | - ).call_function(tx, [self, s], {}) | ||
| 381 | - | ||
| 382 | - torch._dynamo.variables.tensor.TensorVariable.method_record_stream = ( | ||
| 383 | - method_record_stream | ||
| 384 | - ) | ||
| 385 | - | ||
| 386 | - | ||
| 387 | def patch_user_defined_class_variable(): | 359 | def patch_user_defined_class_variable(): |
| 388 | import functools | 360 | import functools |
| 389 | 361 | ||
| @@ -413,5 +385,4 @@ def add_dynamo_methods(): | |||
| 413 | patch_event_variable_python_type() | 385 | patch_event_variable_python_type() |
| 414 | patch_builtin_variable() | 386 | patch_builtin_variable() |
| 415 | patch_npu_stream_context() | 387 | patch_npu_stream_context() |
| 416 | - patch_record_stream() | ||
| 417 | patch_user_defined_class_variable() | 388 | patch_user_defined_class_variable() |