[CUDA][CUDA Graphs] Fix CUDAGraph::reset function (#108896)

Aidyn-A · pytorchmergebot · commit e7bd9c53158f · 2023-09-11T19:49:31.000Z
The following two cases fail due to a small oversight `CUDAGraph::reset()` that causes failures in graph destructor ```Python import torch x = torch.zeros(4, device="cuda") g = torch.cuda.CUDAGraph() with torch.cuda.graph(g): x = x + 1 g.reset() del g ``` that fails with: ``` terminate called after throwing an instance of 'c10::Error' what(): uc >= 0 INTERNAL ASSERT FAILED at ".../pytorch/c10/cuda/CUDACachingAllocator.cpp":2157, please report a bug to PyTorch. ``` and reset and subsequent re-capture ```Python import torch x = torch.zeros(4, device="cuda") g = torch.cuda.CUDAGraph() with torch.cuda.graph(g): x = x + 1 g.reset() with torch.cuda.graph(g): x = x + 1 g.replay() ``` which fails with: ``` Traceback (most recent call last): File "test_graph.py", line 11, in <module> with torch.cuda.graph(g): File ".../pytorch/torch/cuda/graphs.py", line 192, in __enter__ self.cuda_graph.capture_begin( File ".../pytorch/torch/cuda/graphs.py", line 77, in capture_begin super().capture_begin(pool=pool, capture_error_mode=capture_error_mode) RuntimeError: This CUDAGraph instance already owns a captured graph. To capture a new graph, create a new instance. ``` This PR fixes `CUDAGraph::reset()` function for above to use cases. Pull Request resolved: #108896 Approved by: https://github.com/ezyang
diff --git a/aten/src/ATen/cuda/CUDAGraph.cpp b/aten/src/ATen/cuda/CUDAGraph.cpp
@@ -295,9 +295,11 @@ void CUDAGraph::reset() {
   }
   if (has_graph_) {
     C10_CUDA_CHECK_WARN(cudaGraphDestroy(graph_));
+    has_graph_ = false;
   }
   if (has_graph_exec_) {
     C10_CUDA_CHECK_WARN(cudaGraphExecDestroy(graph_exec_));
+    has_graph_exec_ = false;
   }
 #else
   TORCH_CHECK(false, "CUDA graphs may only be used in Pytorch built with CUDA >= 11.0 or ROCM >= 5.3")
diff --git a/test/test_cuda.py b/test/test_cuda.py
@@ -2168,6 +2168,41 @@ def test_graph_capture_simple(self):
 
         self.assertTrue(b.sum().item() == 11000.)
 
+    @unittest.skipIf(not TEST_CUDA_GRAPH, "CUDA >= 11.0 or ROCM >= 5.3 required for graphs")
+    def test_graph_capture_reset_recapture(self):
+        s = torch.cuda.Stream()
+
+        with torch.cuda.stream(s):
+            a = torch.full((1000,), 1, device="cuda")
+            g = torch.cuda.CUDAGraph()
+            torch.cuda.empty_cache()
+            g.capture_begin()
+            b = a
+            for _ in range(10):
+                b = b + 1
+            g.capture_end()
+        torch.cuda.current_stream().wait_stream(s)
+
+        g.replay()
+
+        self.assertTrue(b.sum().item() == 11000.)
+
+        g.reset()
+
+        with torch.cuda.stream(s):
+            g.capture_begin()
+            b.fill_(2.0)
+            for _ in range(10):
+                b = b + 2
+            g.capture_end()
+        torch.cuda.current_stream().wait_stream(s)
+
+        g.replay()
+        self.assertTrue(b.sum().item() == 22000.)
+
+        g.reset()
+        del g
+
     @unittest.skipIf(not TEST_CUDA_GRAPH, "CUDA >= 11.0 or ROCM >= 5.3 required for graphs")
     def test_graph_error(self):
         # We need to run this test in a separate thread as the error we trigger
