cuda: fixes races in mmid and mmf

This commit is contained in:
Aman Gupta
2026-09-06 14:01:34 +08:00
parent 4d9176092d
commit 86e7fc4ffc
2 changed files with 20 additions and 0 deletions
+19
View File
@@ -143,6 +143,7 @@ static __global__ void mul_mat_f(
if (threadIdx.x == 0) {
slot_map[j] = -1;
}
__syncwarp();
if (col_base + j >= ncols_dst_total) {
continue;
@@ -171,10 +172,12 @@ static __global__ void mul_mat_f(
tile_A A[ntA][warp_size / tile_A::J];
#pragma unroll
for (int itA = 0; itA < ntA; ++itA) {
__syncwarp();
#pragma unroll
for (int i = 0; i < tile_A::I; ++i) {
tile_xy[i*tile_k_padded + threadIdx.x] = x[(itA*tile_A::I + i)*stride_row + col];
}
__syncwarp();
#pragma unroll
for (int k0 = 0; k0 < warp_size; k0 += tile_A::J) {
load_ldmatrix(A[itA][k0/tile_A::J], tile_xy + k0, tile_k_padded);
@@ -183,6 +186,7 @@ static __global__ void mul_mat_f(
#pragma unroll
for (int itB = 0; itB < ntB; ++itB) {
__syncwarp();
if constexpr (std::is_same_v<T, float>) {
#pragma unroll
for (int j0 = 0; j0 < tile_B::I; ++j0) {
@@ -212,6 +216,7 @@ static __global__ void mul_mat_f(
} else {
static_assert(std::is_same_v<T, void>, "unsupported type");
}
__syncwarp();
#pragma unroll
for (int k0 = 0; k0 < warp_size; k0 += tile_B::J) {
tile_B B;
@@ -229,6 +234,8 @@ static __global__ void mul_mat_f(
if (nwarps > 1) {
__syncthreads();
} else {
__syncwarp();
}
#pragma unroll
for (int itB = 0; itB < ntB; ++itB) {
@@ -245,6 +252,8 @@ static __global__ void mul_mat_f(
if (nwarps > 1) {
__syncthreads();
} else {
__syncwarp();
}
#pragma unroll
@@ -382,10 +391,12 @@ static __global__ void mul_mat_f_ids(
tile_A A[ntA][warp_size / tile_A::J];
#pragma unroll
for (int itA = 0; itA < ntA; ++itA) {
__syncwarp();
#pragma unroll
for (int i = 0; i < tile_A::I; ++i) {
tile_xy[i*tile_k_padded + threadIdx.x] = x[(itA*tile_A::I + i)*stride_row + col];
}
__syncwarp();
#pragma unroll
for (int k0 = 0; k0 < warp_size; k0 += tile_A::J) {
load_ldmatrix(A[itA][k0/tile_A::J], tile_xy + k0, tile_k_padded);
@@ -419,6 +430,7 @@ static __global__ void mul_mat_f_ids(
int next_buf = 1;
#pragma unroll
for (int itB = 0; itB < ntB; ++itB) {
__syncwarp();
#pragma unroll
for (int j0 = 0; j0 < tile_B::I; ++j0) {
tile_xy[j0*tile_k_padded + threadIdx.x] = vals_buf[curr_buf][j0];
@@ -428,6 +440,7 @@ static __global__ void mul_mat_f_ids(
gather_tile(itB + 1, vals_buf[next_buf]);
}
__syncwarp();
#pragma unroll
for (int k0 = 0; k0 < warp_size; k0 += tile_B::J) {
tile_B B;
@@ -472,6 +485,7 @@ static __global__ void mul_mat_f_ids(
int next_buf = 1;
#pragma unroll
for (int itB = 0; itB < ntB; ++itB) {
__syncwarp();
#pragma unroll
for (int j0 = 0; j0 < tile_B::I; ++j0) {
const float2 tmp = vals_buf[curr_buf][j0];
@@ -482,6 +496,7 @@ static __global__ void mul_mat_f_ids(
gather_tile(itB + 1, vals_buf[next_buf]);
}
__syncwarp();
#pragma unroll
for (int k0 = 0; k0 < warp_size; k0 += tile_B::J) {
tile_B B;
@@ -507,6 +522,8 @@ static __global__ void mul_mat_f_ids(
if (nwarps > 1) {
__syncthreads();
} else {
__syncwarp();
}
#pragma unroll
for (int itB = 0; itB < ntB; ++itB) {
@@ -523,6 +540,8 @@ static __global__ void mul_mat_f_ids(
if (nwarps > 1) {
__syncthreads();
} else {
__syncwarp();
}
#pragma unroll
+1
View File
@@ -101,6 +101,7 @@ static __global__ void mm_ids_helper(
}
}
nex_prev = warp_reduce_sum<warp_size>(nex_prev);
__syncwarp();
for (int itc = threadIdx.x; itc < it_compact; itc += warp_size) {
const mm_ids_helper_store store_it = store[itc];