1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
|
{
lib,
buildPythonPackage,
fetchFromGitHub,
ninja,
setuptools,
torch,
cudaPackages,
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.4";
pyproject = true;
src = fetchFromGitHub {
owner = "Dao-AILab";
repo = "causal-conv1d";
tag = "v${version}";
hash = "sha256-ELuvnKP2g1I2SuaWWiibXh/oDzp4n0vXkm4oeNPOdIw=";
};
build-system = [
ninja
setuptools
torch
];
nativeBuildInputs = [ which ] ++ lib.optionals rocmSupport [ rocmPackages.clr ];
buildInputs =
lib.optionals cudaSupport (
with cudaPackages;
[
cuda_cudart # cuda_runtime.h, -lcudart
cuda_cccl
libcusparse # cusparse.h
libcusolver # cusolverDn.h
cuda_nvcc
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 GPU dependency
pythonImportsCheck = [ "causal_conv1d" ];
env = {
CAUSAL_CONV1D_FORCE_BUILD = "TRUE";
}
// 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 either CUDA or ROCm.
# It doesn't work without either, nor with both.
broken = cudaSupport == rocmSupport;
};
}
|
