# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/callback.compiler.ipynb.

# %% ../../nbs/callback.compiler.ipynb 1
# Contains code from:
# fastai - Apache License 2.0 - Copyright (c) 2023 fast.ai
# PyTorch - PyTorch BSD-style license - Copyright (c) 2013-present PyTorch contributors

# %% ../../nbs/callback.compiler.ipynb 4
from __future__ import annotations
from typing import Dict

from copy import deepcopy
from enum import Enum
from pathlib import Path
import pickle
import warnings

from packaging.version import parse

import torch._dynamo as dynamo
from torch.serialization import FILE_LIKE

from fastai.learner import Learner, save_model, join_path_file, _cast_tensor
from fastai.callback import schedule
from fastai.callback.core import Callback, TrainEvalCallback, CancelFitException

import fastai
if parse(fastai.__version__) < parse('2.7.13'):
    from fastxtend.callback.amp import MixedPrecision
else:
    from fastai.callback.fp16 import MixedPrecision

try:
    from fastxtend.ffcv.loader import Loader
    FFCV = True
except ImportError:
    FFCV = False

from ..imports import *

# %% auto 0
__all__ = ['CompileMode', 'MatMulPrecision', 'CompilerCallback', 'DynamoExplainCallback', 'load_learner']

# %% ../../nbs/callback.compiler.ipynb 7
_min_torch_20 = ismin_torch('2.0')
_only_torch_20 = ismin_torch('2.0') and notmax_torch('2.1.0')
_min_torch_21 = ismin_torch('2.1.0')

if not _min_torch_20:
    warn('Imported `fastxtend.callback.compiler`, which requires a minimum of PyTorch 2.0 to work.')

# %% ../../nbs/callback.compiler.ipynb 8
def is_compiled(model:nn.Module):
    "Check whether a `nn.Module` model has been compiled"
    return (hasattr(model, '_compiled_call_impl') and getattr(model, '_compiled_call_impl') is not None) \
            or isinstance(model, dynamo.OptimizedModule)

# %% ../../nbs/callback.compiler.ipynb 9
class CompileMode(str, Enum):
    "All valid `torch.compile` modes for tab-completion and typo-proofing"
    default        = 'default'
    reduceoverhead = 'reduce-overhead'
    maxautotune    = 'max-autotune'

# %% ../../nbs/callback.compiler.ipynb 11
class MatMulPrecision(str, Enum):
    "All valid `matmul_precision` modes for tab-completion and typo-proofing"
    highest = 'highest'
    high    = 'high'
    medium  = 'medium'

