已合并
[feat] 支持多流api 入图 #30604
fu_hao创建于 2月6日
[feat] 支持多流api 入图 #30604
已合并
共 2 个文件变更+106-3
| @@ -0,0 +1,41 @@ | |||
| 1 | +import unittest | ||
| 2 | +import torch | ||
| 3 | +import torch_npu | ||
| 4 | + | ||
| 5 | + | ||
| 6 | +class TraceStreamEventTests(unittest.TestCase): | ||
| 7 | + def test_dynamo_trace_stream_event(self): | ||
| 8 | + | ||
| 9 | + def cus_func(t): | ||
| 10 | + s = torch.npu.Stream() | ||
| 11 | + tmp = torch.add(t, 2) | ||
| 12 | + event = torch.npu.Event() | ||
| 13 | + event.record() | ||
| 14 | + with torch.npu.stream(s): | ||
| 15 | + event.wait(s) | ||
| 16 | + r = torch.relu(tmp) | ||
| 17 | + r.record_stream(s) | ||
| 18 | + return r | ||
| 19 | + | ||
| 20 | + def my_backend(gm, example_inputs): | ||
| 21 | + graph = gm.graph | ||
| 22 | + print(graph) | ||
| 23 | + fx_target_list = (node.target for node in graph.nodes) | ||
| 24 | + assert_target_list = (torch_npu.npu.streams.Stream, | ||
| 25 | + torch_npu.npu.streams.Event, | ||
| 26 | + "record", | ||
| 27 | + "wait", | ||
| 28 | + torch_npu.utils._dynamo.fake_record_stream) | ||
| 29 | + for target in assert_target_list: | ||
| 30 | + self.assertIn(target, fx_target_list) | ||
| 31 | + return gm | ||
| 32 | + | ||
| 33 | + opt_m = torch.compile(cus_func, backend=my_backend, fullgraph=True, dynamic=False) | ||
| 34 | + i = torch.randn([3, 3], device="npu:0") | ||
| 35 | + with self.assertRaises(RuntimeError) as context: | ||
| 36 | + r = opt_m(i) | ||
| 37 | + self.assertIn("tensor.record_stream is not supported on torch.compile", str(context.exception)) | ||
| 38 | + | ||
| 39 | + | ||
| 40 | +if __name__ == '__main__': | ||
| 41 | + unittest.main() | ||
| @@ -257,12 +257,72 @@ def patch_user_defined_class_variable(): | |||
| 257 | 257 | ||
| 258 | def patched_in_graph_classes(): | 258 | def patched_in_graph_classes(): |
| 259 | result = original_method() | 259 | result = original_method() |
| 260 | - if hasattr(torch, "npu") and hasattr(torch.npu, "Event"): | 260 | + result.add(torch.npu.Event) |
| 261 | - result.add(torch.npu.Event) | 261 | + result.add(torch.npu.Stream) |
| 262 | return result | 262 | return result |
| 263 | UserDefinedClassVariable._in_graph_classes = patched_in_graph_classes | 263 | UserDefinedClassVariable._in_graph_classes = patched_in_graph_classes |
| 264 | 264 | ||
| 265 | 265 | ||
| 266 | +def fake_record_stream(self, s): | ||
| 267 | + """ | ||
| 268 | + let dynamo trace Tensor.record_stream as this emtpy function, | ||
| 269 | + and you can replace it later in your compile backend to an actual function | ||
| 270 | + """ | ||
| 271 | + if isinstance(self, torch._subclasses.fake_tensor.FakeTensor): | ||
| 272 | + return | ||
| 273 | + raise RuntimeError("tensor.record_stream is not supported on torch.compile! " | ||
| 274 | + "You should write a pass to replace torch.npu.fake_record_stream to an actual function in FX graph " | ||
| 275 | + "before aot_autograd.") | ||
| 276 | + | ||
| 277 | + | ||
| 278 | +def patch_record_stream(): | ||
| 279 | + torch.npu.fake_record_stream = fake_record_stream | ||
| 280 | + | ||
| 281 | + def method_record_stream(self, s): | ||
| 282 | + tx = torch._dynamo.symbolic_convert.InstructionTranslator.current_tx() | ||
| 283 | + return torch._dynamo.variables.TorchInGraphFunctionVariable( | ||
| 284 | + torch.npu.fake_record_stream | ||
| 285 | + ).call_function(tx, [self, s], {}) | ||
| 286 | + | ||
| 287 | + torch._dynamo.variables.tensor.TensorVariable.method_record_stream = method_record_stream | ||
| 288 | + | ||
| 289 | + | ||
| 290 | +def patch_variable_builder(): | ||
| 291 | + original_warp = torch._dynamo.variables.builder.VariableBuilder._wrap | ||
| 292 | + | ||
| 293 | + def _patch_wrapper(self, value): | ||
| 294 | + if isinstance(value, torch.npu.Event): | ||
| 295 | + self.install_guards(torch._dynamo.guards.GuardBuilder.ID_MATCH) | ||
| 296 | + torch._dynamo.utils.store_user_object_weakref(value) | ||
| 297 | + event_proxy = self.tx.output.create_proxy( | ||
| 298 | + "call_function", | ||
| 299 | + torch._dynamo.utils.get_user_object_from_id, | ||
| 300 | + (id(value),), | ||
| 301 | + {}, | ||
| 302 | + ) | ||
| 303 | + torch._dynamo.utils.set_example_value(event_proxy.node, value) | ||
| 304 | + out = torch._dynamo.variables.ctx_manager.EventVariable( | ||
| 305 | + event_proxy, | ||
| 306 | + value, | ||
| 307 | + source=self.source, | ||
| 308 | + ) | ||
| 309 | + return out | ||
| 310 | + return original_warp(self, value) | ||
| 311 | + | ||
| 312 | + torch._dynamo.variables.builder.VariableBuilder._wrap = _patch_wrapper | ||
| 313 | + | ||
| 314 | + | ||
| 315 | +def patch_builtin_variable(): | ||
| 316 | + origin_call_id = torch._dynamo.variables.builtin.BuiltinVariable.call_id | ||
| 317 | + | ||
| 318 | + def _wrap_call_id(self, tx, *args): | ||
| 319 | + if torch._dynamo.variables.builtin.istype(args[0], torch._dynamo.variables.ctx_manager.EventVariable): | ||
| 320 | + return torch._dynamo.variables.ConstantVariable.create(id(args[0].value)) | ||
| 321 | + return origin_call_id(self, tx, *args) | ||
| 322 | + | ||
| 323 | + torch._dynamo.variables.builtin.BuiltinVariable.call_id = _wrap_call_id | ||
| 324 | + | ||
| 325 | + | ||
| 266 | def add_dynamo_methods(): | 326 | def add_dynamo_methods(): |
| 267 | UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__ | 327 | UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__ |
| 268 | UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__ | 328 | UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__ |
| @@ -274,4 +334,6 @@ def add_dynamo_methods(): | |||
| 274 | patch_inductor_wrapper() | 334 | patch_inductor_wrapper() |
| 275 | patch_base_schedulernode() | 335 | patch_base_schedulernode() |
| 276 | patch_user_defined_class_variable() | 336 | patch_user_defined_class_variable() |
| 277 | - | 337 | + patch_record_stream() |
| 338 | + patch_variable_builder() | ||
| 339 | + patch_builtin_variable() | ||