submission 44180

# WARNING: This file is generated by prepare.py
# Do not edit this file directly. Instead, edit:
# - pytorch_wrapper.py for the Python wrapper code
# - cuda_impl.cu for the CUDA kernel implementation
# Then run prepare.py to regenerate this file.

from task import input_t, output_t

import torch
from torch.utils.cpp_extension import load_inline
import time
import os
import sys

kernel_cpp = r"""#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <stdio.h>
#include <torch/extension.h>
#include <cooperative_groups.h>
#include <cooperative_groups/memcpy_async.h>

template<int NUM_WARPS>
__global__ void grayscale_kernel(float* input, float* output, int total_pixels) {
    __shared__ float shr[32*3*NUM_WARPS];

    int block_start = blockIdx.x * blockDim.x * 3;
    int lane = threadIdx.x % 32;
    int warp = threadIdx.x / 32;
    int block_offset = 32 * 3 * warp + lane;
    float load_data[3];
#pragma unroll
    for (int i=0; i<3; ++i) {
        int idx = block_offset + block_start + 32 * i;
        if (idx < total_pixels * 3) {
            load_data[i] = input[idx];
        }
    }

#pragma unroll
    for (int i=0; i<3; ++i) {
        shr[block_offset + 32 * i] = load_data[i];
    }
    // __syncwarp();

    // TODO load the data
    int shr_start = 3 * threadIdx.x;
    if (blockIdx.x * blockDim.x + threadIdx.x < total_pixels) {
        float r = shr[shr_start];
        float g = shr[shr_start + 1];
        float b = shr[shr_start + 2];

        output[blockIdx.x * blockDim.x + threadIdx.x] = 0.299f * r + 0.587f * g + 0.114f * b;
    }
}


template<int NUM_WARPS>
__global__ void grayscale_kernel_VECTORIZED(float* input, float* output, int total_pixels) {
    __shared__ float shr[32*3*NUM_WARPS*4];

    // auto block = cooperative_groups::this_thread_block();

    int block_start = blockIdx.x * blockDim.x * 3;
    int lane = threadIdx.x % 32;
    int warp = threadIdx.x / 32;
    int block_offset = 32 * 3 * warp + lane;
    float load_data[12];
#pragma unroll
    for (int i=0; i<3; ++i) {
        int idx = block_offset + block_start + 32 * i;
        if (idx * 4 < total_pixels * 3) {
            // cooperative_groups::memcpy_async(block, shr, input + block_start * 4, 32 * 3 * NUM_WARPS * 4 * sizeof(float));
            reinterpret_cast<int4*>(load_data)[i] = reinterpret_cast<int4*>(input)[idx];
        }
    }

    // cooperative_groups::wait(block);

#pragma unroll
    for (int i=0; i<3; ++i) {
        reinterpret_cast<int4*>(shr)[block_offset + 32 * i] = reinterpret_cast<int4*>(load_data)[i];
    }
    // __syncwarp();

    int shr_start = 3 * threadIdx.x * 4;
    if ((blockIdx.x * blockDim.x + threadIdx.x) * 4 < total_pixels) {
        float result[4];

        for (int j = 0; j<4; ++j) {
            float r = shr[shr_start + j*3];
            float g = shr[shr_start + j*3 + 1];
            float b = shr[shr_start + j*3 + 2];

            result[j] = 0.299f * r + 0.587f * g + 0.114f * b;
        }
        reinterpret_cast<int4*>(output)[blockIdx.x * blockDim.x + threadIdx.x] = reinterpret_cast<int4*>(result)[0];
    }
}

void custom_kernel(torch::Tensor input, torch::Tensor output) {
    int total_pixels = input.size(0) * input.size(1);

    constexpr int NUM_WARPS = 4;

    if (total_pixels % 4 == 0) {
        int threads_per_block = 32 * NUM_WARPS;
        int blocks = (total_pixels + threads_per_block - 1) / (threads_per_block * 4);

        grayscale_kernel_VECTORIZED<NUM_WARPS><<<blocks, threads_per_block>>>((float*) input.data_ptr(), (float*) output.data_ptr(), total_pixels);
    } else {
        int threads_per_block = 32 * NUM_WARPS;
        int blocks = (total_pixels + threads_per_block - 1) / (threads_per_block);

        grayscale_kernel<NUM_WARPS><<<blocks, threads_per_block>>>((float*) input.data_ptr(), (float*) output.data_ptr(), total_pixels);
    }
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("custom_kernel", &custom_kernel, "custom kernel");
}"""

# os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0"
 
cuda_module = load_inline(name="cuda_module", cpp_sources="", cuda_sources=kernel_cpp, with_cuda=True, verbose=True, extra_cuda_cflags=["-O3", "-Xptxas", "-O3"])

def custom_kernel(inp: input_t) -> output_t:
    data, output = inp

    cuda_module.custom_kernel(data, output)
    return output