uint32 overflow in CUDA thread index silently corrupts output when `E * M_out >= 2^32`

[haohq19]

Description

The CUDA kernels in weighting_cuda.cu compute the global thread index as:

const int64_t thread_idx = blockIdx.x * blockDim.x + threadIdx.x;

blockIdx.x and blockDim.x are CUDA built-in unsigned int (32-bit). The multiplication is evaluated in uint32 before the result is widened to int64_t, so when the total number of threads (E * M_out) exceeds 2^32, the high bits are silently dropped.

This affects all four kernels in weighting_cuda.cu:

• spline_weighting_fw_kernel (line 18)
• spline_weighting_bw_x_kernel (line 81)
• spline_weighting_bw_weight_kernel (line 146)
• spline_weighting_bw_basis_kernel (line 206)

When E * M_out >= 2^32:

1. Threads with true global index >= 2^32 have their thread_idx wrap modulo 2^32, computing the spline message for a wrong edge and overwriting the correct result that the early thread already wrote.
2. The output elements corresponding to indices [2^32, E * M_out) are not computed by any thread and occasionally contain NaN causing training crash.

Reproduce

Environment: Ubuntu 22.04, 4090 GPU, PyTorch 2.6.0, torch_spline_conv: 1.2.2, CUDA 12.4

"""Reproduce uint32 thread-index overflow in spline_weighting."""
import torch
from torch_spline_conv import spline_weighting

device = torch.device("cuda")
UINT32_MAX = 1 << 32

# Use out_ch=256 so overflow threshold is E_crit = 2^32/256 = 2^24 = 16,777,216
in_ch, out_ch, K, S = 1, 256, 4, 1
delta = 1024
E_crit = UINT32_MAX // out_ch
E = E_crit + delta   # just over the threshold

torch.manual_seed(42)
x = torch.randn(E, in_ch, device=device)
weight = torch.randn(K, in_ch, out_ch, device=device)
basis = torch.ones(E, S, device=device)
weight_index = torch.zeros(E, S, dtype=torch.long, device=device)

out = spline_weighting(x, weight, basis, weight_index)
torch.cuda.synchronize()

# Reference for just the buggy tail
ref_tail = spline_weighting(
    x[E_crit:], weight, basis[E_crit:], weight_index[E_crit:]
)
torch.cuda.synchronize()

buggy_tail = out[E_crit:]
max_err = (buggy_tail - ref_tail).abs().max().item()
print(f"Tail max abs error: {max_err:.6e}")   # expect ~10, not ~0
print(f"Tail matches ref: {torch.allclose(buggy_tail, ref_tail, atol=1e-5)}")

Output:

Tail max abs error: 1.063440e+01
Tail matches ref: False