submission 230383

  import torch
- import triton
- import triton.language as tl
+ from torch.utils.cpp_extension import load_inline
  from task import input_t, output_t
  
  # =============================================================================
- # Phase 1.3: 明示的分解版（並列乗算 + 最後に加算）
+ # Inline CUDA: float4 ベクトルロード版 RGB to Grayscale (Phase 3: Inline PTX)
  # =============================================================================
- # 目標: 2.44ms → 2.38ms
+ # 目標: 2.47ms (Triton) → 2.38ms
  #
- # 仮説: 累積加算(gray += x * coef)はFMA依存チェーンを形成
- #       並列乗算(term_x = x * coef)後に加算することでILP向上
- #       3つの乗算が並列実行可能になる
+ # Phase 3 変更点:
+ #   - Inline PTX で ld.global.v4.f32 を明示的に使用
+ #   - LDG.128 命令を完全に強制
+ #   - コンパイラの最適化判断を完全にバイパス
  #
- # Phase 1.1 結果: G→R→B = 2.47ms（効果なし）
- # Phase 1.2 結果: B→R→G = 2.47ms（効果なし）
- #
- # 変更点:
- # - 3つのチャンネルを並列でロード（独立Block Pointer）
- # - 3つの乗算を並列で実行
- # - 最後に3項を加算
+ # PTX 命令:
+ #   ld.global.v4.f32 {%0, %1, %2, %3}, [%4];
+ #   - 128-bit (16 bytes) 一括ロード
+ #   - 4 つの float を同時に取得
  
- @triton.autotune(
-     configs=[
-         # 中ブロック
-         triton.Config({'BLOCK_SIZE': 2048}, num_warps=4, num_stages=3),
-         triton.Config({'BLOCK_SIZE': 2048}, num_warps=4, num_stages=4),
-         triton.Config({'BLOCK_SIZE': 2048}, num_warps=4, num_stages=5),
+ cuda_src = """
+ #include <cuda_runtime.h>
  
-         # 大ブロック
-         triton.Config({'BLOCK_SIZE': 4096}, num_warps=4, num_stages=4),
-         triton.Config({'BLOCK_SIZE': 4096}, num_warps=4, num_stages=5),
-         triton.Config({'BLOCK_SIZE': 4096}, num_warps=8, num_stages=3),
-         triton.Config({'BLOCK_SIZE': 4096}, num_warps=8, num_stages=4),
-         triton.Config({'BLOCK_SIZE': 4096}, num_warps=8, num_stages=5),
+ // 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)
+     );
+ }
  
-         # 超大ブロック
-         triton.Config({'BLOCK_SIZE': 8192}, num_warps=8, num_stages=3),
-         triton.Config({'BLOCK_SIZE': 8192}, num_warps=8, num_stages=4),
-         triton.Config({'BLOCK_SIZE': 8192}, num_warps=8, num_stages=5),
-         triton.Config({'BLOCK_SIZE': 8192}, num_warps=16, num_stages=3),
-         triton.Config({'BLOCK_SIZE': 8192}, num_warps=16, num_stages=4),
-     ],
-     key=['n_elements'],
- )
- @triton.jit
- def grayscale_kernel(
-     input_ptr, output_ptr,
-     n_elements,
-     BLOCK_SIZE: tl.constexpr
- ):
-     pid = tl.program_id(axis=0)
+ // 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;
  
-     # 独立した3つのBlock Pointer（並列ロード可能）
-     ptr_r = tl.make_block_ptr(
-         base=input_ptr,
-         shape=(n_elements, 3),
-         strides=(3, 1),
-         offsets=(pid * BLOCK_SIZE, 0),
-         block_shape=(BLOCK_SIZE, 1),
-         order=(1, 0)
-     )
-     ptr_g = tl.make_block_ptr(
-         base=input_ptr,
-         shape=(n_elements, 3),
-         strides=(3, 1),
-         offsets=(pid * BLOCK_SIZE, 1),
-         block_shape=(BLOCK_SIZE, 1),
-         order=(1, 0)
-     )
-     ptr_b = tl.make_block_ptr(
-         base=input_ptr,
-         shape=(n_elements, 3),
-         strides=(3, 1),
-         offsets=(pid * BLOCK_SIZE, 2),
-         block_shape=(BLOCK_SIZE, 1),
-         order=(1, 0)
-     )
+     if (pixel_base >= n_pixels) return;
  
-     # 並列ロード（依存関係なし）
-     r = tl.load(ptr_r, boundary_check=(0,), padding_option="zero")
-     g = tl.load(ptr_g, boundary_check=(0,), padding_option="zero")
-     b = tl.load(ptr_b, boundary_check=(0,), padding_option="zero")
+     int base = pixel_base * 3;
+     const float* ptr = input + base;
  
-     # 並列乗算（依存関係なし、ILP最大化）
-     term_r = r * 0.2989
-     term_g = g * 0.5870
-     term_b = b * 0.1140
+     // 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]
  
-     # 最後に加算
-     gray = term_r + term_g + term_b
+     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);
  
-     # Output (1Dポインタ)
-     offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-     mask = offs < n_elements
-     tl.store(output_ptr + offs, tl.reshape(gray, (BLOCK_SIZE,)), mask=mask)
+     // 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
-     n_elements = output_tensor.numel()
  
-     grid = lambda META: (triton.cdiv(n_elements, META['BLOCK_SIZE']),)
+     module = _get_module()
+     module.rgb_to_grayscale(input_tensor, output_tensor)
  
-     grayscale_kernel[grid](
-         input_tensor, output_tensor,
-         n_elements
-     )
- 
      return output_tensor