submission 230430

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

# =============================================================================
# Inline CUDA: float4 ベクトルロード版 RGB to Grayscale (Phase 3: Inline PTX)
# =============================================================================
# 目標: 2.47ms (Triton) → 2.38ms
#
# Phase 3 変更点:
#   - Inline PTX で ld.global.v4.f32 を明示的に使用
#   - LDG.128 命令を完全に強制
#   - コンパイラの最適化判断を完全にバイパス
#
# PTX 命令:
#   ld.global.v4.f32 {%0, %1, %2, %3}, [%4];
#   - 128-bit (16 bytes) 一括ロード
#   - 4 つの float を同時に取得

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)
    );
}

// RGB to Grayscale カーネル（Inline 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 計算
    float4 gray;
    gray.x = d0_x * 0.2989f + d0_y * 0.5870f + d0_z * 0.1140f;  // Pixel 0
    gray.y = d0_w * 0.2989f + d1_x * 0.5870f + d1_y * 0.1140f;  // Pixel 1
    gray.z = d1_z * 0.2989f + d1_w * 0.5870f + d2_x * 0.1140f;  // Pixel 2
    gray.w = d2_y * 0.2989f + d2_z * 0.5870f + d2_w * 0.1140f;  // Pixel 3

    // float4 で出力
    *reinterpret_cast<float4*>(output + pixel_base) = gray;
}

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_float4_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