submission 233446

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

# =============================================================================
# Inline CUDA: B200 TMA + LDS.128 + 1024ピクセル/タイル版
# =============================================================================
# 目標: 598.336μs以下（B200 1位）
#
# v2_14からの最適化 (EXP-01):
#   タイルサイズを512→1024ピクセルに拡大
#   - TMA発行回数が半減（524,288 → 262,144）
#   - ブロック起動オーバーヘッドが半減
#   - スレッド数を128→256に拡大
#
# ベース設計 (v2_14から継承):
#   - TMA (cp.async.bulk) による非同期バルクコピー
#   - blockIdx.x による通常グリッド起動（cudaMalloc排除）
#   - ld.shared.v4.f32 (LDS.128) で共有メモリから一括ロード
#   - st.global.v4.f32 (STG.128) による出力ベクトル化
#   - 1スレッド4ピクセル処理でスループット最大化
#
# 共有メモリ使用量: 12KB/ブロック (B200: 228KB/SM, 余裕あり)
# =============================================================================

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

// =============================================================================
// 定数定義 (EXP-01: 1024ピクセル/タイル)
// =============================================================================
constexpr int PIXELS_PER_TILE   = 1024;      // 1024ピクセル/タイル (512→1024)
constexpr int BYTES_PER_TILE    = 12288;     // 1024 × 3 × 4 = 12288バイト (16の倍数)
constexpr int FLOATS_PER_TILE   = 3072;      // 1024 × 3 = 3072 floats

constexpr int THREADS_PER_BLOCK = 256;       // 256スレッド/ブロック (128→256)
constexpr int PIXELS_PER_THREAD = 4;         // 4ピクセル/スレッド維持 (256 × 4 = 1024)

// =============================================================================
// PTXヘルパー関数
// =============================================================================

// mbarrier初期化
__device__ __forceinline__ void mbarrier_init(uint64_t* mbar, int arrival_count) {
    asm volatile(
        "mbarrier.init.shared.b64 [%0], %1;"
        :: "r"((uint32_t)__cvta_generic_to_shared(mbar)), "r"(arrival_count)
    );
}

// mbarrier: 期待するトランザクションバイト数を設定 + 到着
__device__ __forceinline__ void mbarrier_arrive_expect_tx(uint64_t* mbar, int tx_bytes) {
    asm volatile(
        "mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;"
        :: "r"((uint32_t)__cvta_generic_to_shared(mbar)), "r"(tx_bytes)
    );
}

// mbarrier: フェーズビットで待機
__device__ __forceinline__ void mbarrier_wait(uint64_t* mbar, int phase) {
    int done = 0;
    while (!done) {
        asm volatile(
            "{"
            ".reg .pred p;"
            "mbarrier.try_wait.parity.shared.b64 p, [%1], %2;"
            "selp.b32 %0, 1, 0, p;"
            "}"
            : "=r"(done)
            : "r"((uint32_t)__cvta_generic_to_shared(mbar)), "r"(phase)
        );
    }
}

// cp.async.bulk: グローバル→共有メモリのバルクコピー
__device__ __forceinline__ void cp_async_bulk(
    void* smem_ptr,
    const void* gmem_ptr,
    int bytes,
    uint64_t* mbar
) {
    asm volatile(
        "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];"
        :: "r"((uint32_t)__cvta_generic_to_shared(smem_ptr)),
           "l"(gmem_ptr),
           "r"(bytes),
           "r"((uint32_t)__cvta_generic_to_shared(mbar))
    );
}

// 共有メモリからのベクトルロード (LDS.128)
__device__ __forceinline__ void load_shared_float4(
    uint32_t smem_addr,
    float& v0, float& v1, float& v2, float& v3
) {
    asm volatile(
        "ld.shared.v4.f32 {%0, %1, %2, %3}, [%4];"
        : "=f"(v0), "=f"(v1), "=f"(v2), "=f"(v3)
        : "r"(smem_addr)
    );
}

// float4ベクトルストア（STG.128）
__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"
    );
}

