submission 78259

from task import input_t, output_t
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline


def custom_kernel(data: input_t) -> output_t:
    cuda_source = """
    #include <torch/types.h>
    #include <cuda.h>
    #include <cuda_runtime.h>
    #include <cmath>

    #define BLOCK_SIZE 32
    #define THREADS_PER_BLOCK 1024

    __global__ void convolution_2d_kernel(const float* matrix, const float* conv_kernel, float* result,
                                        int batch_size, int kernel_size, int in_channels, int out_channels, 
                                        int in_grid, int out_grid, int sm_writes_per_block, int sm_writes_per_thread) {
        int z = blockIdx.z;
        int row = blockIdx.y * blockDim.y + threadIdx.y;
        int col = blockIdx.x * blockDim.x + threadIdx.x;

        int channel_idx = z % out_channels;
        int batch_idx = z / out_channels;
        int output_idx = batch_idx * out_channels * out_grid * out_grid + 
                            channel_idx * out_grid * out_grid + row * out_grid + col;

        extern __shared__ float conv_sram[];
        if (z >= out_channels * batch_size) return;
        
        int sram_index = threadIdx.y * blockDim.x + threadIdx.x;
        float element_value = 0.0f;

        for (int c = 0; c < in_channels; c++) {
            for (int i = 0; i < sm_writes_per_thread; i++) {
                int element_index = i * THREADS_PER_BLOCK + sram_index;
                if ( element_index < sm_writes_per_block) {
                    int conv_index = channel_idx * in_channels * kernel_size * kernel_size + c * kernel_size * kernel_size + element_index;
                    conv_sram[element_index] = conv_kernel[conv_index];
                }
            }


            __syncthreads();
            if (row < out_grid && col < out_grid) {
                for (int y = 0; y < kernel_size; y++) {
                    for (int x = 0; x < kernel_size; x++) {
                        int matrix_index = batch_idx * in_channels * in_grid * in_grid + 
                            c * in_grid * in_grid + (row + y) * in_grid + (col + x);
                        element_value += conv_sram[y * kernel_size + x] * matrix[matrix_index];
                    }
                }
            }
            __syncthreads(); 
        }

        if (row < out_grid & col < out_grid) {
            result[output_idx] = element_value;
        }
    }

    torch::Tensor convolution_2d(torch::Tensor matrix, torch::Tensor conv_kernel) {
        const auto batch = matrix.size(0);
        const auto in_channels = matrix.size(1);
        const auto in_grid = matrix.size(2);
        const auto kernel_size = conv_kernel.size(2);
        const auto out_grid = in_grid + 1 - kernel_size;
        const auto out_channels = conv_kernel.size(0);

        const int sm_writes_per_block = kernel_size * kernel_size;
        const int sm_writes_per_thread = (sm_writes_per_block + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
        auto result = torch::zeros({batch, out_channels, out_grid, out_grid}, matrix.options());

        dim3 threads_per_block(BLOCK_SIZE, BLOCK_SIZE, 1);
        dim3 number_of_blocks((out_grid + threads_per_block.x - 1) / threads_per_block.x,
                            (out_grid + threads_per_block.y - 1) / threads_per_block.y,
                            out_channels * batch); //flattened output channels and batch
        int shared_memory_size = sizeof(float) * kernel_size * kernel_size;

        convolution_2d_kernel<<<number_of_blocks, threads_per_block, shared_memory_size>>>(
            matrix.data_ptr<float>(), conv_kernel.data_ptr<float>(), result.data_ptr<float>(),
            batch, kernel_size, in_channels, out_channels, in_grid, out_grid, sm_writes_per_block,
            sm_writes_per_thread);

        return result;
    }

    """
    import os
    os.environ['TORCH_CUDA_ARCH_LIST'] = '8.9'

    cpp_source = "torch::Tensor convolution_2d(torch::Tensor matrix, torch::Tensor conv_kernel);"
    convolution_2d_extension = load_inline(
        name='convolution_2d_extension',
        cpp_sources=cpp_source,
        cuda_sources=cuda_source,
        functions=['convolution_2d'],
        with_cuda=True,
        extra_cuda_cflags=["-O3"],
        build_directory=None,
        # extra_cuda_cflags=['--expt-relaxed-constexpr']
    )  
    
    input_tensor, kernel, output = data
    print(input_tensor.shape)
    print(kernel.shape)
    output[...] = convolution_2d_extension.convolution_2d(input_tensor, kernel)
    return output