fused_fftconv2d_lowering

fused_fftconv2d_lowering(allow_reduced_precision=True)

Context manager that installs FusedFFTConv2dLowering globally.

Prefer fused_fftconv2d_options() when you control the torch.compile call — it is scoped to a single callable instead of patching global inductor config. Reach for this context manager when something else compiles the model for you (a trainer, a framework entry point) and you cannot pass options.

torch.compile caches compiled artifacts, so enter this before the first compiled call on the model you want rewritten; a function already compiled without the pass keeps its cached code.

Parameters:

allow_reduced_precision (bool) – See FusedFFTConv2dLowering.

Example

>>> import torch
>>> from nvsubquadratic.ops.fftconv_lowering import (
...     fused_fftconv2d_lowering,
...     lowering_stats,
... )
>>> with fused_fftconv2d_lowering():
...     compiled = torch.compile(model)
...     out = compiled(x)
>>> lowering_stats().get("rewritten", 0)
1

Note

Inductor exposes a single pre_grad_custom_pass slot. Any pass already registered is saved and restored on exit, but the two do not compose while this one is active.