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

[elogir]

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.