# %% ../../nbs/callback.compiler.ipynb 13
class CompilerCallback(Callback):
    "A callback for using `torch.compile` (beta) with fastai"
    order = TrainEvalCallback.order + 1 # Compiling needs to occur on the GPU, but before distributed training starts

    def __init__(self,
        fullgraph:bool=False, # Prevent breaking model into subgraphs
        dynamic:bool=False, # Use dynamic shape tracing
        backend:str|Callable='inductor', # `torch.compile` backend to use
        mode:str|CompileMode|None=None, # `torch.compile` mode to use
        options:Dict[str, Union[str,int,bool]]|None=None, # Extra options to pass to compile backend
        matmul_precision:str|MatMulPrecision='high', # Set Ampere and newer matmul precision
        recompile:bool=False, # Force a compiled model to recompile. Use when freezing/unfreezing a compiled model.
        verbose:bool=False, # Verbose output
    ):
        mode = CompileMode(mode).value if mode is not None else mode
        matmul_precision = MatMulPrecision(matmul_precision).value
        if mode is not None and options is not None:
            raise ValueError(f"Both {mode=} or {options=} cannot be set at the same time.")
        store_attr(but='recompile')
        self._recompile = recompile

    def before_fit(self):
        if not _min_torch_20:
            self.run = False
            warn("Attempting to use `CompilerCallback` without PyTorch 2.0 or greater. Disabling.")
            return

        if torch.cuda.get_device_capability() >= (8, 0) and torch.get_float32_matmul_precision() != self.matmul_precision:
            if self.verbose and self.matmul_precision != 'highest':
                print(f"Your GPU has modern tensor cores, automatically enabling by setting `torch.set_float32_matmul_precision('{self.matmul_precision}')`")
            torch.set_float32_matmul_precision(self.matmul_precision)

        if hasattr(self.learn, 'progressive_resize') and _only_torch_20:
            warn("Using `ProgressiveResize` and `torch.compile` at the same time with PyTorch 2.0 will result in a new compile every size change.")
        msg = ""
        if self.dynamic and _only_torch_20:
            msg += "Using PyTorch 2.0 `compile` with dynamic shapes is under active development and might fail. Upgrade to PyTorch 2.1.\n"
        if self.mode == 'max-autotune':
            msg += "Using `torch.compile` with `mode='max-autotune'` is under active development and might fail\n"
        if msg != "":
            msg += "See https://pytorch.org/docs/master/compile/index.html#troubleshooting-and-gotchas for more details"
            warn(msg)
        if self.mode == 'reduce-overhead':
            warn("Using `torch.compile` & fastai with `mode='reduce-overhead'` currently doesn't appear to train.")

        if self._recompile and is_compiled(self.learn.model):
            if self.verbose:
                print("Recompiling model")
            self._reset_compiled()
        self._recompile = False

        if not is_compiled(self.learn.model):
            if not isinstance(self.learn.model, nn.Module):
                warn("Model is not ")
            if hasattr(self.learn.model, 'compile'):
                self.learn.model.compile(fullgraph=self.fullgraph, dynamic=self.dynamic,
                                         backend=self.backend, mode=self.mode, options=self.options)
            else:
                compiled_model = torch.compile(self.learn.model, fullgraph=self.fullgraph,
                                               dynamic=self.dynamic, backend=self.backend,
                                               mode=self.mode, options=self.options)
                self.learn.model = compiled_model._orig_mod
                self.learn.model._orig_forward = self.learn.model.forward
                self.learn.model.forward = compiled_model.dynamo_ctx(self.learn.model.forward)
                self.learn.model._compiled_call_impl = True
        else:
            warn("Model is already compiled. To recompile pass `recomple=True` to CompilerCallback.")

    def _reset_compiled(self):
        if is_compiled(self.learn.model):
            dynamo.reset()
            self.learn.model._compiled_call_impl = None
            if hasattr(self.learn.model, '_orig_forward'):
                self.learn.model.forward = self.learn.model._orig_forward
            if isinstance(self.learn.model, dynamo.OptimizedModule):
                self.learn.model = self.learn.model._orig_mod

