pkgs/development/python-modules/causal-conv1d/default.nix

   1 {
   2   lib,
   3   buildPythonPackage,
   4   fetchFromGitHub,
   5   ninja,
   6   setuptools,
   7   torch,
   8   cudaPackages,
   9   rocmPackages,
  10   config,
  11   cudaSupport ? config.cudaSupport,
  12   which,
  13 }:
  14
  15 buildPythonPackage rec {
  16   pname = "causal-conv1d";
  17   version = "1.4.0";
  18   pyproject = true;
  19
  20   src = fetchFromGitHub {
  21     owner = "Dao-AILab";
  22     repo = "causal-conv1d";
  23     rev = "refs/tags/v${version}";
  24     hash = "sha256-p5x5u3zEmEMN3mWd88o3jmcpKUnovTvn7I9jIOj/ie0=";
  25   };
  26
  27   build-system = [
  28     ninja
  29     setuptools
  30     torch
  31   ];
  32
  33   nativeBuildInputs = [ which ];
  34
  35   buildInputs = (
  36     lib.optionals cudaSupport (
  37       with cudaPackages;
  38       [
  39         cuda_cudart # cuda_runtime.h, -lcudart
  40         cuda_cccl
  41         libcusparse # cusparse.h
  42         libcusolver # cusolverDn.h
  43         cuda_nvcc
  44         libcublas
  45       ]
  46     )
  47   );
  48
  49   dependencies = [
  50     torch
  51   ];
  52
  53   # pytest tests not enabled due to nvidia GPU dependency
  54   pythonImportsCheck = [ "causal_conv1d" ];
  55
  56   env = {
  57     CAUSAL_CONV1D_FORCE_BUILD = "TRUE";
  58   } // lib.optionalAttrs cudaSupport { CUDA_HOME = "${lib.getDev cudaPackages.cuda_nvcc}"; };
  59
  60   meta = with lib; {
  61     description = "Causal depthwise conv1d in CUDA with a PyTorch interface";
  62     homepage = "https://github.com/Dao-AILab/causal-conv1d";
  63     license = licenses.bsd3;
  64     maintainers = with maintainers; [ cfhammill ];
  65     # The package requires CUDA or ROCm, the ROCm build hasn't
  66     # been completed or tested, so broken if not using cuda.
  67     broken = !cudaSupport;
  68   };
  69 }