Tuning helpers

torch_lattice.utils.tune.tune(model, data_loader, n_samples=100, collect_fn=<function <lambda>>, enable_fp16=False, save_dir='.torch-lattice-tune', tune_tag='temp', force_retune=False, dataflow_range=None, dataflow_prune=False, tune_with_bwd=False, verbose=True, skip_warning=False)

Tune sparse convolution backend configuration for a model.

Parameters:

• model (Module) – Module to profile for convolution backend configuration.

• data_loader (Iterable) – Iterable that yields representative training samples.

• n_samples (int) – Number of samples used while profiling candidate configs.

• collect_fn (Callable) – Function that converts one data-loader item into model input. The tuned call is equivalent to model(collect_fn(data)) unless the callable returns a structure consumed by the model itself.

• enable_fp16 (bool) – Profile with half precision and CUDA autocast enabled.

• save_dir (str) – Directory used to cache tuned configuration files.

• tune_tag (str) – Cache file name under save_dir.

• force_retune (bool) – Ignore an existing cache file and profile again.

• dataflow_range (List) – Candidate convolution dataflows. When omitted, forward-only tuning checks implicit GEMM and Fetch-on-Demand; backward tuning uses implicit GEMM.

• dataflow_prune (bool) – Select the best dataflow before tuning lower-level config thresholds.

• tune_with_bwd (bool) – Include backward timing in the tuning objective.

• verbose (bool) – Print tuning progress and cache information.

• skip_warning (bool) – Suppress iterator and backend-mode warnings.