#!POPCORN leaderboard amd-fp8-mm
#!POPCORN gpus MI300

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

os.environ["CXX"] = "clang++"
os.environ["PYTORCH_ROCM_ARCH"] = "gfx942"

cuda_src = r"""
#include <iostream>
#include <hip/hip_runtime.h>
#include <hip/amd_detail/amd_hip_bf16.h>
#include <hip/amd_detail/amd_hip_fp8.h>

__host__ __device__ __forceinline__ int cdiv(int a, int b) {
  return (a + b - 1) / b;
}

#define BLOCK_DIM 128

template <const uint32_t BN, const uint32_t BK, const uint32_t BM,
          const uint32_t WITERN, const uint32_t WITERM, const uint32_t SM_COUNT>
__global__ void fp8_mm_kernel_N(const __hip_fp8_e4m3_fnuz *A,
                                const __hip_fp8_e4m3_fnuz *B,
                                const float *A_scale, const float *B_scale,
                                float *C, uint32_t N, uint32_t K, uint32_t M) {
  using uint8x16 =
      __attribute__((__vector_size__(16 * sizeof(uint8_t)))) uint8_t;
  using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;

  static constexpr uint32_t VECTOR_SIZE = 4;
  static constexpr uint32_t WARPSIZE = 64;
  static constexpr uint32_t WN = 32 * WITERN;
  static constexpr uint32_t WM = 32 * WITERM;
  static constexpr uint32_t numThreads = (BN * BM) / (16 * WITERN * WITERM);
  static constexpr uint32_t strideA = (numThreads / BN) * VECTOR_SIZE * 4;
  static constexpr uint32_t strideB = (numThreads / BM) * VECTOR_SIZE * 4;

  static_assert(numThreads % BN == 0, "BN should be a multiple of numThreads");
  static_assert(numThreads % BM == 0, "BM should be a multiple of numThreads");
  static_assert(BK <= 128 && BM <= 128, "Range above 128 is not supported");

  uint32_t numTiles = cdiv(N, BN) * cdiv(M, BM);
  uint32_t subTilesPerTile = cdiv(K, BK);
  uint32_t totalSubTiles = numTiles * subTilesPerTile;
  uint32_t subTilesToProcess = cdiv(totalSubTiles, SM_COUNT);
  uint32_t startSubTileIdx = subTilesToProcess * blockIdx.x;
  uint32_t endSubTileIdx = startSubTileIdx + subTilesToProcess;
  uint32_t startTileIdx = startSubTileIdx / subTilesPerTile;
  uint32_t subTile = startSubTileIdx;

  for (uint32_t tileIdx = startTileIdx, subTilesProcessed = 0;
       tileIdx < numTiles && subTilesProcessed < subTilesToProcess; ++tileIdx) {
    uint32_t rowOffsetC = (tileIdx % cdiv(N, BN)) * BN;
    uint32_t colOffsetC = (tileIdx / cdiv(N, BN)) * BM;
    uint32_t colOffsetA = rowOffsetC;
    uint32_t colOffsetB = colOffsetC;
    uint32_t M_scale = cdiv(M, BLOCK_DIM);

    uint32_t innerColA = threadIdx.x % BN;
    uint32_t innerRowA = (threadIdx.x / BN) * VECTOR_SIZE * 4;
    uint32_t innerColB = threadIdx.x % BM;
    uint32_t innerRowB = (threadIdx.x / BM) * VECTOR_SIZE * 4;

    uint32_t laneIdx = threadIdx.x % WARPSIZE;
    uint32_t warpIdx = threadIdx.x / WARPSIZE;
    uint32_t warpColOffset = (warpIdx % (BM / WM)) * WM;
    uint32_t warpRowOffset = (warpIdx / (BM / WM)) * WN;
    uint32_t warpX = laneIdx % 32;
    uint32_t warpY = laneIdx / 32;

    __shared__ __hip_fp8_e4m3_fnuz As[BN][BK + 4], Bs[BM][BK + 4];
    __shared__ float Ws[BN + 1];

    floatx16 d[WITERN][WITERM] = {0};

    for (subTile = subTile % subTilesPerTile;
         subTile < subTilesPerTile && subTilesProcessed < subTilesToProcess;
         ++subTile, ++subTilesProcessed) {
      uint32_t tileOffset = subTile * BK;
      // load from global to shared memory in coalesced manner
      for (uint32_t innerRowOffsetA = 0; innerRowOffsetA < BK;
           innerRowOffsetA += strideA) {
        if ((tileOffset + innerRowOffsetA + innerRowA) < K &&
            (colOffsetA + innerColA) < N &&
            (innerRowOffsetA + innerRowA) < BK) {
          uint8x16 pack;

          pack[0] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA) * N +
                 (colOffsetA + innerColA)]);
          pack[1] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 1) * N +
                 (colOffsetA + innerColA)]);
          pack[2] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 2) * N +
                 (colOffsetA + innerColA)]);
          pack[3] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 3) * N +
                 (colOffsetA + innerColA)]);
          pack[4] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 4) * N +
                 (colOffsetA + innerColA)]);
          pack[5] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 5) * N +
                 (colOffsetA + innerColA)]);
          pack[6] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 6) * N +
                 (colOffsetA + innerColA)]);
          pack[7] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 7) * N +
                 (colOffsetA + innerColA)]);
          pack[8] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 8) * N +
                 (colOffsetA + innerColA)]);
          pack[9] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 9) * N +
                 (colOffsetA + innerColA)]);
          pack[10] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 10) * N +
                 (colOffsetA + innerColA)]);
          pack[11] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 11) * N +
                 (colOffsetA + innerColA)]);
          pack[12] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 12) * N +
                 (colOffsetA + innerColA)]);
          pack[13] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 13) * N +
                 (colOffsetA + innerColA)]);
          pack[14] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 14) * N +
                 (colOffsetA + innerColA)]);
          pack[15] = *reinterpret_cast<const uint8_t *>(
              &A[(tileOffset + innerRowOffsetA + innerRowA + 15) * N +
                 (colOffsetA + innerColA)]);

          *reinterpret_cast<uint8x16 *>(
              &As[innerColA][innerRowOffsetA + innerRowA]) = pack;
        } else if ((innerRowOffsetA + innerRowA) < BK) {
          *reinterpret_cast<uint8x16 *>(
              &As[innerColA][innerRowOffsetA + innerRowA]) = {0};
        }
      }
      if (threadIdx.x < (BN / VECTOR_SIZE)) {
        *reinterpret_cast<float4 *>(&Ws[threadIdx.x * VECTOR_SIZE]) =
            *reinterpret_cast<const float4 *>(
                &A_scale[(tileOffset / BLOCK_DIM) * N +
                         (colOffsetA + threadIdx.x * VECTOR_SIZE)]);
      }
      for (uint32_t innerRowOffsetB = 0; innerRowOffsetB < BK;
           innerRowOffsetB += strideB) {
        if ((tileOffset + innerRowOffsetB + innerRowB) < K &&
            (colOffsetB + innerColB) < M &&
            (innerRowOffsetB + innerRowB) < BK) {
          uint8x16 pack;

          pack[0] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB) * M +
                 (colOffsetB + innerColB)]);
          pack[1] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 1) * M +
                 (colOffsetB + innerColB)]);
          pack[2] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 2) * M +
                 (colOffsetB + innerColB)]);
          pack[3] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 3) * M +
                 (colOffsetB + innerColB)]);
          pack[4] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 4) * M +
                 (colOffsetB + innerColB)]);
          pack[5] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 5) * M +
                 (colOffsetB + innerColB)]);
          pack[6] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 6) * M +
                 (colOffsetB + innerColB)]);
          pack[7] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 7) * M +
                 (colOffsetB + innerColB)]);
          pack[8] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 8) * M +
                 (colOffsetB + innerColB)]);
          pack[9] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 9) * M +
                 (colOffsetB + innerColB)]);
          pack[10] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 10) * M +
                 (colOffsetB + innerColB)]);
          pack[11] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 11) * M +
                 (colOffsetB + innerColB)]);
          pack[12] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 12) * M +
                 (colOffsetB + innerColB)]);
          pack[13] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 13) * M +
                 (colOffsetB + innerColB)]);
          pack[14] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 14) * M +
                 (colOffsetB + innerColB)]);
          pack[15] = *reinterpret_cast<const uint8_t *>(
              &B[(tileOffset + innerRowOffsetB + innerRowB + 15) * M +
                 (colOffsetB + innerColB)]);

          *reinterpret_cast<uint8x16 *>(
              &Bs[innerColB][innerRowOffsetB + innerRowB]) = pack;
        } else if ((innerRowOffsetB + innerRowB) < BK) {
          *reinterpret_cast<uint8x16 *>(
              &Bs[innerColB][innerRowOffsetB + innerRowB]) = {0};
        }
      }
      if (threadIdx.x == numThreads - 1) { // load using different warp
        Ws[BN] = B_scale[(tileOffset / BLOCK_DIM) * M_scale +
                         (colOffsetB / BLOCK_DIM)];
      }

      __syncthreads();

      long a[WITERN], b[WITERM];
      float b_scale = Ws[BN];
      floatx16 c[WITERN][WITERM] = {0};

      for (uint32_t BKOffset = 0; BKOffset < BK; BKOffset += 16) {
        for (uint32_t wn = 0; wn < WITERN; ++wn) {
          a[wn] = *reinterpret_cast<long *>(
              &As[warpRowOffset + wn * 32 + warpX][BKOffset + warpY * 8]);
        }
        for (uint32_t wm = 0; wm < WITERM; ++wm) {
          b[wm] = *reinterpret_cast<long *>(
              &Bs[warpColOffset + wm * 32 + warpX][BKOffset + warpY * 8]);
        }
        for (uint32_t wn = 0; wn < WITERN; ++wn) {
          for (uint32_t wm = 0; wm < WITERM; ++wm) {
            c[wn][wm] = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
                a[wn], b[wm], c[wn][wm], 0, 0, 0);
          }
        }
      }
      for (uint32_t wn = 0; wn < WITERN; ++wn) {
        for (uint32_t wm = 0; wm < WITERM; ++wm) {
          for (uint32_t j = 0; j < 4; ++j) {
            for (uint32_t i = 0; i < 4; ++i) {
              d[wn][wm][i + j * 4] +=
                  c[wn][wm][i + j * 4] *
                  Ws[warpRowOffset + wn * 32 + j * 8 + warpY * 4 + i] * b_scale;
            }
          }
        }
      }

      __syncthreads();
    }

    for (uint32_t wn = 0; wn < WITERN; ++wn) {
      for (uint32_t wm = 0; wm < WITERM; ++wm) {
        for (uint32_t j = 0; j < 4; ++j) {
          for (uint32_t i = 0; i < 4; ++i) {
            if ((rowOffsetC + warpRowOffset + wn * 32 + j * 8 + warpY * 4 + i) <
                    N &&
                (colOffsetC + warpColOffset + wm * 32 + warpX) < M) {
              atomicAdd(&C[(rowOffsetC + warpRowOffset + wn * 32 + j * 8 +
                            warpY * 4 + i) *
                               M +
                           (colOffsetC + warpColOffset + wm * 32 + warpX)],
                        d[wn][wm][i + j * 4]);
            }
          }
        }
      }
    }
  }
}

__global__ void bf16ToFloat(float *in, __hip_bfloat16 *out, uint32_t N) {
  using uint16x4 =
      __attribute__((__vector_size__(4 * sizeof(uint16_t)))) uint16_t;

  uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx * 4 < N) {
    float4 val_in = *reinterpret_cast<float4 *>(&in[idx * 4]);
    __hip_bfloat16 x = val_in.x;
    __hip_bfloat16 y = val_in.y;
    __hip_bfloat16 z = val_in.z;
    __hip_bfloat16 w = val_in.w;

    uint16x4 val_out;
    val_out[0] = *reinterpret_cast<uint16_t *>(&x);
    val_out[1] = *reinterpret_cast<uint16_t *>(&y);
    val_out[2] = *reinterpret_cast<uint16_t *>(&z);
    val_out[3] = *reinterpret_cast<uint16_t *>(&w);
    *reinterpret_cast<uint16x4 *>(&out[idx * 4]) = val_out;
  }
}

at::Tensor fp8_mm(at::Tensor A, at::Tensor B, at::Tensor A_scale,
                  at::Tensor B_scale, at::Tensor C) {
  int N = A.size(0), K = A.size(1), M = B.size(0);
  at::Tensor Cf32 = at::zeros(
      {N, M}, at::TensorOptions().device(at::kCUDA).dtype(at::kFloat));

  const uint32_t BK = 128;
  const uint32_t BN = 128;
  const uint32_t BM = 128;
  const uint32_t WITERN = 1;
  const uint32_t WITERM = 1;
  const uint32_t SM_COUNT = 304;
  const uint32_t COL_SCHEDULING = cdiv(N, BN) >= cdiv(M, BM);
  dim3 numThreads((BN * BM) / (16 * WITERN * WITERM));
  dim3 numBlocks(SM_COUNT);
  fp8_mm_kernel_N<BN, BK, BM, WITERN, WITERM, SM_COUNT>
      <<<numBlocks, numThreads>>>((__hip_fp8_e4m3_fnuz *)A.data_ptr(),
                                  (__hip_fp8_e4m3_fnuz *)B.data_ptr(),
                                  A_scale.data_ptr<float>(),
                                  B_scale.data_ptr<float>(),
                                  (float *)Cf32.data_ptr(), N, K, M);
                                  
  return Cf32;
}
"""

cpp_src = r"""
at::Tensor fp8_mm(at::Tensor A, at::Tensor B, at::Tensor A_scale,
                  at::Tensor B_scale, at::Tensor C);
"""

if sys.stdout is None:
    sys.stdout = open("/dev/stdout", "w")
if sys.stderr is None:
    sys.stderr = open("/dev/stderr", "w")

module = load_inline(
    name="fp8_mm",
    cpp_sources=[cpp_src],
    cuda_sources=[cuda_src],
    functions=["fp8_mm"],
    verbose=True,
    extra_cuda_cflags=[
        "-Ofast",
        "--offload-arch=gfx942",
        "-std=c++20",
        "-ffp-contract=fast",
        "-lhip_hcc",
        "-mcumode",
    ],
)


def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    return module.fp8_mm(a, b, a_scale, b_scale, c)

python eval.py leaderboard /tmp/tmp8ita2xcg
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip [skipped, already hipified]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 1
[1/1] clang++ main.o hip.cuda.o -shared -L/usr/local/lib/python3.10/dist-packages/torch/lib -lc10 -lc10_hip -ltorch_cpu -ltorch_hip -ltorch -ltorch_python -L/opt/rocm-6.3.1/lib -lamdhip64 -o fp8_mm.so
