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