summaryrefslogtreecommitdiffstats
path: root/tools/cuda-bridge/embed_ptx.py
blob: 71b0e5f9b13810d8015f5511cd7773ae0ba74cfd (plain)
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
#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
#
# embed_ptx.py - turn a PTX text file into a NUL-terminated C byte array.
# Used by the Makefile so cuda_bridge carries its kernels without any runtime
# dependency on files or on a CUDA toolkit.

import sys


def main() -> int:
    if len(sys.argv) != 3:
        print("usage: embed_ptx.py <in.ptx> <out.h>", file=sys.stderr)
        return 2
    data = open(sys.argv[1], "rb").read()
    if b"\x00" in data:
        print("embed_ptx: input contains NUL", file=sys.stderr)
        return 1
    rows = []
    for off in range(0, len(data), 12):
        chunk = data[off:off + 12]
        rows.append("\t" + " ".join(f"0x{b:02x}," for b in chunk))
    body = "\n".join(rows)
    text = (
        "/* SPDX-License-Identifier: BSD-3-Clause */\n"
        "/* generated by embed_ptx.py from %s - do not edit */\n"
        "static const unsigned char kernels_ptx[] = {\n"
        "%s\n"
        "\t0x00\n"
        "};\n"
        "static const unsigned int kernels_ptx_len = %d;\n"
    ) % (sys.argv[1], body, len(data))
    with open(sys.argv[2], "w") as fh:
        fh.write(text)
    return 0


if __name__ == "__main__":
    sys.exit(main())