| 123456789101112131415161718192021222324252627 |
- #pragma once
- #include <c10/core/DeviceType.h>
- #include <c10/macros/Macros.h>
- #include <ATen/detail/MTIAHooksInterface.h>
- #include <optional>
- // This file defines the top level Accelerator concept for PyTorch.
- // A device is an accelerator per the definition here if:
- // - It is mutually exclusive with all other accelerators
- // - It performs asynchronous compute via a Stream/Event system
- // - It provides a set of common APIs as defined by AcceleratorHooksInterface
- //
- // As of today, accelerator devices are (in no particular order):
- // CUDA, MTIA, PrivateUse1
- // We want to add once all the proper APIs are supported and tested:
- // HIP, MPS, XPU
- namespace at {
- // Ensures that only one accelerator is available (at
- // compile time if possible) and return it.
- // When checked is true, the returned optional always has a value.
- TORCH_API std::optional<c10::DeviceType> getAccelerator(bool checked = false);
- } // namespace at
|