From a6b0de3beb49afbf87b2b81dc4e601ee3a64ecc4 Mon Sep 17 00:00:00 2001 From: "Adam J. Stewart" Date: Mon, 27 Jun 2022 09:21:49 -0700 Subject: [PATCH] py-torch: add M1 GPU support (#31283) --- var/spack/repos/builtin/packages/py-torch/package.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/var/spack/repos/builtin/packages/py-torch/package.py b/var/spack/repos/builtin/packages/py-torch/package.py index 9c7abe41e7c..c4883b95eba 100644 --- a/var/spack/repos/builtin/packages/py-torch/package.py +++ b/var/spack/repos/builtin/packages/py-torch/package.py @@ -6,6 +6,7 @@ import os import sys +from spack.operating_systems.mac_os import macos_version 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('magma', default=not is_darwin, description='Use MAGMA', when='+cuda') 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=cray') 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('magma') enable_or_disable('metal') + enable_or_disable('mps') enable_or_disable('breakpad') enable_or_disable('nccl')