Add unified memory APIs for torch.accelerator #152932

https://github.com/pytorch/pytorch/pull/152932

为torch.accelerator增加了如下API：

  • empty_cache
  • max_memory_allocated
  • max_memory_reserved
  • memory_allocated
  • memory_reserved
  • memory_stats
  • reset_accumulated_memory_stats
  • reset_peak_memory_stats

这个pr进一步说明了PyTorch正在去CUDA化，从而转变为一系列加速器提供计算框架，即使他们的github主页仍然写着：

Tensors and Dynamic neural networks in Python with strong GPU acceleration

这些用法很长一段时间是为cuda提供的，调用方式主要为torch.cuda.empty_cache等，接口主要作用是查看device上的显存占用如memory_allocated和memory_reserved，以及清理和重置相关状态。 一起来看下相关函数的调用流程，以cuda为例，从源码分析这个接口是怎么一步一步获取底层的状态的。首先所有memory相关的数据都是从memory_stats这个字典内提取对于的value

def memory_reserved(device_index: _device_t = None, /) -> int:
    r"""Return the current :ref:`accelerator<accelerators>` device memory managed by the caching allocator
    in bytes for a given device index.
    Args:
        device_index (:class:`torch.device`, str, int, optional): the index of the device to target.
            If not given, use :func:`torch.accelerator.current_device_index` by default.
            If a :class:`torch.device` or str is provided, its type must match the current
            :ref:`accelerator<accelerators>` device type.
    """
    return memory_stats(device_index).get("reserved_bytes.all.current", 0)

而memory_stas通过底层_accelerator_getDeviceStats接口获取相关device的状态

stats = torch._C._accelerator_getDeviceStats(device_index)
torch._C._accelerator_getDeviceStats是一个pybind绑定的C接口函数
at::accelerator::getDeviceStats(device_index);

继续看调用栈

TORCH_API inline at::CachingDeviceAllocator::DeviceStats getDeviceStats(
    c10::DeviceIndex device_index) {
  const auto device_type = getAccelerator(true).value();
  return at::getDeviceAllocator(device_type)->getDeviceStats(device_index);
}

以cuda为例

inline c10::CachingDeviceAllocator::DeviceStats getDeviceStats(
    c10::DeviceIndex device) {
  return get()->getDeviceStats(device);
}

其中get方法来自allocator，allocator是CUDA Cacing Allocator中的单例实现，所有方法都需要通过该接口实现

inline CUDAAllocator* get() {
  return allocator.load();
}

cuda的getDeviceStats通过下面代码实现

  DeviceStats getDeviceStats(c10::DeviceIndex device) override {
    assertValidDevice(device);
    return device_allocator[device]->getStats();
  }

总之，现在，所有加速器都可以通过统一的接口返回从PyTorch申请的内存，以及对其状态进行管理，大一统进行中……