1/* mesh-ptx-link: links LLVM bitcode into one PTX module in the cluster float
2 * mode (f32 FTZ, unfused), keeping only the exported symbols.
3 *
4 *   mesh-ptx-link --mcpu sm_86 [-O0..3] [--debug] [--export SYM]... -o OUT.ptx IN.bc...
5 *
6 * As alpha-rust-mesh scripts/ptx-link.sh (rustc's linker for programs): the
7 * commands of rustc's llvm-bitcode-linker (llvm-link, opt, llc) with the float
8 * mode added. Tools: $MESH_LLVM_LINK, $MESH_OPT, $MESH_LLC, or found on PATH
9 * (also name-21 .. name-15) or under /usr/lib/llvm-N/bin. Any argument it
10 * does not know is refused. */
11#include "link.h"
12
13#include <errno.h>
14#include <stdio.h>
15#include <stdlib.h>
16#include <string.h>
17#include <sys/wait.h>
18#include <unistd.h>
19
20static int run(char **argv) {
21    pid_t pid = fork();
22    if (pid < 0) { fprintf(stderr, "mesh-ptx-link: fork: %s\n", strerror(errno)); return 1; }
23    if (pid == 0) {
24        execv(argv[0], argv);
25        fprintf(stderr, "mesh-ptx-link: %s: %s\n", argv[0], strerror(errno));
26        _exit(127);
27    }
28    int st;
29    while (waitpid(pid, &st, 0) < 0)
30        if (errno != EINTR) return 1;
31    if (WIFEXITED(st) && WEXITSTATUS(st) == 0) return 0;
32    fprintf(stderr, "mesh-ptx-link: %s failed\n", argv[0]);
33    return 1;
34}
35
36static int usage(void) {
37    fprintf(stderr, "usage: mesh-ptx-link --mcpu sm_NN [-O0..3] [--debug] [--export SYM]... -o OUT.ptx IN.bc...\n");
38    return 2;
39}
40
41int main(int argc, char **argv) {
42    const char **inputs = calloc((size_t)argc, sizeof *inputs), **exports = calloc((size_t)argc, sizeof *exports);
43    if (!inputs || !exports) return 1;
44    mesh_ptx_link l = {.opt = 3};
45    size_t ni = 0, ne = 0;
46    for (int i = 1; i < argc; i++) {
47        const char *a = argv[i];
48        if ((!strcmp(a, "--mcpu") || !strcmp(a, "--target-cpu")) && i + 1 < argc) l.cpu = argv[++i];
49        else if ((!strcmp(a, "-o") || !strcmp(a, "--output")) && i + 1 < argc) l.output = argv[++i];
50        else if (!strcmp(a, "--export") && i + 1 < argc) exports[ne++] = argv[++i];
51        else if (!strcmp(a, "--debug")) l.debug = true;
52        else if (a[0] == '-' && a[1] == 'O' && a[2] >= '0' && a[2] <= '3' && !a[3]) l.opt = a[2] - '0';
53        else if (a[0] == '-') { fprintf(stderr, "mesh-ptx-link: unknown option %s\n", a); return usage(); }
54        else inputs[ni++] = a;
55    }
56    if (!l.output || !l.cpu || ni == 0) { fprintf(stderr, "mesh-ptx-link: need -o, --mcpu and inputs\n"); return usage(); }
57    if (!mesh_ptx_link_valid_arch(l.cpu)) { fprintf(stderr, "mesh-ptx-link: --mcpu %s: want sm_NN (e.g. sm_86, sm_121)\n", l.cpu); return 1; }
58    l.inputs = inputs;
59    l.ninputs = ni;
60    if (ne) { l.exports = exports; l.nexports = ne; }
61    char *llvm_link = mesh_ptx_link_find_tool("llvm-link", "MESH_LLVM_LINK");
62    char *opt = mesh_ptx_link_find_tool("opt", "MESH_OPT");
63    char *llc = mesh_ptx_link_find_tool("llc", "MESH_LLC");
64    const char *missing = !llvm_link ? "llvm-link ($MESH_LLVM_LINK)" : !opt ? "opt ($MESH_OPT)" : !llc ? "llc ($MESH_LLC)" : NULL;
65    if (missing) { fprintf(stderr, "mesh-ptx-link: no %s: install LLVM's tools\n", missing); return 1; }
66    l.llvm_link = llvm_link;
67    l.opt_tool = opt;
68    l.llc = llc;
69    char *syms = mesh_ptx_link_symbols(&l), *symfile = mesh_ptx_link_side(&l, ".symbols.txt");
70    FILE *f = fopen(symfile, "w");
71    if (!f || fputs(syms, f) < 0 || fclose(f) != 0) { fprintf(stderr, "mesh-ptx-link: %s: %s\n", symfile, strerror(errno)); return 1; }
72    char **steps[3] = {mesh_ptx_link_argv_link(&l), mesh_ptx_link_argv_opt(&l), mesh_ptx_link_argv_llc(&l)};
73    int rc = 0;
74    for (int i = 0; i < 3 && rc == 0; i++) rc = run(steps[i]);
75    for (int i = 0; i < 3; i++) mesh_ptx_link_free_argv(steps[i]);
76    free(syms), free(symfile), free(llvm_link), free(opt), free(llc), free(inputs), free(exports);
77    return rc;
78}