Host SIGSEGV in cuLaunchKernel when cuModuleUnload runs concurrently, with Tile IR modules only.

Hello, I'm not deeply familiar with CUDA internals, so I'm not certain this is Tile IR specific, but I couldn't manage to reproduce it with a regular CUDA kernel, so I figured I'd post here. Happy to be told this is the wrong place.

Two threads doing load → launch → unload on the same context segfault inside libcuda, on the host, below cuLaunchKernel. Usually within a few seconds.

The loop is just:

cuModuleLoadData(&m, cubin);
cuModuleGetFunction(&f, m, "repro_kernel");
cuLaunchKernel(f, 1,1,1, 1,1,1, 0, 0, args, 0);
cuCtxSynchronize();               // nothing from this module is in flight
if (n & 3) cuModuleUnload(m);     // keep one in four

Environment: RTX 5090 (sm_120), driver 595.84, CUDA 13.2, cuda-tile 1.5.0

Here is the full reproducer:

import sys

import cuda.tile as ct
from cuda.tile.compilation import (
    ArrayConstraint,
    CallingConvention,
    KernelSignature,
    export_kernel,
)

M = N = 128


@ct.kernel
def repro_kernel(a, b, c):
    ct.store(c, (0, 0), ct.load(a, (0, 0), (M, N)) + ct.load(b, (0, 0), (M, N)))


arr = lambda: ArrayConstraint(
    ct.float32,
    2,
    index_dtype=ct.int32,
    stride_lower_bound_incl=0,
    alias_groups=[],
    may_alias_internally=False,
)
sig = KernelSignature(
    [arr(), arr(), arr()], CallingConvention.cutile_python_v2(), symbol="repro_kernel"
)
out = sys.argv[1] if len(sys.argv) > 1 else "kernel.cubin"
export_kernel(repro_kernel, [sig], out, gpu_code="sm_120", output_format="cubin")
print("wrote", out)

#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

typedef struct CUctx_st *CUcontext;
typedef struct CUmod_st *CUmodule;
typedef struct CUfunc_st *CUfunction;
typedef unsigned long long CUdeviceptr;
extern "C" {
int cuInit(unsigned);
int cuDeviceGet(int *, int);
int cuCtxCreate_v2(CUcontext *, unsigned, int);
int cuCtxSetCurrent(CUcontext);
int cuMemAlloc_v2(CUdeviceptr *, size_t);
int cuModuleLoadData(CUmodule *, const void *);
int cuModuleGetFunction(CUfunction *, CUmodule, const char *);
int cuModuleUnload(CUmodule);
int cuCtxSynchronize(void);
int cuLaunchKernel(CUfunction, unsigned, unsigned, unsigned, unsigned, unsigned,
                   unsigned, unsigned, void *, void **, void **);
int cudaDeviceSynchronize(void);
}

static const int M = 128;
static void *g_cubin;
static CUcontext g_ctx;
static CUdeviceptr g_buf[3];
static int g_dim[4] = {M, M, M, 1};  // shape0, shape1, stride0, stride1
static void *g_args[15];             // per array: pointer, then those four
static unsigned long long g_launch, g_load, g_unload;
#define ADD(v) __atomic_add_fetch(&(v), 1, __ATOMIC_RELAXED)

static void *worker(void *) {
  cuCtxSetCurrent(g_ctx);
  for (int n = 0;; ++n) {
    CUmodule m;
    if (cuModuleLoadData(&m, g_cubin)) continue;
    ADD(g_load);
    CUfunction f;
    if (cuModuleGetFunction(&f, m, "repro_kernel")) { 
	cuModuleUnload(m);
	continue;
    }
    if (!cuLaunchKernel(f, 1, 1, 1, 1, 1, 1, 0, 0, g_args, 0)) ADD(g_launch);
    
    cuCtxSynchronize();
    
    if ((n & 3) && !cuModuleUnload(m)) ADD(g_unload);
  }
}

int main(int argc, char **argv) {
  const char *path = argc > 1 ? argv[1] : "kernel.cubin";
  int secs = argc > 2 ? atoi(argv[2]) : 60;
  setvbuf(stdout, 0, _IOLBF, 0);

  FILE *fp = fopen(path, "rb");
  if (!fp) { perror(path); return 2; }
  fseek(fp, 0, SEEK_END);
  long sz = ftell(fp);
  rewind(fp);
  g_cubin = malloc(sz);
  if (fread(g_cubin, 1, sz, fp) != (size_t)sz) { perror("read"); return 2; }
  fclose(fp);

  int dev;
  if (cuInit(0) || cuDeviceGet(&dev, 0) || cuCtxCreate_v2(&g_ctx, 0, dev)) {
    fprintf(stderr, "CUDA init failed\n");
    return 2;
  }
  for (int i = 0, a = 0; i < 3; ++i) {    // kernel ABI: per array a pointer,
    cuMemAlloc_v2(&g_buf[i], M * M * 4);  // then shape0, shape1, stride0,
    g_args[a++] = &g_buf[i];              // stride1, as int32
    for (int q = 0; q < 4; ++q) g_args[a++] = &g_dim[q];
  }
  printf("%s (%ld bytes), 2 threads, nothing preloaded\n", path, sz);

  pthread_t t[2];
  for (int i = 0; i < 2; ++i) pthread_create(&t[i], 0, worker, 0);

  for (int i = 0; i < secs; ++i) {
    sleep(1);
    printf("  launches=%llu loads=%llu unloads=%llu\n",
           __atomic_load_n(&g_launch, __ATOMIC_RELAXED),
           __atomic_load_n(&g_load, __ATOMIC_RELAXED),
           __atomic_load_n(&g_unload, __ATOMIC_RELAXED));
  }
  printf("SURVIVED %ds\n", secs);
  _exit(0);
}

Commands:

g++ -O2 -pthread repro.cc -o repro -lcuda
./repro kernel.cubin 20

I also tried with

extern "C" __global__ void repro_kernel(float *a, int as0, int as1, int at0,
                                        int at1, float *b, int bs0, int bs1,
                                        int bt0, int bt1, float *c, int cs0,
                                        int cs1, int ct0, int ct1) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < as0 * as1) c[i] = a[i] + b[i];
}

but it never reproduced the bug on my side.

Start with the repro.cc reproducer and the Python script that generates kernel.cubin, then run the provided g++ and ./repro commands on the stated RTX 5090, driver, CUDA, and cuda-tile versions. Compare the Tile IR module with the regular CUDA kernel case and determine whether the failure is attributable to cuda-tile; done means the failure is reproduced and its ownership or next investigation is established.