diff --git a/pkgs/development/python-modules/causal-conv1d/default.nix b/pkgs/development/python-modules/causal-conv1d/default.nix index ec6a888cc9ea..8902d04fc0e9 100644 --- a/pkgs/development/python-modules/causal-conv1d/default.nix +++ b/pkgs/development/python-modules/causal-conv1d/default.nix @@ -9,19 +9,21 @@ rocmPackages, config, cudaSupport ? config.cudaSupport, + rocmSupport ? config.rocmSupport, + rocmGpuTargets ? rocmPackages.clr.localGpuTargets or rocmPackages.clr.gpuTargets, which, }: buildPythonPackage rec { pname = "causal-conv1d"; - version = "1.5.2"; + version = "1.5.4"; pyproject = true; src = fetchFromGitHub { owner = "Dao-AILab"; repo = "causal-conv1d"; tag = "v${version}"; - hash = "sha256-B2I5QiJl0p5d1BeQcMbJBAYUb10HzqFd88QMM8Rerm0="; + hash = "sha256-ELuvnKP2g1I2SuaWWiibXh/oDzp4n0vXkm4oeNPOdIw="; }; build-system = [ @@ -30,9 +32,9 @@ buildPythonPackage rec { torch ]; - nativeBuildInputs = [ which ]; + nativeBuildInputs = [ which ] ++ lib.optionals rocmSupport [ rocmPackages.clr ]; - buildInputs = ( + buildInputs = lib.optionals cudaSupport ( with cudaPackages; [ @@ -44,26 +46,47 @@ buildPythonPackage rec { libcublas ] ) - ); + ++ lib.optionals rocmSupport ( + with rocmPackages; + [ + rocm-core + rocm-device-libs + rocm-runtime + rocm-comgr + hipblas + rocblas + hipcub + rocprim + ] + ); dependencies = [ torch ]; - # pytest tests not enabled due to nvidia GPU dependency + # pytest tests not enabled due to GPU dependency pythonImportsCheck = [ "causal_conv1d" ]; env = { CAUSAL_CONV1D_FORCE_BUILD = "TRUE"; } - // lib.optionalAttrs cudaSupport { CUDA_HOME = "${lib.getDev cudaPackages.cuda_nvcc}"; }; + // lib.optionalAttrs cudaSupport { CUDA_HOME = "${lib.getDev cudaPackages.cuda_nvcc}"; } + // lib.optionalAttrs rocmSupport { + ROCM_PATH = "${rocmPackages.clr}"; + HIP_ARCHITECTURES = builtins.concatStringsSep "," rocmGpuTargets; + CPLUS_INCLUDE_PATH = lib.makeSearchPath "include" [ + rocmPackages.hipcub + rocmPackages.rocprim + ]; + }; meta = { description = "Causal depthwise conv1d in CUDA with a PyTorch interface"; homepage = "https://github.com/Dao-AILab/causal-conv1d"; license = lib.licenses.bsd3; - # The package requires CUDA or ROCm, the ROCm build hasn't - # been completed or tested, so broken if not using cuda. - broken = !cudaSupport; + + # The package requires either CUDA or ROCm. + # It doesn't work without either, nor with both. + broken = cudaSupport != rocmSupport; }; }