# %% ../../nbs/callback.compiler.ipynb 16
class DynamoExplainCallback(Callback):
    "A callback to automate finding graph breaks with PyTorch Compile's Dynamo Explain"
    order = MixedPrecision.order+1 # DynamoExplain occurs on the GPU before any training starts

    def __init__(self,
        print_results:bool=True, # Print enabled `torch._dynamo.explain` output(s)
        out_guards:bool=False, # Print the `out_guards` output
        ops_per_graph:bool=False, # Print the `ops_per_graph` output
        break_reasons:bool=False, # Print the `break_reasons` output
    ):
        self.print_results = print_results
        self.print_out_guards = out_guards
        self.print_ops_per_graph = ops_per_graph
        self.print_break_reasons = break_reasons

    def before_fit(self):
        if not _min_torch_20:
            self.run = False
            warn("Attempting to use `DynamoExplainCallback` without PyTorch 2.0 or greater. Canceling training.")
            raise CancelFitException()

        self.explanation, self.out_guards, self.graphs, self.ops_per_graph, self.break_reasons, self.explanation_verbose = '','','','','',''
        states = get_random_states()
        try:
            if FFCV and isinstance(self.dls.train, Loader) and self.dls.train.async_tfms:
                # With `async_tfms`, `Loader` needs to initialize all `Loader.batches_ahead` Cuda streams
                # for the training dataloader. Since FFCV doesn't support seeded transforms and the reset
                # random state only seeds the dataset order, this shouldn't effect training outcome.
                b = self.dls.train.one_batch(batches_ahead=True)
            elif hasattr(self.dls.valid, 'one_batch'):
                b = self.dls.valid.one_batch()
            else:
                b = next(iter(self.dls.valid))
                model_device = next(self.model.parameters()).device
                b = to_device(b, model_device)
            i = getattr(self.dls, 'n_inp', 1 if len(b)==1 else len(b)-1)
            self.learn.xb, self.learn.yb = b[:i], b[i:]

            if hasattr(self.learn, 'mixed_precision'):
                self.learn.mixed_precision.autocast.__enter__()

            if _only_torch_20:
                self.explanation, self.out_guards, self.graphs, self.ops_per_graph, self.break_reasons, self.explanation_verbose \
                    = dynamo.explain(self.learn.model, *_cast_tensor(self.learn.xb))
            else:
                self.explain_output = dynamo.explain(self.learn.model)(*_cast_tensor(self.learn.xb))

            if hasattr(self.learn, 'mixed_precision'):
                self.learn.mixed_precision.autocast.__exit__(None, None, None)

            self.learn.opt.zero_grad()
        finally:
            set_random_states(**states)

        if self.print_results:
            print('\nDynamo Explain Report\n')
            if _min_torch_21:
                print_copy = deepcopy(self.explain_output)
                if not self.print_ops_per_graph:
                    print_copy.ops_per_graph = None
                if not self.print_out_guards:
                    print_copy.out_guards = None
                print(print_copy)
                print_copy = None
            else:
                output = "Explanation:\n"
                output += f"  {self.explanation}\n"
                output += "Break Reasons:\n"
                for idx, break_reason in enumerate(self.break_reasons):
                    output += f"  Break Reason {idx+1}:\n"
                    output += f"    Reason: {break_reason.reason}\n"
                    output += "    User Stack:\n"
                    for frame_summary in break_reason.user_stack:
                        output += f"      {frame_summary}\n"

                if self.ops_per_graph is not None and self.print_ops_per_graph:
                    output += "Ops per Graph:\n"
                    for idx, ops in enumerate(self.ops_per_graph):
                        output += f"  Ops {idx+1}:\n"
                        for op in ops:
                            output += f"    {op}\n"

                if self.out_guards is not None and self.print_out_guards:
                    output += "Out Guards:\n"
                    for i, guard in enumerate(self.out_guards[0]):
                        output += f"  Guard {i+1}:\n"
                        output += f"    {str(guard)}"

                print(output)

            print('\n')

        raise CancelFitException()

# %% ../../nbs/callback.compiler.ipynb 19
@patch
def compile(self:Learner,
    fullgraph:bool=False, # Prevent breaking model into subgraphs
    dynamic:bool=False, # Use dynamic shape tracing. Sets to `False` if PyTorch < 2.1
    backend:str|Callable='inductor', # `torch.compile` backend to use
    mode:str|CompileMode|None=None, # `torch.compile` mode to use
    options:Dict[str, Union[str,int,bool]]|None=None, # Extra options to pass to compile backend
    matmul_precision:str|MatMulPrecision='high', # Set Ampere and newer matmul precision
    recompile:bool=False, # Force a compiled model to recompile. Use when freezing/unfreezing a compiled model.
    verbose:bool=False, # Verbose output
):
    "Set `Learner` to compile model using `torch.compile` via `CompilerCallback`"
    return self.add_cb(CompilerCallback(fullgraph=fullgraph, dynamic=dynamic if _min_torch_21 else False,
                                        backend=backend, mode=mode, options=options,
                                        matmul_precision=matmul_precision,
                                        recompile=recompile, verbose=verbose))

