From b32cb5765c598e59f2f91d99807b4331f4eba569 Mon Sep 17 00:00:00 2001 From: renjithravindrankannath <94420380+renjithravindrankannath@users.noreply.github.com> Date: Fri, 12 Aug 2022 23:17:20 -0700 Subject: Add new dependencies for rocm variant for py-torch recipe (#32100) * Cmake module path updated for ROCm 5.2 * nccl is already set below for PyTorch 1.6+ * Threadpool is set below for PyTorch 1.6+ --- var/spack/repos/builtin/packages/py-torch/package.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/var/spack/repos/builtin/packages/py-torch/package.py b/var/spack/repos/builtin/packages/py-torch/package.py index e9de9d257b..0ba0736b9d 100644 --- a/var/spack/repos/builtin/packages/py-torch/package.py +++ b/var/spack/repos/builtin/packages/py-torch/package.py @@ -195,6 +195,7 @@ class PyTorch(PythonPackage, CudaPackage): depends_on("rocfft") depends_on("rocblas") depends_on("miopen-hip") + depends_on("rocminfo") # https://github.com/pytorch/pytorch/issues/60332 # depends_on('xnnpack@2022-02-16', when='@1.12:+xnnpack') # depends_on('xnnpack@2021-06-21', when='@1.10:1.11+xnnpack') @@ -427,7 +428,6 @@ class PyTorch(PythonPackage, CudaPackage): env.set("ROCFFT_PATH", self.spec["rocfft"].prefix) env.set("HIPFFT_PATH", self.spec["hipfft"].prefix) env.set("HIPSPARSE_PATH", self.spec["hipsparse"].prefix) - env.set("THRUST_PATH", self.spec["rocthrust"].prefix.include) env.set("HIP_PATH", self.spec["hip"].prefix) env.set("HIPRAND_PATH", self.spec["rocrand"].prefix) env.set("ROCRAND_PATH", self.spec["rocrand"].prefix) @@ -437,6 +437,8 @@ class PyTorch(PythonPackage, CudaPackage): env.set("HIPCUB_PATH", self.spec["hipcub"].prefix) env.set("ROCTHRUST_PATH", self.spec["rocthrust"].prefix) env.set("ROCTRACER_PATH", self.spec["roctracer-dev"].prefix) + if self.spec.satisfies("^hip@5.2.0:"): + env.set("CMAKE_MODULE_PATH", self.spec["hip"].prefix.lib.cmake.hip) enable_or_disable("cudnn") if "+cudnn" in self.spec: -- cgit v1.2.3-70-g09d2