[hardware][misc] introduce platform abstraction (#6080)

This commit is contained in:
youkaichao
2024-07-02 20:12:22 -07:00
committed by GitHub
parent 9d6a8daa87
commit 482045ee77
16 changed files with 113 additions and 29 deletions

View File

@@ -0,0 +1,18 @@
from typing import Optional
import torch
from .interface import Platform, PlatformEnum
current_platform: Optional[Platform]
if torch.version.cuda is not None:
from .cuda import CudaPlatform
current_platform = CudaPlatform()
elif torch.version.hip is not None:
from .rocm import RocmPlatform
current_platform = RocmPlatform()
else:
current_platform = None
__all__ = ['Platform', 'PlatformEnum', 'current_platform']