CompiledModule

class qlip.compiler.module.CompiledModule(model, *, builder_config, **kwargs)

Bases: CompiledModule

Compiled module class for compilation.

Parameters

  • model (torch.nn.Module) – Model to compile.

  • builder_config (BuilderConfig) – Builder configuration.

  • **kwargs (dict) – Additional keyword arguments, see CompiledModule for more details.

Variables

  • builder_config (BuilderConfig) – Builder configuration.

  • ioconfig (IOConfig) – Input and output configuration.

compile(device='cuda', original_device='meta', **kwargs)

Compile model.

Parameters

  • device (str) – Device to compile model, by default “cuda”

  • original_device (str) – Original device of the compiled model, by default “meta”

  • **kwargs (dict) – Additional keyword arguments, backend specific.

collect_shapes(value=True)

Toggle shape collection mode.

Parameters

  • value (bool) – Whether to collect shapes.

collect_inputs(value=True)

Toggle input collection mode.

Parameters

  • value (bool) – Whether to collect inputs for export.

set_skip_n(n)

Set number of forward passes to skip during collection (shapes or inputs).

Parameters

  • n (int) – Number of forward passes to skip.

flush_shapes(type='static', opt='mode')

Flush collected shapes to static or dynamic axes profiles.

Parameters

  • type (str) – Type of shapes to flush. Can be “static” or “dynamic”.

  • opt (str) – Optimal shape for dynamic axes profile. Can be “mode”, “min” or “max”.

forward(*args, **kw)

Forward pass.

Run compiled model if available, otherwise collect shapes.

enable(val)

Enable or disable compiled module functionality.

export_mode()

Toggle export mode for model.

Used by ~qlip.compiler.base.BaseExporter.export methods.

extra_repr()

extra_repr is used by torch.nn.Module to print the model structure.

load(engine_path, device='cuda', original_device='meta', **kwargs)

Load session into compiled module.

Parameters

  • engine_path (str | Path) – Path to engine file.

  • device (str) – Device to load session to.

  • original_device (str) – Move original model to this device.

  • **kwargs (dict) – Additional keyword arguments, backend specific.

property original_device: torch.device

Device of the original model.

set_inference_config(config)

Set inference configuration for lazy initialization before loading engine.

Parameters

  • config (SessionConfig) – Configuration for inference session.

unload()

Unload compiled engine.