full model: on CUDA, capture the steps of one audio sample into a CUDA graph replayed on static buffers

@@ -106,11 +106,13 @@ class Simulation:

106106

107107 self._step = step

108108 self._compile_pending = bool(compile)

109- 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

109+ self._compile_mode = compile if isinstance(compile, str) else "default"

110+ # On CUDA the steps of one audio sample are captured into a CUDA graph

111+ # (see _try_compile) and replayed on static buffers: self._graph, the

112+ # buffer that receives the transverse velocity at the sample instant,

113+ # and the stream the graph was captured on.

114+ self._graph = None

115+ self._sample = None

114116

115117 # ---- progress ----

116118

@@ -145,16 +147,23 @@ class Simulation:

145147 S = self.static

146148 if self._compile_pending:

147149 self._try_compile()

148- mark = self._mark_step

149- for s in range(steps):

150- local = self.total_steps + s

151- mark()

152- state_and_stick = self._step(S, *state, stick)

153- state, stick = state_and_stick[:7], state_and_stick[7]

154- if local % sps == 0:

150+ if self._graph is not None:

151+ # One graph replay per audio sample; the static buffers advance in place.

152+ for local in range(self.total_steps, self.total_steps + steps, sps):

153+ self._graph.replay()

155154 slot = (local // sps) % H

156- self._hist[slot] = state[6].reshape(-1)

155+ self._hist[slot] = self._sample.reshape(-1)

157156 slots.append(slot)

157+ state, stick = self._state, self._stick

158+ else:

159+ for s in range(steps):

160+ local = self.total_steps + s

161+ state_and_stick = self._step(S, *state, stick)

162+ state, stick = state_and_stick[:7], state_and_stick[7]

163+ if local % sps == 0:

164+ slot = (local // sps) % H

165+ self._hist[slot] = state[6].reshape(-1)

166+ slots.append(slot)

158167 self._state, self._stick = state, stick

159168 self.total_steps += steps

160169 if slots:

@@ -167,23 +176,60 @@ class Simulation:

167176 self.seconds += time.perf_counter() - t0

168177

169178 def _try_compile(self) -> None:

170- """Wraps the step in ``torch.compile`` (which fuses its many small tensor

171- operations; about 6 times faster on the CPU and needed on a GPU) and

172- checks it on the current state; falls back to the eager step with a

173- warning when no compiler is available."""

179+ """Wraps the step in ``torch.compile``, which fuses its some two hundred

180+ small tensor operations into a few kernels (about 6 times faster than

181+ the plain step on the CPU), and on CUDA captures the steps of one audio

182+ sample into a CUDA graph, so that a sample costs one graph launch

183+ rather than a Python call per kernel (a further 4 to 5 times). Falls

184+ back, with a warning, to whatever part of that works."""

174185 import warnings

175186

176187 self._compile_pending = False

177188 try:

178189 compiled = torch.compile(step, dynamic=False, mode=self._compile_mode)

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()

182190 compiled(self.static, *self._state, self._stick) # compiles; the result is discarded and the state is unchanged

183191 self._step = compiled

184192 except Exception as e: # noqa: BLE001

185193 warnings.warn(f"torch.compile failed ({type(e).__name__}: {str(e).splitlines()[0][:200]}); running the uncompiled step, which is slower", RuntimeWarning, stacklevel=3)

186194 self._step = step

195+ return

196+ if self.device.type != "cuda":

197+ return

198+ try:

199+ self._capture_graph()

200+ except Exception as e: # noqa: BLE001

201+ warnings.warn(f"CUDA graph capture failed ({type(e).__name__}: {str(e).splitlines()[0][:200]}); running the compiled step without a graph, which is slower", RuntimeWarning, stacklevel=3)

202+ self._graph = None

203+

204+ def _capture_graph(self) -> None:

205+ """Captures ``steps_per_sample`` compiled steps into a CUDA graph that

206+ reads and writes the state in place. The transverse velocity after the

207+ first step of the graph, the sample instant, is copied to ``_sample``."""

208+ S, sps = self.static, self.steps_per_sample

209+ state = [x.clone() for x in self._state]

210+ stick = self._stick.clone()

211+ sample = torch.empty_like(state[6])

212+ side = torch.cuda.Stream()

213+ side.wait_stream(torch.cuda.current_stream())

214+ with torch.cuda.stream(side):

215+ for _ in range(3): # warm-up on a side stream, as capture requires

216+ self._step(S, *state, stick)

217+ torch.cuda.current_stream().wait_stream(side)

218+ torch.cuda.synchronize()

219+ graph = torch.cuda.CUDAGraph()

220+ with torch.cuda.graph(graph):

221+ cur, st = tuple(state), stick

222+ for i in range(sps):

223+ out = self._step(S, *cur, st)

224+ cur, st = out[:7], out[7]

225+ if i == 0:

226+ sample.copy_(cur[6])

227+ for a, b in zip(state, cur):

228+ a.copy_(b)

229+ stick.copy_(st)

230+ torch.cuda.synchronize()

231+ self._state, self._stick = tuple(state), stick

232+ self._graph, self._sample = graph, sample

187233

188234 def _audio(self, slots: list[int]) -> np.ndarray:

189235 """The retarded-time Rayleigh sum for the batch's samples: the acceleration

@@ -249,7 +295,16 @@ class Simulation:

249295 if x is not None:

250296 cur[k] = torch.as_tensor(np.asarray(x, dtype=np.float64), dtype=self.dtype, device=self.device) * self.static.ins

251297 cur[3] = torch.zeros_like(cur[2])

252- self._state = tuple(cur)

298+ self._assign_state(tuple(cur), self._stick)

299+

300+ def _assign_state(self, state, stick) -> None:

301+ """Replaces the state; into the graph's static buffers when a graph has been captured."""

302+ if self._graph is not None:

303+ for a, b in zip(self._state, state):

304+ a.copy_(b)

305+ self._stick.copy_(stick)

306+ else:

307+ self._state, self._stick = tuple(state), stick

253308

254309 def step_once(self, strike: bool = False) -> None:

255310 """Advances the state by one step without recording; without the drumstick unless ``strike``."""

@@ -259,7 +314,7 @@ class Simulation:

259314

260315 S = replace(S, strike=False)

261316 out = step(S, *self._state, self._stick)

262- self._state, self._stick = out[:7], out[7]

317+ self._assign_state(out[:7], out[7])

263318 self.total_steps += 1

264319

265320

@@ -284,7 +339,8 @@ def simulate(

284339 default, with the rounding safeguards of the kernel) or ``"float64"``.

285340 ``compile`` wraps the step in ``torch.compile`` (``True``, the default,

286341 with a fallback to the plain step when no compiler is available; ``False``;

287- or a ``torch.compile`` mode string). ``fps`` is the movie frame rate and

342+ or a ``torch.compile`` mode string) and, on CUDA, captures the steps of one

343+ audio sample into a CUDA graph. ``fps`` is the movie frame rate and

288344 also the batch rate of the run. ``loss`` overrides the fitted

289345 ``(sigma0, sigma1, sigma_in)`` of the two-parameter loss operator.

290346 """