PKGBUILD


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

# Maintainer: Daniel Bershatsky <bepshatsky@yandex.ru>

pkgname=('cutlass' 'python-cutlass')
pkgver=3.4.0
pkgrel=1
pkgdesc='CUDA Templates for Linear Algebra Subroutines'
arch=('x86_64')
url='https://github.com/NVIDIA/cutlass'
license=('BSD-3-Clause')
groups=()
depends=('cuda' 'cudnn')
makedepends=('cmake' 'git' 'ninja' 'pybind11' 'python-build'
             'python-installer' 'python-setuptools' 'python-wheel')
optdepends=()
source=("$pkgname-$pkgver.tar.gz::https://github.com/NVIDIA/$pkgname/archive/refs/tags/v$pkgver.tar.gz"
        'cutlass.diff')
sha256sums=('49f4b854acc2a520126ceefe4f701cfe8c2b039045873e311b1f10a8ca5d5de1'
            'SKIP')

prepare() {
    cd cutlass-$pkgver
    patch -p1 -i ../cutlass.diff
}

build() {
    export PATH="/opt/cuda/bin:$PATH"
    cmake -S $pkgname-$pkgver -B $pkgname-$pkgver/build -G Ninja \
        -DCUTLASS_NVCC_ARCHS='70;75;80;86;89;90;90a' \
        -DCUTLASS_ENABLE_GTEST_UNIT_TESTS=OFF \
        -DCUTLASS_ENABLE_TESTS=OFF \
        -DCUTLASS_ENABLE_CUBLAS=ON \
        -DCUTLASS_ENABLE_CUDNN=ON \
        -DCUTLASS_ENABLE_EXAMPLES=OFF \
        -DCUTLASS_INSTALL_TESTS=OFF \
        -DCUTLASS_UNITY_BUILD_ENABLED=ON
    cmake --build $pkgname-$pkgver/build

    # Make python-pycute wheel.
    cd $srcdir/$pkgname-$pkgver/python
    rm -rf build *.egg-info
    python setup_pycute.py bdist_wheel

    # Make python-cutlass-library wheel.
    cd $srcdir/$pkgname-$pkgver/python
    rm -rf build *.egg-info
    python setup_library.py bdist_wheel

    # Make python-cutlass wheel.
    cd $srcdir/$pkgname-$pkgver/python
    rm -rf build *.egg-info
    python -m build -nw .
}

package_cutlass() {
    cmake --install $srcdir/$pkgname-$pkgver/build --prefix $pkgdir/usr
    install -Dm644 \
        $srcdir/$pkgname-$pkgver/LICENSE.txt \
        $pkgdir/usr/share/licenses/$pkgname/LICENSE
    rm -rfv $pkgdir/usr/test
}

package_python-cutlass() {
    pkgdesc='Python interface to cutlass library'
    arch=('any')
    url='https://github.com/NVIDIA/cutlass'
    license=('BSD-3-Clause')
    groups=('nvidia')
    depends=('python-pydot' 'python-scikit-build' 'python-treelib')
    optdepends=()

    python -m installer \
        --compile-bytecode 1 \
        --destdir $pkgdir \
        cutlass-$pkgver/python/dist/pycute-*.whl

    python -m installer \
        --compile-bytecode 1 \
        --destdir $pkgdir \
        cutlass-$pkgver/python/dist/cutlass_library-*.whl

    python -m installer \
        --compile-bytecode 1 \
        --destdir $pkgdir \
        cutlass-$pkgver/python/dist/cutlass-*.whl
}