leetgpu_conv3d_kernel.cu
#include "solve.h"
#include <cuda_runtime.h>
#include <stdio.h>
__global__ void conv3d_kernel(const float* input, const float* kernel, float* output,
int input_depth, int input_rows, int input_cols,
int kernel_depth, int kernel_rows, int kernel_cols,
int output_depth, int output_rows, int output_cols) {
int ocol = blockIdx.x * blockDim.x + threadIdx.x;
int orow = blockIdx.y * blockDim.y + threadIdx.y;
int od = blockIdx.z * blockDim.z + threadIdx.z;
if (ocol < output_cols && orow < output_rows && od < output_depth) {
float sum = 0.0f;
for (int kd = 0; kd < kernel_depth; kd++) {
for (int kr = 0; kr < kernel_rows; kr++) {
for (int kc = 0; kc < kernel_cols; kc++) {
int id = od + kd;
int ir = orow + kr;
int ic = ocol + kc;
int input_idx = id * (input_rows * input_cols) + ir * input_cols + ic;
int kernel_idx = kd * (kernel_rows * kernel_cols) + kr * kernel_cols + kc;
sum += input[input_idx] * kernel[kernel_idx];
}
}
}
int output_idx = od * (output_rows * output_cols) + orow * output_cols + ocol;
output[output_idx] = sum;
}
}
void solve(const float* input, const float* kernel, float* output,
int input_depth, int input_rows, int input_cols,
int kernel_depth, int kernel_rows, int kernel_cols) {
int output_depth = input_depth - kernel_depth + 1;
int output_rows = input_rows - kernel_rows + 1;
int output_cols = input_cols - kernel_cols + 1;
size_t input_size = input_depth * input_rows * input_cols * sizeof(float);
size_t kernel_size = kernel_depth * kernel_rows * kernel_cols * sizeof(float);
size_t output_size = output_depth * output_rows * output_cols * sizeof(float);
float *d_input, *d_kernel, *d_output;
cudaMalloc((void**)&d_input, input_size);
cudaMalloc((void**)&d_kernel, kernel_size);
cudaMalloc((void**)&d_output, output_size);
cudaMemcpy(d_input, input, input_size, cudaMemcpyHostToDevice);
cudaMemcpy(d_kernel, kernel, kernel_size, cudaMemcpyHostToDevice);
dim3 blockDim(8, 8, 8);
dim3 gridDim((output_cols + blockDim.x - 1) / blockDim.x,
(output_rows + blockDim.y - 1) / blockDim.y,
(output_depth + blockDim.z - 1) / blockDim.z);
conv3d_kernel<<<gridDim, blockDim>>>(d_input, d_kernel, d_output,
input_depth, input_rows, input_cols,
kernel_depth, kernel_rows, kernel_cols,
output_depth, output_rows, output_cols);
cudaDeviceSynchronize();
cudaMemcpy(output, d_output, output_size, cudaMemcpyDeviceToHost);
cudaFree(d_input);
cudaFree(d_kernel);
cudaFree(d_output);
}
leetgpu_attention.cu
#include "solve.h"
#include <cuda_runtime.h>
#include <math.h>
__global__ void compute_scores(const float* Q, const float* K, float* scores,
int N, int d_model, int d_k) {
int head = blockIdx.z;
int i = blockIdx.y * blockDim.y + threadIdx.y;
int j = blockIdx.x * blockDim.x + threadIdx.x;
if(i < N && j < N) {
float sum = 0.0f;
int offset = head * d_k;
for (int k = 0; k < d_k; k++) {
float q_val = Q[i * d_model + offset + k];
float k_val = K[j * d_model + offset + k];
sum += q_val * k_val;
}
// Scale the dot product by sqrt(d_k)
float scale = 1.0f / sqrtf((float)d_k);
// Store into scores array (each head gets a contiguous N*N block)
scores[ head * N * N + i * N + j ] = sum * scale;
}
}
__global__ void softmax_kernel(float* scores, int N) {
int head = blockIdx.z;
int row = blockIdx.y * blockDim.y + threadIdx.y;
if (row < N) {
int base = head * N * N + row * N;
float max_val = -1e20f;
for (int j = 0; j < N; j++) {
float val = scores[ base + j ];
if(val > max_val)
max_val = val;
}
float sum = 0.0f;
for (int j = 0; j < N; j++) {
float exp_val = expf(scores[ base + j ] - max_val);
scores[ base + j ] = exp_val;
sum += exp_val;
}
for (int j = 0; j < N; j++) {
scores[ base + j ] /= sum;
}
}
}
__global__ void compute_output_kernel(const float* scores, const float* V, float* out,
int N, int d_model, int d_k) {
int head = blockIdx.z;
int i = blockIdx.y * blockDim.y + threadIdx.y;
int k = blockIdx.x * blockDim.x + threadIdx.x; // k in [0, d_k)
if (i < N && k < d_k) {
float sum = 0.0f;
int offset = head * d_k;
for (int j = 0; j < N; j++) {
float attn = scores[ head * N * N + i * N + j ];
float v_val = V[j * d_model + offset + k];
sum += attn * v_val;
}
out[i * d_model + offset + k] = sum;
}
}
void solve(const float* Q, const float* K, const float* V,
float* output, int N, int d_model, int h) {
int d_k = d_model / h;
float *d_Q, *d_K, *d_V, *d_scores, *d_output;
size_t modelBytes = N * d_model * sizeof(float);
size_t scoreBytes = h * N * N * sizeof(float);
cudaMalloc((void**)&d_Q, modelBytes);
cudaMalloc((void**)&d_K, modelBytes);
cudaMalloc((void**)&d_V, modelBytes);
cudaMalloc((void**)&d_scores, scoreBytes);
cudaMalloc((void**)&d_output, modelBytes);
cudaMemcpy(d_Q, Q, modelBytes, cudaMemcpyHostToDevice);
cudaMemcpy(d_K, K, modelBytes, cudaMemcpyHostToDevice);
cudaMemcpy(d_V, V, modelBytes, cudaMemcpyHostToDevice);
dim3 blockDim1(16, 16);
dim3 gridDim1((N + blockDim1.x - 1) / blockDim1.x,
(N + blockDim1.y - 1) / blockDim1.y,
h);
compute_scores<<<gridDim1, blockDim1>>>(d_Q, d_K, d_scores, N, d_model, d_k);
cudaDeviceSynchronize();
dim3 blockDim2(1, 256);
dim3 gridDim2(1, (N + blockDim2.y - 1) / blockDim2.y, h);
softmax_kernel<<<gridDim2, blockDim2>>>(d_scores, N);
cudaDeviceSynchronize();
dim3 blockDim3(16, 16);
dim3 gridDim3((d_k + blockDim3.x - 1) / blockDim3.x,
(N + blockDim3.y - 1) / blockDim3.y,
h);
compute_output_kernel<<<gridDim3, blockDim3>>>(d_scores, d_V, d_output, N, d_model, d_k);
cudaDeviceSynchronize();
cudaMemcpy(output, d_output, modelBytes, cudaMemcpyDeviceToHost);
cudaFree(d_Q);
cudaFree(d_K);
cudaFree(d_V);
cudaFree(d_scores);
cudaFree(d_output);
}