# %% ../../nbs/callback.compiler.ipynb 23
@patch
def export(self:Learner,
    fname:FILE_LIKE='export.pkl', # Learner export file name, path, bytes, or IO
    pickle_module:Any=pickle, # Module used for pickling metadata and objects
    pickle_protocol:int=2 # Pickle protocol used
):
    "Export the content of `self` without the items and the optimizer state for inference"
    if rank_distrib(): return # don't export if child proc
    self._end_cleanup()
    old_dbunch = self.dls
    self.dls = self.dls.new_empty()
    state = self.opt.state_dict() if self.opt is not None else None
    self.opt = None
    compiled_forward = None
    # torch.compiled models currently cannot be pickled.
    if _only_torch_20 and is_compiled(self.model):
        compiled_forward = self.model.forward
        self.model.forward = self.model._orig_forward
        delattr(self.model, '_orig_forward')
    with warnings.catch_warnings():
        # To avoid the warning that come from PyTorch about model not being checked
        warnings.simplefilter("ignore")
        torch.save(self, self.path/fname, pickle_module=pickle_module, pickle_protocol=pickle_protocol)
    if _min_torch_20 and compiled_forward is not None:
        self.model._orig_forward = self.model.forward
        self.model.forward = compiled_forward
    self.create_opt()
    if state is not None: self.opt.load_state_dict(state)
    self.dls = old_dbunch

# %% ../../nbs/callback.compiler.ipynb 24
def load_learner(
    fname:FILE_LIKE, # File name, path, bytes, or IO
    cpu:bool=True, # Load model to CPU
    pickle_module=pickle # Module used for unpickling metadata and objects
):
    "Load a `Learner` object in `fname`, by default putting it on the `cpu`"
    distrib_barrier()
    map_loc = 'cpu' if cpu else default_device()
    try: res = torch.load(fname, map_location=map_loc, pickle_module=pickle_module)
    except AttributeError as e:
        e.args = [f"Custom classes or functions exported with your `Learner` not available in namespace.\Re-declare/import before loading:\n\t{e.args[0]}"]
        raise
    if cpu:
        res.dls.cpu()
        if hasattr(res, 'channels_last'):
            res = res.to_contiguous(to_fp32=True)
        elif hasattr(res, 'mixed_precision'):
            res = res.to_fp32()
        elif hasattr(res, 'non_native_mixed_precision'):
            res = res.to_non_native_fp32()
        if hasattr(res, 'compiler'):
            res = res.remove_cb(CompilerCallback)
    return res

# %% ../../nbs/callback.compiler.ipynb 27
@patch
def freeze_to(self:Learner, n:int):
    "Freeze parameter groups up to `n`"
    if self.opt is None:
        self.create_opt()
    self.opt.freeze_to(n)
    self.opt.clear_state()
    if _min_torch_20 and is_compiled(self.model):
        if hasattr(self, 'compiler'):
            self.compiler._recompile = True
        else:
            warn("Freezing or unfreezing a compiled model isn't supported."\
                 "\nThe model must be recompiled to take effect."\
                 "\nPass `CompilerCallback(..., recompile=True)` to `Learner.cbs`"\
                 "\nor call `torch._dynamo.reset() and recompile model.")

# %% ../../nbs/callback.compiler.ipynb 30
@patch
@delegates(Learner.fit_one_cycle)
def fine_tune(self:Learner,
    epochs:int, # Number of unfrozen epochs to train
    base_lr:float=2e-3, # Base learning rate, model head unfrozen learning rate
    freeze_epochs:int=1, # Number of frozen epochs to train
    lr_mult:Numeric=100, # Model stem unfrozen learning rate: `base_lr/lr_mult`
    pct_start:float=0.3, # Start unfrozen learning rate cosine annealing
    div:Numeric=5.0, # Initial unfrozen learning rate: `base_lr/div`
    compile_frozen:bool=False, # Compile model during frozen finetuning if `CompilerCallback` is used
    **kwargs
):
    "Fine tune with `Learner.freeze` for `freeze_epochs`, then with `Learner.unfreeze` for `epochs`, using discriminative LR."
    self.freeze()
    if _min_torch_20 and hasattr(self, 'compiler') and not compile_frozen:
        self.compiler.run = is_compiled(self.model)
    self.fit_one_cycle(freeze_epochs, slice(base_lr), pct_start=0.99, **kwargs)
    base_lr /= 2
    self.unfreeze()
    self.fit_one_cycle(epochs, slice(base_lr/lr_mult, base_lr), pct_start=pct_start, div=div, **kwargs)
