submission 232981

  ⋯ 2 unchanged lines
  from task import input_t, output_t
  
  # =============================================================================
- # Inline CUDA: B200最適化版 RGB to Grayscale (__launch_bounds__ 調整版)
+ # Inline CUDA: B200 TMA + LDS.128 + blockIdx.x (シンプル版)
  # =============================================================================
  # 目標: 598.336μs以下（B200 1位）
  #
- # 最適化: __launch_bounds__(256, 6) でOccupancy向上
- #   - maxThreadsPerBlock = 256: ブロックあたり最大256スレッド
- #   - minBlocksPerMultiprocessor = 6: SM当たり最低6ブロックを保証
- #   → v2_13の(256, 4)から調整、より高いOccupancyを狙う
+ # 最適化戦略:
+ #   1. TMA (cp.async.bulk) による非同期バルクコピー
+ #   2. blockIdx.x による通常グリッド起動（cudaMalloc排除）
+ #   3. ld.shared.v4.f32 (LDS.128) で共有メモリから一括ロード
+ #   4. st.global.v4.f32 (STG.128) による出力ベクトル化
+ #   5. 1スレッド4ピクセル処理でスループット最大化
  #
- # v2_13からの変更点:
- #   - __launch_bounds__(256, 4) → __launch_bounds__(256, 6)
- #   - ロジックは完全に同一
- #
- # PTX 命令:
- #   ld.global.v4.f32 {%0, %1, %2, %3}, [%4];  - 128-bit ロード
- #   st.global.v4.f32 [%0], {%1, %2, %3, %4};  - 128-bit ストア
+ # 前回の失敗から学んだ教訓:
+ #   - cudaMalloc/cudaFree は致命的なオーバーヘッド → 排除
+ #   - Persistent Threads + atomicAdd は逆効果 → 排除
+ #   - 12回の個別LDS命令 → LDS.128 × 3回に統合
+ # =============================================================================
  
  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
+ // =============================================================================
+ // 定数定義
+ // =============================================================================
+ constexpr int PIXELS_PER_TILE   = 512;       // 512ピクセル/タイル
+ constexpr int BYTES_PER_TILE    = 6144;      // 512 × 3 × 4 = 6144バイト (16の倍数)
+ constexpr int FLOATS_PER_TILE   = 1536;      // 512 × 3 = 1536 floats
+ 
+ constexpr int THREADS_PER_BLOCK = 128;       // 128スレッド/ブロック
+ constexpr int PIXELS_PER_THREAD = 4;         // 4ピクセル/スレッド (128 × 4 = 512)
+ 
+ // =============================================================================
+ // 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(
-         "ld.global.v4.f32 {%0, %1, %2, %3}, [%4];"
-         : "=f"(r0), "=f"(r1), "=f"(r2), "=f"(r3)
-         : "l"(ptr)
+         "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))
      );
  }
  
- // Inline PTX による float4 ストア
+ // 【新規】共有メモリからのベクトルロード (LDS.128)
+ // smem_addrはバイトアドレス（__cvta_generic_to_shared済み）
+ __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
+     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"
+         :: "l"(ptr), "f"(v0), "f"(v1), "f"(v2), "f"(v3) : "memory"
      );
  }
  
- // RGB to Grayscale カーネル（__launch_bounds__ 調整版）
- // __launch_bounds__(256, 6): SM当たり6ブロック（1536スレッド）を保証
- __global__ __launch_bounds__(256, 6) void grayscale_kernel(
+ // =============================================================================
+ // TMA + LDS.128 カーネル（blockIdx.x版、シンプル）
+ // =============================================================================
+ __global__ __launch_bounds__(128, 8) void grayscale_tma_lds128_kernel(
      const float* __restrict__ input,
      float* __restrict__ output,
-     int n_pixels
+     const int n_pixels
  ) {
-     int tid = blockIdx.x * blockDim.x + threadIdx.x;
-     int pixel_base = tid * 4;
+     // 共有メモリ: シングルバッファ + mbarrier
+     __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;
+ 
+     // 境界チェック（最終タイルのみ必要だが、完全に割り切れるので実質不要）
      if (pixel_base >= n_pixels) return;
  
-     int base = pixel_base * 3;
-     const float* ptr = input + base;
+     // ========================================
+     // Step 1: mbarrier初期化（Thread 0のみ）
+     // ========================================
+     if (tid == 0) {
+         mbarrier_init(&mbar, 1);
+     }
+     __syncthreads();
  
-     // 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]
+     // ========================================
+     // 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);
+     }
  
-     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);
+     // ========================================
+     // Step 3: TMA完了待機
+     // ========================================
+     mbarrier_wait(&mbar, 0);
  
-     // 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
+     // ========================================
+     // Step 4: LDS.128で共有メモリから一括ロード
+     // ========================================
+     // 共有メモリアドレス計算
+     // tid * 12 floats = tid * 48 bytes (常に16の倍数)
+     uint32_t smem_addr = (uint32_t)__cvta_generic_to_shared(&smem[tid * 12]);
  
-     // PTX で 128-bit ストア
-     store_float4_ptx(output + pixel_base, gray_x, gray_y, gray_z, gray_w);
+     // float4 × 3回でレジスタに取り込み
+     // 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計算
+     // ========================================
+     // マッピング（v2_13と同一）:
+     // 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);
  }
  
+ // =============================================================================
+ // ホストインターフェース（cudaMalloc排除！）
+ // =============================================================================
  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;
+     // シンプルなブロック数計算（cudaMalloc不要！）
+     int blocks = (n_pixels + PIXELS_PER_TILE - 1) / PIXELS_PER_TILE;
  
-     grayscale_kernel<<<blocks, threads>>>(
+     grayscale_tma_lds128_kernel<<<blocks, THREADS_PER_BLOCK>>>(
          input.data_ptr<float>(),
          output.data_ptr<float>(),
          n_pixels
  ⋯ 15 unchanged lines
      global _module
      if _module is None:
          _module = load_inline(
-             name='grayscale_cuda_b200_launch_bounds_v2',
+             name='grayscale_cuda_tma_lds128',
              cuda_sources=[cuda_src],
              cpp_sources=[cpp_src],
              functions=['rgb_to_grayscale'],
-             extra_cuda_cflags=['-O3', '--use_fast_math'],
+             extra_cuda_cflags=[
+                 '-O3',
+                 '--use_fast_math',
+                 '-std=c++17',
+                 '-arch=sm_90',
+             ],
              verbose=False
          )
      return _module