full model: mark CUDA-graph steps so the compiled kernel runs with reduce-overhead
1 changed file+12−1
src/cymsim/full/simulation.pymodified+12−1View file
@@ -107,6 +107,10 @@ class Simulation:
107107 self._step = step
108108 self._compile_pending = bool(compile)
109109 self._compile_mode = compile if isinstance(compile, str) else ("reduce-overhead" if self.device.type == "cuda" else "default")
110+ # With CUDA graphs ("reduce-overhead") the outputs of one step are the
111+ # inputs of the next, and torch must be told where a step begins so that
112+ # it may reuse its output buffers after they have been copied.
113+ self._mark_step = lambda: None
110114
111115 # ---- progress ----
112116
@@ -141,8 +145,10 @@ class Simulation:
141145 S = self.static
142146 if self._compile_pending:
143147 self._try_compile()
148+ mark = self._mark_step
144149 for s in range(steps):
145150 local = self.total_steps + s
151+ mark()
146152 state_and_stick = self._step(S, *state, stick)
147153 state, stick = state_and_stick[:7], state_and_stick[7]
148154 if local % sps == 0:
@@ -170,7 +176,12 @@ class Simulation:
170176 self._compile_pending = False
171177 try:
172178 compiled = torch.compile(step, dynamic=False, mode=self._compile_mode)
173- compiled(self.static, *self._state, self._stick)
179+ if self.device.type == "cuda" and self._compile_mode in ("reduce-overhead", "max-autotune"):
180+ self._mark_step = torch.compiler.cudagraph_mark_step_begin
181+ self._mark_step()
182+ out = compiled(self.static, *self._state, self._stick)
183+ self._state, self._stick = tuple(x.clone() for x in out[:7]), out[7].clone()
184+ self.total_steps += 1
174185 self._step = compiled
175186 except Exception as e: # noqa: BLE001
176187 warnings.warn(f"torch.compile failed ({type(e).__name__}: {str(e).splitlines()[0][:200]}); running the uncompiled step, which is slower", RuntimeWarning, stacklevel=3)