submission 230556

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

# =============================================================================
# Inline CUDA: B200最適化版 RGB to Grayscale
# =============================================================================
# 目標: 598.336μs以下（B200 1位）
#
# v2_11ベースの最適化:
#   - threads = 128 に変更（より多くのブロックでSM占有率向上）
#   - 入力: ld.global.v4.f32 (PTX 128-bit ロード)
#   - 出力: st.global.v4.f32 (PTX 128-bit ストア)
#
# 設計根拠:
#   - B200はBlackwell世代でSM数が多い
#   - より多くのブロックを生成することでSM占有率を向上
#   - A100での最適値256に対し、B200では128が最適な可能性
#
# 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 カーネル（B200最適化版）
__global__ 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();

    // B200最適化: threads = 128（より多くのブロック生成）
    int threads = 128;
    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_optimized',
            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