submission 44151

- import torch
- import triton
- import triton.language as tl
- #!POPCORN leaderboard identity_py
+ # 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
   
- @triton.jit
- def kernel(inp, out, numel, BLOCK_SIZE: tl.constexpr):
-     pid = tl.program_id(0)
-     offs = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE
-     tl.arange(0, 4)
-     red = tl.load(inp + offs * 3, mask=(offs * 3 < numel))
-     green = tl.load(inp + offs * 3 + 1, mask=(offs * 3 + 1 < numel))
-     blue = tl.load(inp + offs * 3 + 2, mask=(offs * 3 + 2 < numel))
-     res = 0.2989 * red + 0.5870 * green + 0.1140 * blue
-     tl.store(out + offs, res, mask=(offs * 3 < numel))
+ import torch
+ from torch.utils.cpp_extension import load_inline
+ import time
+ import os
+ import sys
   
- # User kernel implementation.
- def custom_kernel(input: input_t) -> output_t:
-     data, output = input
-     numel = data.numel()
-     BLOCK_SIZE = 128
-     grid = (triton.cdiv(numel, 3 * BLOCK_SIZE),)
-     kernel[grid](data, output, numel, BLOCK_SIZE)
+ kernel_cpp = r"""#include <cuda_runtime.h>
+ #include <device_launch_parameters.h>
+ #include <stdio.h>
+ #include <torch/extension.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];
+  
+     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) {
+             reinterpret_cast<int4*>(load_data)[i] = reinterpret_cast<int4*>(input)[idx];
+         }
+     }
+  
+ #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 = 2;
+  
+     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");
+ }"""
+  
+ cuda_module = load_inline(name="cuda_module", cpp_sources="", cuda_sources=kernel_cpp, with_cuda=True, verbose=True)
+  
+ def custom_kernel(inp: input_t) -> output_t:
+     data, output = inp
+  
+     cuda_module.custom_kernel(data, output)
      return output
   No newline at end of file