submission 230715

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

# =============================================================================
# Inline CUDA: B200最適化版 RGB to Grayscale (__launch_bounds__)
# =============================================================================
# 目標: 598.336μs以下（B200 1位）
#
# 最適化: __launch_bounds__(256, 4) でOccupancy向上
#   - maxThreadsPerBlock = 256: ブロックあたり最大256スレッド
#   - minBlocksPerMultiprocessor = 4: SM当たり最低4ブロックを保証
#   → コンパイラがレジスタ使用量を制限してこれを達成
#
# v2_11からの変更点:
#   - カーネル宣言に __launch_bounds__(256, 4) を追加
#   - ロジックは完全に同一
#
# PTX 命令:
#   ld.global.v4.f32 {%0, %1, %2, %3}, [%4];  - 128-bit ロード
#   st.global.v4.f32 [%0], {%1, %2, %3, %4};  - 128-bit ストア

cuda_src = """
#include <cuda_runtime.h>

// Inline PTX による float4 ロード
__device__ __forceinline__ void load_float4_ptx(
    const float* ptr,
    float& r0, float& r1, float& r2, float& r3
) {
    asm volatile(
        "ld.global.v4.f32 {%0, %1, %2, %3}, [%4];"
        : "=f"(r0), "=f"(r1), "=f"(r2), "=f"(r3)
        : "l"(ptr)
    );
}

// Inline PTX による float4 ストア
__device__ __forceinline__ void store_float4_ptx(
    float* ptr,
    float v0, float v1, float v2, float v3
) {
    asm volatile(
        "st.global.v4.f32 [%0], {%1, %2, %3, %4};"
        :
        : "l"(ptr), "f"(v0), "f"(v1), "f"(v2), "f"(v3)
        : "memory"
    );
}

// RGB to Grayscale カーネル（__launch_bounds__版）
// __launch_bounds__(256, 4): SM当たり4ブロック（1024スレッド）を保証
__global__ __launch_bounds__(256, 4) void grayscale_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int n_pixels
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int pixel_base = tid * 4;

    if (pixel_base >= n_pixels) return;

    int base = pixel_base * 3;
    const float* ptr = input + base;

    // PTX で 128-bit ロード × 3 回
    float d0_x, d0_y, d0_z, d0_w;  // [R0, G0, B0, R1]
    float d1_x, d1_y, d1_z, d1_w;  // [G1, B1, R2, G2]
    float d2_x, d2_y, d2_z, d2_w;  // [B2, R3, G3, B3]

    load_float4_ptx(ptr,      d0_x, d0_y, d0_z, d0_w);
    load_float4_ptx(ptr + 4,  d1_x, d1_y, d1_z, d1_w);
    load_float4_ptx(ptr + 8,  d2_x, d2_y, d2_z, d2_w);

    // Grayscale 計算
    float gray_x = d0_x * 0.2989f + d0_y * 0.5870f + d0_z * 0.1140f;  // Pixel 0
    float gray_y = d0_w * 0.2989f + d1_x * 0.5870f + d1_y * 0.1140f;  // Pixel 1
    float gray_z = d1_z * 0.2989f + d1_w * 0.5870f + d2_x * 0.1140f;  // Pixel 2
    float gray_w = d2_y * 0.2989f + d2_z * 0.5870f + d2_w * 0.1140f;  // Pixel 3

    // PTX で 128-bit ストア
    store_float4_ptx(output + pixel_base, gray_x, gray_y, gray_z, gray_w);
}

torch::Tensor rgb_to_grayscale(torch::Tensor input, torch::Tensor output) {
    int n_pixels = output.numel();

    int threads = 256;
    int pixels_per_block = threads * 4;
    int blocks = (n_pixels + pixels_per_block - 1) / pixels_per_block;

    grayscale_kernel<<<blocks, threads>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        n_pixels
    );

    return output;
}
"""

cpp_src = """
#include <torch/extension.h>

torch::Tensor rgb_to_grayscale(torch::Tensor input, torch::Tensor output);
"""

_module = None

def _get_module():
    global _module
    if _module is None:
        _module = load_inline(
            name='grayscale_cuda_b200_launch_bounds',
            cuda_sources=[cuda_src],
            cpp_sources=[cpp_src],
            functions=['rgb_to_grayscale'],
            extra_cuda_cflags=['-O3', '--use_fast_math'],
            verbose=False
        )
    return _module


def custom_kernel(data: input_t) -> output_t:
    input_tensor, output_tensor = data

    module = _get_module()
    module.rgb_to_grayscale(input_tensor, output_tensor)

    return output_tensor