submission 230431

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

# =============================================================================
# Inline CUDA: 入出力両方 Inline PTX 版 RGB to Grayscale
# =============================================================================
# 目標: 2.39ms → 2.38ms (1位: 2389.197μs = 2.389ms)
#
# 変更点:
#   - 入力: ld.global.v4.f32 (継続)
#   - 出力: st.global.v4.f32 (新規追加)
#   - 完全に PTX レベルでメモリアクセスを制御
#
# 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 カーネル（入出力両方 PTX 版）
__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();

    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_full_ptx',
            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