Hugging Face's logo

Spaces:
yizhangliu
/
Grounded-Segment-Anything
Running on T4

App Files Community
11
Grounded-Segment-Anything / transformers_4_35_0 /kernels /mra /cuda_kernel.cu
liuyizhang
add transformers_4_35_0
1ce5e18
raw history blame contribute delete
No virus
11.7 kB
#include "cuda_kernel.h"
//////////////////////////////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////////////////////////
__global__ void index_max_cuda_kernel(
float *index_vals, // [batch_size, 32, num_block]
int *indices, // [batch_size, num_block]
float *max_vals, // [batch_size, A_num_block * 32]
float *max_vals_scatter, // [batch_size, 32, num_block]
long batch_size,
long A_num_block,
long B_num_block,
long num_block
) {
long batch_idx = blockIdx.x;
long thread_idx = threadIdx.x;
long num_thread = blockDim.x;
extern __shared__ float buffer[];
int *max_buffer = (int*)buffer;
for (int i = 0; i < A_num_block * 32; i = i + num_thread) {
int idx = i + thread_idx;
if (idx < A_num_block * 32) {
max_buffer[idx] = -1e8;
}
}
__syncthreads();
int *indices_pt = &indices[batch_idx * num_block];
float *index_vals_pt = &index_vals[batch_idx * num_block * 32];
for (int idx_start = 0; idx_start < 32 * num_block; idx_start = idx_start + num_thread) {
int idx = idx_start + thread_idx;
int A_block_idx = indices_pt[idx % num_block] / B_num_block;
atomicMax(&max_buffer[A_block_idx * 32 + idx / num_block], (int)(index_vals_pt[idx] * 1000));
}
__syncthreads();
float *max_vals_pt = &max_vals[batch_idx * A_num_block * 32];
for (int i = 0; i < A_num_block * 32; i = i + num_thread) {
int idx = i + thread_idx;
if (idx < A_num_block * 32) {
max_vals_pt[idx] = (float)max_buffer[idx] / 1000.;
}
}
float *max_vals_scatter_pt = &max_vals_scatter[batch_idx * num_block * 32];
for (int idx_start = 0; idx_start < 32 * num_block; idx_start = idx_start + num_thread) {
int idx = idx_start + thread_idx;
int A_block_idx = indices_pt[idx % num_block] / B_num_block;
max_vals_scatter_pt[idx] = (float)max_buffer[A_block_idx * 32 + idx / num_block] / 1000.;
}
}
__global__ void mm_to_sparse_cuda_kernel(
float *dense_A, // [batch_size, A_num_block, dim, 32]
float *dense_B, // [batch_size, B_num_block, dim, 32]
int *indices, // [batch_size, num_block]
float *sparse_C, // [batch_size, num_block, 32, 32]
long batch_size,
long A_num_block,
long B_num_block,
long dim,
long num_block
) {
long batch_idx = blockIdx.y;
long block_idx = blockIdx.x * blockDim.y + threadIdx.y;
long thread_idx = threadIdx.x;
__shared__ float buffer[4096];
float *A_buffer = &buffer[threadIdx.y * 1024]; // [2, 8, 32]
float *B_buffer = &buffer[threadIdx.y * 1024 + 512]; // [2, 8, 32]
long batch_idx__block_idx = batch_idx * num_block + block_idx;
long AB_block_idx = indices[batch_idx__block_idx];
float *dense_A_pt = &dense_A[(batch_idx * A_num_block + AB_block_idx / B_num_block) * dim * 32];
float *dense_B_pt = &dense_B[(batch_idx * B_num_block + AB_block_idx % B_num_block) * dim * 32];
int reg_1_idx = thread_idx / 8; // [0000000011111111222222223333333344444444555555556666666677777777]
int reg_2_idx = thread_idx % 8; // [0123456701234567012345670123456701234567012345670123456701234567]
float reg_1[8];
float reg_2[8];
float reg_array[16] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
#pragma unroll
for (int i = 0; i < 4; i++) {
A_buffer[i * 64 + thread_idx] = dense_A_pt[i * 64 + thread_idx];
B_buffer[i * 64 + thread_idx] = dense_B_pt[i * 64 + thread_idx];
}
__syncthreads();
#pragma unroll
for (int i = 0; i < 4; i++) {
reg_1[i] = A_buffer[reg_1_idx * 4 + i];
reg_2[i] = B_buffer[reg_2_idx * 4 + i];
}
for (int dim_stride = 1; dim_stride < (dim / 8); dim_stride++) {
#pragma unroll
for (int i = 0; i < 4; i++) {
A_buffer[(dim_stride % 2) * 256 + i * 64 + thread_idx] = dense_A_pt[dim_stride * 256 + i * 64 + thread_idx];
B_buffer[(dim_stride % 2) * 256 + i * 64 + thread_idx] = dense_B_pt[dim_stride * 256 + i * 64 + thread_idx];
}
#pragma unroll
for (int mini_dim_idx = 1; mini_dim_idx < 8; mini_dim_idx++) {
#pragma unroll
for (int i = 0; i < 4; i++) {
reg_1[(mini_dim_idx % 2) * 4 + i] = A_buffer[((dim_stride - 1) % 2) * 256 + mini_dim_idx * 32 + reg_1_idx * 4 + i];
reg_2[(mini_dim_idx % 2) * 4 + i] = B_buffer[((dim_stride - 1) % 2) * 256 + mini_dim_idx * 32 + reg_2_idx * 4 + i];
}
#pragma unroll
for (int i = 0; i < 4; i++) {
#pragma unroll
for (int j = 0; j < 4; j++) {
reg_array[i * 4 + j] += reg_1[((mini_dim_idx - 1) % 2) * 4 + i] * reg_2[((mini_dim_idx - 1) % 2) * 4 + j];
}
}
}
__syncthreads();
#pragma unroll
for (int i = 0; i < 4; i++) {
reg_1[i] = A_buffer[(dim_stride % 2) * 256 + reg_1_idx * 4 + i];
reg_2[i] = B_buffer[(dim_stride % 2) * 256 + reg_2_idx * 4 + i];
}
#pragma unroll
for (int i = 0; i < 4; i++) {
#pragma unroll
for (int j = 0; j < 4; j++) {
reg_array[i * 4 + j] += reg_1[4 + i] * reg_2[4 + j];
}
}
}
#pragma unroll
for (int mini_dim_idx = 1; mini_dim_idx < 8; mini_dim_idx++) {
#pragma unroll
for (int i = 0; i < 4; i++) {
reg_1[(mini_dim_idx % 2) * 4 + i] = A_buffer[256 + mini_dim_idx * 32 + reg_1_idx * 4 + i];
reg_2[(mini_dim_idx % 2) * 4 + i] = B_buffer[256 + mini_dim_idx * 32 + reg_2_idx * 4 + i];
}
#pragma unroll
for (int i = 0; i < 4; i++) {
#pragma unroll
for (int j = 0; j < 4; j++) {
reg_array[i * 4 + j] += reg_1[((mini_dim_idx - 1) % 2) * 4 + i] * reg_2[((mini_dim_idx - 1) % 2) * 4 + j];
}
}
}
#pragma unroll
for (int i = 0; i < 4; i++) {
#pragma unroll
for (int j = 0; j < 4; j++) {
reg_array[i * 4 + j] += reg_1[4 + i] * reg_2[4 + j];
}
}
__syncthreads();
float *C_buffer = &buffer[threadIdx.y * 1024]; // [32, 32]
#pragma unroll
for (int i = 0; i < 4; i++) {
#pragma unroll
for (int j = 0; j < 4; j++) {
C_buffer[(reg_2_idx * 4 + j) * 32 + reg_1_idx * 4 + i] = reg_array[i * 4 + j];
}
}
__syncthreads();
float *sparse_C_pt = &sparse_C[batch_idx__block_idx * 1024];
#pragma unroll
for (int i = 0; i < 16; i++) {
sparse_C_pt[i * 64 + thread_idx] = C_buffer[i * 64 + thread_idx];
}
}
__global__ void sparse_dense_mm_cuda_kernel(
float *sparse_A, // [batch_size, num_block, 32, 32]
int *indices, // [batch_size, num_block]
float *dense_B, // [batch_size, B_num_block, dim, 32]
float *dense_C, // [batch_size, A_num_block, dim, 32]
long batch_size,
long A_num_block,
long B_num_block,
long dim,
long num_block
) {
long batch_idx = blockIdx.y;
long block_idx = blockIdx.x * blockDim.y + threadIdx.y;
long thread_idx = threadIdx.x;
__shared__ float buffer[6144];
float *A_buffer = &buffer[threadIdx.y * 3072]; // [32, 32]
float *B_buffer = &buffer[threadIdx.y * 3072 + 1024]; // [32, 64]
long batch_idx__block_idx = batch_idx * num_block + block_idx;
float *sparse_A_pt = &sparse_A[batch_idx__block_idx * 1024];
#pragma unroll
for (int i = 0; i < 8; i++) {
A_buffer[i * 128 + thread_idx] = sparse_A_pt[i * 128 + thread_idx];
}
long AB_block_idx = indices[batch_idx__block_idx];
float *dense_B_pt = &dense_B[(batch_idx * B_num_block + AB_block_idx % B_num_block) * 32 * dim];
float *dense_C_pt = &dense_C[(batch_idx * A_num_block + AB_block_idx / B_num_block) * 32 * dim];
// [0000000011111111222222223333333344444444555555556666666677777777]
// [0123456701234567012345670123456701234567012345670123456701234567]
int reg_1_idx = thread_idx / 8;
int reg_2_idx = thread_idx % 8;
float reg_1[8];
float reg_2[8];
float reg_array[16];
for (int dim_stride = 0; dim_stride < dim; dim_stride = dim_stride + 64) {
#pragma unroll
for (int i = 0; i < 16; i++) {
B_buffer[i * 128 + thread_idx] = dense_B_pt[dim_stride * 32 + i * 128 + thread_idx];
}
#pragma unroll
for (int i = 0; i < 16; i++) {
reg_array[i] = 0;
}
__syncthreads();
#pragma unroll
for (int i = 0; i < 4; i++) {
reg_1[i] = B_buffer[(reg_1_idx * 4 + i) * 32];
reg_2[i] = A_buffer[reg_2_idx * 4 + i];
}
#pragma unroll
for (int mini_dim_idx = 1; mini_dim_idx < 32; mini_dim_idx++) {
#pragma unroll
for (int i = 0; i < 4; i++) {
reg_1[(mini_dim_idx % 2) * 4 + i] = B_buffer[(reg_1_idx * 4 + i) * 32 + mini_dim_idx];
reg_2[(mini_dim_idx % 2) * 4 + i] = A_buffer[mini_dim_idx * 32 + reg_2_idx * 4 + i];
}
#pragma unroll
for (int i = 0; i < 4; i++) {
#pragma unroll
for (int j = 0; j < 4; j++) {
reg_array[i * 4 + j] += reg_1[((mini_dim_idx - 1) % 2) * 4 + i] * reg_2[((mini_dim_idx - 1) % 2) * 4 + j];
}
}
}
#pragma unroll
for (int i = 0; i < 4; i++) {
#pragma unroll
for (int j = 0; j < 4; j++) {
reg_array[i * 4 + j] += reg_1[4 + i] * reg_2[4 + j];
}
}
__syncthreads();
float *C_buffer = &buffer[threadIdx.y * 3072 + 1024]; // [64, 32]
#pragma unroll
for (int i = 0; i < 4; i++) {
#pragma unroll
for (int j = 0; j < 4; j++) {
C_buffer[(reg_1_idx * 4 + i) * 32 + reg_2_idx * 4 + j] = reg_array[i * 4 + j];
}
}
__syncthreads();
#pragma unroll
for (int i = 0; i < 16; i++) {
atomicAdd(&dense_C_pt[dim_stride * 32 + i * 128 + thread_idx], C_buffer[i * 128 + thread_idx]);
}
__syncthreads();
}
}
__global__ void reduce_sum_cuda_kernel(
float *sparse_A, // [batch_size, num_block, 32, 32]
int *indices, // [batch_size, num_block]
float *dense_C, // [batch_size, A_num_block, 32]
long batch_size,
long A_num_block,
long B_num_block,
long num_block
) {
long batch_idx = blockIdx.y;
long block_idx = blockIdx.x * blockDim.y + threadIdx.y;
long thread_idx = threadIdx.x;
long batch_idx__block_idx = batch_idx * num_block + block_idx;
long AB_block_idx = indices[batch_idx__block_idx];
float *sparse_A_pt = &sparse_A[batch_idx__block_idx * 1024];
float reg_array[16];
float value = 0;
#pragma unroll
for (int i = 0; i < 8; i++) {
reg_array[i] = sparse_A_pt[i * 32 + thread_idx];
}
#pragma unroll
for (int stride = 8; stride < 32; stride = stride + 8) {
#pragma unroll
for (int i = 0; i < 8; i++) {
reg_array[(stride + i) % 16] = sparse_A_pt[(stride + i) * 32 + thread_idx];
}
#pragma unroll
for (int i = 0; i < 8; i++) {
value = value + reg_array[(stride - 8 + i) % 16];
}
}
#pragma unroll
for (int i = 0; i < 8; i++) {
value = value + reg_array[8 + i];
}
float *dense_C_pt = &dense_C[(batch_idx * A_num_block + AB_block_idx / B_num_block) * 32];
atomicAdd(&dense_C_pt[thread_idx], value);
}
__global__ void scatter_cuda_kernel(
float *dense_A, // [batch_size, A_num_block, 32]
int *indices, // [batch_size, num_block]
float *sparse_C, // [batch_size, num_block, 32, 32]
long batch_size,
long A_num_block,
long B_num_block,
long num_block
) {
long batch_idx = blockIdx.y;
long block_idx = blockIdx.x * blockDim.y + threadIdx.y;
long thread_idx = threadIdx.x;
long batch_idx__block_idx = batch_idx * num_block + block_idx;
long AB_block_idx = indices[batch_idx__block_idx];
float *dense_A_pt = &dense_A[(batch_idx * A_num_block + AB_block_idx / B_num_block) * 32];
float *sparse_C_pt = &sparse_C[(batch_idx * num_block + block_idx) * 1024];
float value = dense_A_pt[thread_idx];
#pragma unroll
for (int i = 0; i < 32; i++) {
sparse_C_pt[i * 32 + thread_idx] = value;
}
}