CroqTileThe Next-Gen GPU & DSA Language Achieves 5x Productivity

__co__ void spmm(global f8_e4m3 [M, PACKED_K] lhs_packed, global u8 [M, META_COLS] lhs_meta, global f8_e4m3 [N, K] rhs, global f16 [M, N] output) {
  parallel {block_m, block_n} by [cdiv(M, SPMM_WARP_M), cdiv(N, SPMM_WARP_N)] : block {
    shared event full[STAGES], empty[STAGES];
    shared f8_e4m3 [STAGES * SPMM_WARP_M, SPMM_PACKED_TILE_K] lhs_s;
    shared f8_e4m3 [STAGES * SPMM_WARP_N, SPMM_TILE_K] rhs_s;
    parallel p1 by 2 : group-4 {
      inthreads.async (p1 == 0) {
        foreach {iv_k} in [cdiv(K, SPMM_TILE_K)] {
          stage = iv_k % STAGES;
          wait empty[stage];
          tma.copy.async<full[stage]>.swiz<32> lhs_packed.subspan(SPMM_WARP_M, SPMM_PACKED_TILE_K).at(block_m, iv_k) => lhs_s.subspan(SPMM_WARP_M, SPMM_PACKED_TILE_K).at(stage, 0);
          trigger full[stage];
        }
      }
      inthreads.async (p1 == 1) {
        mc = mma.fill.f16 0.0f;
        foreach {iv_k} in [cdiv(K, SPMM_TILE_K)] {
          wait full[iv_k % STAGES];
          ma = mma.load.swiz<32> lhs_s.at(iv_k % STAGES, 0);
          mb = mma.load.swiz<64> rhs_s.at(iv_k % STAGES, 0);
          me = mma.load lhs_meta_s;
          mma.row.row.sp mc, ma, mb, me;
          trigger empty[iv_k % STAGES];
        }
        mma.store mc, output_s;
        tma.copy output_s => output.subspan(SPMM_WARP_M, SPMM_WARP_N).at(block_m, block_n);
      }
    }
  }
}