// =============================================================================
// TMA + LDS.128 + 1024ピクセル/タイル カーネル
// =============================================================================
// __launch_bounds__(256, 4): SM当たり4ブロック（共有メモリ12KB×4=48KB < 228KB）
__global__ __launch_bounds__(256, 4) void grayscale_1024tile_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    const int n_pixels
) {
    // 共有メモリ: 12KB (3072 floats) + mbarrier (8B)
    __shared__ __align__(128) float smem[FLOATS_PER_TILE];
    __shared__ __align__(8) uint64_t mbar;

    const int tid = threadIdx.x;
    const int tile_idx = blockIdx.x;
    const int pixel_base = tile_idx * PIXELS_PER_TILE + tid * PIXELS_PER_THREAD;

    // 境界チェック（268,435,456 / 1024 = 262,144 で割り切れるが、安全のため維持）
    if (pixel_base >= n_pixels) return;

    // ========================================
    // Step 1: mbarrier初期化（Thread 0のみ）
    // ========================================
    if (tid == 0) {
        mbarrier_init(&mbar, 1);
    }
    __syncthreads();  // mbarrier初期化の可視性を保証

    // ========================================
    // Step 2: TMAロード発行（Thread 0のみ）
    // ========================================
    if (tid == 0) {
        const float* src = input + tile_idx * PIXELS_PER_TILE * 3;
        mbarrier_arrive_expect_tx(&mbar, BYTES_PER_TILE);
        cp_async_bulk(smem, src, BYTES_PER_TILE, &mbar);
    }

    // ========================================
    // Step 3: TMA完了待機
    // ========================================
    mbarrier_wait(&mbar, 0);

    // ========================================
    // Step 4: LDS.128で共有メモリから一括ロード
    // ========================================
    // tid * 12 floats = tid * 48 bytes (常に16の倍数)
    uint32_t smem_addr = (uint32_t)__cvta_generic_to_shared(&smem[tid * 12]);

    // s0 = [R0, G0, B0, R1]
    // s1 = [G1, B1, R2, G2]
    // s2 = [B2, R3, G3, B3]
    float s0_x, s0_y, s0_z, s0_w;
    float s1_x, s1_y, s1_z, s1_w;
    float s2_x, s2_y, s2_z, s2_w;

    load_shared_float4(smem_addr,      s0_x, s0_y, s0_z, s0_w);
    load_shared_float4(smem_addr + 16, s1_x, s1_y, s1_z, s1_w);  // +16バイト = +4floats
    load_shared_float4(smem_addr + 32, s2_x, s2_y, s2_z, s2_w);  // +32バイト = +8floats

    // ========================================
    // Step 5: Grayscale計算 (fmafチェーン)
    // ========================================
    // マッピング:
    // Pixel 0: R=s0_x, G=s0_y, B=s0_z
    // Pixel 1: R=s0_w, G=s1_x, B=s1_y
    // Pixel 2: R=s1_z, G=s1_w, B=s2_x
    // Pixel 3: R=s2_y, G=s2_z, B=s2_w
    float gray0 = fmaf(s0_x, 0.2989f, fmaf(s0_y, 0.5870f, s0_z * 0.1140f));
    float gray1 = fmaf(s0_w, 0.2989f, fmaf(s1_x, 0.5870f, s1_y * 0.1140f));
    float gray2 = fmaf(s1_z, 0.2989f, fmaf(s1_w, 0.5870f, s2_x * 0.1140f));
    float gray3 = fmaf(s2_y, 0.2989f, fmaf(s2_z, 0.5870f, s2_w * 0.1140f));

    // ========================================
    // Step 6: STG.128でグローバルメモリに出力
    // ========================================
    store_float4_ptx(output + pixel_base, gray0, gray1, gray2, gray3);
}

// =============================================================================
// ホストインターフェース
// =============================================================================
torch::Tensor rgb_to_grayscale(torch::Tensor input, torch::Tensor output) {
    int n_pixels = output.numel();

    // ブロック数計算: 268,435,456 / 1024 = 262,144 ブロック
    int blocks = (n_pixels + PIXELS_PER_TILE - 1) / PIXELS_PER_TILE;

    grayscale_1024tile_kernel<<<blocks, THREADS_PER_BLOCK>>>(
        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_1024tile',
            cuda_sources=[cuda_src],
            cpp_sources=[cpp_src],
            functions=['rgb_to_grayscale'],
            extra_cuda_cflags=[
                '-O3',
                '--use_fast_math',
                '-std=c++17',
                '-arch=sm_90',
            ],
            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