[dynamo] Fix read/write conflicts in a cuda test (#145658)

Prior to this patch, the `test_cuda_event_created_outside_of_graph`
is flaky in CI, and that's because we have read and write to the same
`foo` tensor buffer from 2 different streams. This patch eliminates that
by adding a synchronization to wait till read finishes before starting
the write.

Fixes #133837, #133828.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/145658
Approved by: https://github.com/yifuwang
This commit is contained in:
Ryan Guo 2025-01-24 15:54:02 -08:00 committed by PyTorch MergeBot
parent c986eba560
commit 66631bc84b

View File

@ -461,6 +461,9 @@ class CtxManagerTests(torch._dynamo.test_case.TestCase):
torch.mm(x, x, out=foo)
event.record()
out = fn(foo)
# let `fn` finish reading `foo` before writing to it in the next
# iteration or `run_iters` call.
torch.cuda.current_stream().synchronize()
return out
ref = run_iters(func, compile=False)