py-torch: add M1 GPU support (#31283)
This commit is contained in:
parent
11d71ca85e
commit
a6b0de3beb
@ -6,6 +6,7 @@
|
|||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
|
|
||||||
|
from spack.operating_systems.mac_os import macos_version
|
||||||
from spack.package import *
|
from spack.package import *
|
||||||
|
|
||||||
|
|
||||||
@ -59,6 +60,7 @@ class PyTorch(PythonPackage, CudaPackage):
|
|||||||
variant('kineto', default=True, description='Use Kineto profiling library', when='@1.8:')
|
variant('kineto', default=True, description='Use Kineto profiling library', when='@1.8:')
|
||||||
variant('magma', default=not is_darwin, description='Use MAGMA', when='+cuda')
|
variant('magma', default=not is_darwin, description='Use MAGMA', when='+cuda')
|
||||||
variant('metal', default=is_darwin, description='Use Metal for Caffe2 iOS build')
|
variant('metal', default=is_darwin, description='Use Metal for Caffe2 iOS build')
|
||||||
|
variant('mps', default=is_darwin and macos_version() >= Version('12.3'), description='Use MPS for macOS build', when='@1.12: platform=darwin')
|
||||||
variant('nccl', default=True, description='Use NCCL', when='+cuda platform=linux')
|
variant('nccl', default=True, description='Use NCCL', when='+cuda platform=linux')
|
||||||
variant('nccl', default=True, description='Use NCCL', when='+cuda platform=cray')
|
variant('nccl', default=True, description='Use NCCL', when='+cuda platform=cray')
|
||||||
variant('nccl', default=True, description='Use NCCL', when='+rocm platform=linux')
|
variant('nccl', default=True, description='Use NCCL', when='+rocm platform=linux')
|
||||||
@ -373,6 +375,7 @@ def enable_or_disable(variant, keyword='USE', var=None, newer=False):
|
|||||||
enable_or_disable('kineto')
|
enable_or_disable('kineto')
|
||||||
enable_or_disable('magma')
|
enable_or_disable('magma')
|
||||||
enable_or_disable('metal')
|
enable_or_disable('metal')
|
||||||
|
enable_or_disable('mps')
|
||||||
enable_or_disable('breakpad')
|
enable_or_disable('breakpad')
|
||||||
|
|
||||||
enable_or_disable('nccl')
|
enable_or_disable('nccl')
|
||||||
|
Loading…
Reference in New Issue
Block a user