Function test_matmul

extra/gemm/amd_asm_matmul.py:442–497

()

THREADS = 128

def test_matmul():
  dev = Device[Device.DEFAULT]
  print(f"Device arch: {dev.renderer.target.arch}")

  insts = build_kernel(N)

  rng = np.random.default_rng(42)
  a = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
  b = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
  c = Tensor.empty(N, N)
  Tensor.realize(a, b, c)

  grid, local = (N // BLOCK_N, N // BLOCK_M, 1), (THREADS, 1, 1)
  print(f"Grid: {grid}, Local: {local}")

  dname:str = Device.DEFAULT
  def asm_kernel(A:UOp, B:UOp, C:UOp) -> UOp:
    gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
    lidxs = [UOp.special(n, f"lidx{i}") for i,n in enumerate(local)]
    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536)), addrspace=AddrSpace.LOCAL), (), 'lds')
    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
                                                                                  estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
  linear = c.schedule_linear()

  ets = []
  with Context(DEBUG=2):
    for _ in range(getenv("CNT", 5)):
      start = GlobalCounters.time_sum_s
      run_linear(linear)
      ets.append(GlobalCounters.time_sum_s - start)
  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")

  if getenv("VERIFY", 1):
    GlobalCounters.reset()
    with Context(DEBUG=2): tc = (a @ b).realize()
    with Context(DEBUG=0): err = (c - tc).square().mean().item()
    print(f"mean squared error {err}")
    if err != err or err > 1e-06:
      c_np, tc_np = c.numpy(), tc.numpy()
      for bi in range(N // 128):
        for bj in range(N // 128):
          blk_c = c_np[bi*128:(bi+1)*128, bj*128:(bj+1)*128]
          blk_ref = tc_np[bi*128:(bi+1)*128, bj*128:(bj+1)*128]
          blk_diff = blk_c - blk_ref
          zero_rows = [i for i in range(128) if np.all(np.abs(blk_c[i,:]) < 1e-10)]
          nz_rows = [i for i in range(128) if i not in zero_rows]
          nz_mse = float(np.mean(blk_diff[nz_rows,:]**2)) if nz_rows else 0
          print(f"Block ({bi},{bj}): zero_rows={zero_rows}, nz_rows_mse={nz_mse:.2e}")
          # show first few non-zero row comparisons
          if nz_rows and nz_mse > 1e-6:
            for r in nz_rows[:3]:
              print(f"  row {r} asm[0:8]:  {blk_c[r,:8]}")
              print(f"  row {r} ref[0:8]:  {blk_ref[r,:8]}")
      raise RuntimeError("matmul is wrong!")

if __name__ == "__main__":