conv1d.cu

#include <cuda_runtime.h> 
#include <iostream> 

#define KERNEL_SIZE 2048

#ifdef ALIGNED 
__device__ __align__(64) float cmem_kernel[KERNEL_SIZE];
#else
__device__ float cmem_kernel[KERNEL_SIZE]; 
#endif 

#ifndef THREADS_PER_BLOCK
#define THREADS_PER_BLOCK 256
#endif

__global__ void naive_conv1d_kernel(const float* input, const float* kernel, float* output, int input_size, int kernel_size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int output_size = input_size - kernel_size + 1;

    if (idx < output_size) {
        float sum = 0.0f;
        for (int j = 0; j < kernel_size; ++j) {
            sum += input[idx + j] * kernel[j];
        }
        output[idx] = sum;
    }
}

void naive_conv1d(const float* input, const float* kernel, float* output, int input_size, int kernel_size){    
    int output_size = input_size - kernel_size + 1;

    int blocksPerGrid = (output_size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;

    naive_conv1d_kernel<<<blocksPerGrid, THREADS_PER_BLOCK>>>(input, kernel, output, input_size, kernel_size);
    cudaDeviceSynchronize();
}

// ############################################################################################

__global__ void cmem_conv1d_kernel(const float* input, const float* kernel, float* output, int input_size, int kernel_size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int output_size = input_size - kernel_size + 1;

    if (idx < output_size) {
        float sum = 0.0f;
        for (int j = 0; j < kernel_size; ++j) {
            sum += input[idx + j] * cmem_kernel[j];
        }
        output[idx] = sum;
    }
}

void cmem_conv1d(const float* input, const float* kernel, float* output, int input_size, int kernel_size){
    int output_size = input_size - kernel_size + 1;

    int blocksPerGrid = (output_size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
    cudaMemcpyToSymbol(cmem_kernel, kernel, kernel_size * sizeof(float));

    cmem_conv1d_kernel<<<blocksPerGrid, THREADS_PER_BLOCK>>>(input, kernel, output, input_size, kernel_size);
    cudaDeviceSynchronize();
}

// ############################################################################################

__global__ void smem_conv1d_kernel(const float* input, const float* kernel, float*  output,
                                      int input_size, int kernel_size) {
    
    // idx of the first output element that the block is responsible for
    const size_t block_start =  blockDim.x * blockIdx.x ; 
    // idx of the output element that this thread is responsible for
    const size_t tid = threadIdx.x + block_start ; 
    
    extern __shared__ float smem_input[]; // shared memory 

    const size_t smem_len = blockDim.x + kernel_size - 1; // block region + halo region
    
    // all threads collaborate to load data from gmem to smem 
    for (size_t i = threadIdx.x; i < smem_len ; i += blockDim.x) {
        size_t global_idx = block_start + i;
        if (global_idx < input_size) 
            smem_input[i] = input[global_idx];
    }
    
    // wait for all threads to finish the transfer
    __syncthreads(); 

    if (tid > input_size - kernel_size)
        return; 
    
    // compute the convolution
    float sum = 0 ; 
    for(size_t i = 0 ; i < kernel_size ; i++){
        sum += smem_input[threadIdx.x +i] * cmem_kernel[i]; 
    }
    output[tid] = sum ; 
    
}


void smem_conv1d(const float* input, const float* kernel, float* output, int input_size, int kernel_size) {
    int output_size = input_size - kernel_size + 1;

    int blocksPerGrid = (output_size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;

    size_t smem_size = (THREADS_PER_BLOCK + kernel_size) * sizeof(float);
    
    cudaMemcpyToSymbol(cmem_kernel, kernel, kernel_size * sizeof(float));

    smem_conv1d_kernel<<<blocksPerGrid, THREADS_PER_BLOCK, smem_size>>>(input, kernel, output, input_size, kernel_size);
    cudaDeviceSynchronize();
}


// ############################################################################################

// only define the constants for the register block version 
#define OUTPUTS_PER_THREAD 8

__global__ void reg_block_conv1d_kernel(const float* input, const float* kernel, float*  output,
                                      int input_size, int kernel_size)  {
    
    const int block_output_size = blockDim.x * OUTPUTS_PER_THREAD;
    int output_size = input_size - kernel_size + 1;
    
    const size_t block_output_idx_start = (size_t)blockIdx.x * blockDim.x * OUTPUTS_PER_THREAD;
    const size_t thread_output_idx_start = block_output_idx_start + threadIdx.x * OUTPUTS_PER_THREAD;

    extern __shared__ float smem_input[];
    const int smem_size = block_output_size + kernel_size - 1;

    for (int i = threadIdx.x; i < smem_size; i += blockDim.x) {
        size_t global_idx = block_output_idx_start + i;
        if (global_idx < input_size) {
            smem_input[i] = input[global_idx];
        }
    }

    __syncthreads(); 

    if (thread_output_idx_start >= output_size) {
        return;
    }

    float sum[OUTPUTS_PER_THREAD] = {};
    float reg_input[OUTPUTS_PER_THREAD];
    

    #pragma unroll
    for (int n = 0; n < OUTPUTS_PER_THREAD; ++n) 
        if (threadIdx.x * OUTPUTS_PER_THREAD + n < smem_size )
            reg_input[n] = smem_input[threadIdx.x * OUTPUTS_PER_THREAD + n] ; 
 
    for (int k = 0; k < kernel_size; ++k) {
        const float k_val = cmem_kernel[k];

        #pragma unroll 
        for (int n = 0; n < OUTPUTS_PER_THREAD; ++n){
            sum[n] += k_val * reg_input[n]; 
        }

        #pragma unroll
        for (int n = 0; n < OUTPUTS_PER_THREAD-1; ++n){
            reg_input[n] = reg_input[n+1]; 
        }
        
        reg_input[OUTPUTS_PER_THREAD-1] = smem_input[threadIdx.x * OUTPUTS_PER_THREAD + k + OUTPUTS_PER_THREAD]; 
    }

    #pragma unroll
    for (int n = 0; n < OUTPUTS_PER_THREAD; ++n) {
        if (thread_output_idx_start + n < output_size) {
            output[thread_output_idx_start + n] = sum[n]; 
        }
    }
}

void reg_block_conv1d(const float* input, const float* kernel, float* output, int input_size, int kernel_size) {
    int output_size = input_size - kernel_size + 1;
    constexpr int block_output_size = THREADS_PER_BLOCK * OUTPUTS_PER_THREAD;

    int blocksPerGrid = (output_size + block_output_size - 1) / block_output_size;

    int shared_mem_size = block_output_size + kernel_size +  OUTPUTS_PER_THREAD ; //for padding

    cudaMemcpyToSymbol(cmem_kernel, kernel, kernel_size * sizeof(float));

    reg_block_conv1d_kernel<<<
        blocksPerGrid,
        THREADS_PER_BLOCK,
        shared_mem_size * sizeof(float)
    >>>(input, kernel, output, input_size, kernel_size);

    
    cudaDeviceSynchronize();
}
  // Max supported kernel size (adjust if needed)

__global__ void vectorized_conv1d_kernel(const float*  input,
                                         const float*  kernel,
                                         float* output,
                                         int input_size, int kernel_size,
                                         const int smem_size) {
    const int block_output_size = blockDim.x * OUTPUTS_PER_THREAD;
    const int output_size = input_size - kernel_size + 1;

    const size_t block_output_idx_start = (size_t)blockIdx.x * block_output_size;
    const size_t thread_output_idx_start = block_output_idx_start + threadIdx.x * OUTPUTS_PER_THREAD;

    extern __shared__ float smem_input[];  // shared memory buffer

    const size_t smem_start_idx = block_output_idx_start;
    const float4* input_vec4 = reinterpret_cast<const float4*>(input);
    float4* smem_vec4 = reinterpret_cast<float4*>(smem_input);

    const int smem_size_vec4 = smem_size / 4;

    // Load from global memory into shared memory using float4
    for (int i = threadIdx.x; i < smem_size_vec4; i += blockDim.x) {
        size_t global_idx = smem_start_idx + i * 4;

        if (global_idx + 3 < input_size) {
            smem_vec4[i] = input_vec4[global_idx / 4];
        } else {
            float buffer[4] = {};
            for (int k = 0; k < 4; ++k) {
                if (global_idx + k < input_size) {
                    buffer[k] = input[global_idx + k];
                }
            }
            smem_vec4[i] = make_float4(buffer[0], buffer[1], buffer[2], buffer[3]);
        }
    }

    __syncthreads();

    // Early return if this thread won't produce valid outputs
    if (thread_output_idx_start >= output_size) return;

    // Load first float4 (register input init)
    float4 vec = smem_vec4[threadIdx.x];
    float reg_input[OUTPUTS_PER_THREAD] = {vec.x, vec.y, vec.z, vec.w};
    float reg_input_next[OUTPUTS_PER_THREAD] = {};
    float sum[OUTPUTS_PER_THREAD] = {};

    const int n_outer = (kernel_size + OUTPUTS_PER_THREAD - 1) / OUTPUTS_PER_THREAD;

    for (int i = 0; i < n_outer; ++i) {
        // Guard access to avoid out-of-bounds in smem_vec4
        int vec_idx = threadIdx.x + i + 1;
        if (vec_idx < smem_size_vec4) {
            vec = smem_vec4[vec_idx];
        } else {
            vec = make_float4(0.f, 0.f, 0.f, 0.f);
        }

        reg_input_next[0] = vec.x;
        reg_input_next[1] = vec.y;
        reg_input_next[2] = vec.z;
        reg_input_next[3] = vec.w;

        #pragma unroll
        for (int k = 0; k < OUTPUTS_PER_THREAD; ++k) {
            int j = i * OUTPUTS_PER_THREAD + k;
            if (j < kernel_size) {
                float k_val = cmem_kernel[j];
                #pragma unroll
                for (int n = 0; n < OUTPUTS_PER_THREAD; ++n) {
                    sum[n] += k_val * reg_input[n];
                }
                #pragma unroll
                for (int n = 0; n < OUTPUTS_PER_THREAD - 1; ++n) {
                    reg_input[n] = reg_input[n + 1];
                }
                reg_input[OUTPUTS_PER_THREAD - 1] = reg_input_next[k];
            }
        }
    }

    // Write results
    #pragma unroll
    for (int n = 0; n < OUTPUTS_PER_THREAD; ++n) {
        int out_idx = thread_output_idx_start + n;
        if (out_idx < output_size) {
            output[out_idx] = sum[n];
        }
    }
}


void vectorized_conv1d(const float*  input, const float* kernel, float* output, int input_size, int kernel_size) {
    int output_size = input_size - kernel_size + 1;

    constexpr int block_output_size = THREADS_PER_BLOCK * OUTPUTS_PER_THREAD;

    int blocksPerGrid = (output_size + block_output_size - 1) / block_output_size;

    int shared_mem_size = block_output_size + kernel_size - 1 ; 
    // ensure shared memory is a multiple of 4 for float4 alignment 
    shared_mem_size = 4 * ((shared_mem_size + 3) / 4);  

    cudaMemcpyToSymbol(cmem_kernel, kernel, kernel_size * sizeof(float));

    vectorized_conv1d_kernel<<<
        blocksPerGrid,
        THREADS_PER_BLOCK,
        shared_mem_size * sizeof(float)
    >>>(input, kernel, output, input_size, kernel_size, shared_mem_size);

    
    cudaDeviceSynchronize();
}




// ############################################################################################ 



__global__ void micro_opt_conv1d_kernel(const float* input, const float* kernel, float*  output,
                                      int input_size, int kernel_size)  {
    
    const int block_output_size = blockDim.x * OUTPUTS_PER_THREAD;
    int output_size = input_size - kernel_size + 1;
    
    const size_t block_output_idx_start = (size_t)blockIdx.x * block_output_size;

    const size_t thread_output_idx_start = block_output_idx_start + threadIdx.x * OUTPUTS_PER_THREAD;
    
    extern __shared__ float smem_input[];
    const int smem_size = block_output_size + kernel_size - 1;

    for (int i = threadIdx.x; i < smem_size; i += blockDim.x) {
        size_t global_idx = block_output_idx_start + i;
        if (global_idx < input_size) {
            smem_input[i] = __ldg(input+global_idx);
        }
    }

    __syncthreads(); 

    if (thread_output_idx_start >= output_size) {
        return;
    }

    float sum[OUTPUTS_PER_THREAD];
    float reg_input[OUTPUTS_PER_THREAD];
    
    #pragma unroll
    for (int n = 0; n < OUTPUTS_PER_THREAD; ++n) {
        sum[n] = 0.0f;
    }

    #pragma unroll
    for (int n = 0; n < OUTPUTS_PER_THREAD; ++n) 
        if (threadIdx.x * OUTPUTS_PER_THREAD + n < smem_size )
            reg_input[n] = smem_input[threadIdx.x * OUTPUTS_PER_THREAD + n] ; 
 
    for (int k = 0; k < kernel_size; ++k) {
        const float k_val = cmem_kernel[k];

        #pragma unroll 
        for (int n = 0; n < OUTPUTS_PER_THREAD; ++n){
            sum[n] += k_val * reg_input[n]; 
        }

        #pragma unroll
        for (int n = 0; n < OUTPUTS_PER_THREAD-1; ++n){
            reg_input[n] = reg_input[n+1]; 
        }
        
        reg_input[OUTPUTS_PER_THREAD-1] = smem_input[threadIdx.x * OUTPUTS_PER_THREAD + k + OUTPUTS_PER_THREAD]; 
    }

    #pragma unroll
    for (int n = 0; n < OUTPUTS_PER_THREAD; ++n) {
        if (thread_output_idx_start + n < output_size) {

            output[thread_output_idx_start + n] = sum[n];
            //asm volatile("st.global.cs.f32 [%0], %1;" :: "l"(output + thread_output_idx_start + n), "f"(sum[n])) ; 
        }
    }
}

void micro_opt_conv1d(const float* input, const float* kernel, float* output, int input_size, int kernel_size) {
    
    int output_size = input_size - kernel_size + 1;
    constexpr int block_output_size = THREADS_PER_BLOCK * OUTPUTS_PER_THREAD;

    int blocksPerGrid = (output_size + block_output_size - 1) / block_output_size;

    int shared_mem_size = block_output_size + kernel_size +  OUTPUTS_PER_THREAD ; //for padding

    cudaMemcpyToSymbol(cmem_kernel, kernel, kernel_size * sizeof(float));

    micro_opt_conv1d_kernel<<<
        blocksPerGrid,
        THREADS_PER_BLOCK,
        shared_mem_size * sizeof(float)
    >>>(input, kernel, output, input_size, kernel_size);

    cudaDeviceSynchronize();
}