import torch
import intel_extension_for_pytorch._C as core

# Expand torch cpu autocast to support fp32 dtype for int8 path,
# for pytorch torch.cpu.amp.autocast, if give dtyps is fp32, the autocast will be disabled.
class _autocast(torch.cpu.amp.autocast):
    def __enter__(self):
        self.prev = torch.is_autocast_cpu_enabled()
        self.prev_fast_dtype = core.get_autocast_dtype()
        torch.set_autocast_cpu_enabled(self._enabled)
        core.set_autocast_dtype(self.fast_dtype)
        torch.autocast_increment_nesting()

    def __exit__(self, *args):
        # Drop the cache when we exit to a nesting level that's outside any instance of autocast.
        if torch.autocast_decrement_nesting() == 0:
            core.clear_autocast_cache()
        torch.set_autocast_cpu_enabled(self.prev)
        core.set_autocast_dtype(self.prev_fast_dtype)
        return False


torch.cpu.amp.autocast = _autocast
