Function test_matmul

extra/gemm/rdna4_asm_matmul.py:203–245

()

N = getenv("N", 4096)

def test_matmul():
  dev = Device[Device.DEFAULT]
  arch = getattr(dev.renderer, 'arch', 'gfx1200')
  print(f"Device arch: {arch}")
  insts = build_kernel(N, arch)

  rng = np.random.default_rng(42)
  a = Tensor(rng.random((N, N), dtype=np.float32).astype(np.float16))
  b = Tensor(rng.random((N, N), dtype=np.float32).astype(np.float16))
  c = Tensor.empty(N, N, dtype=dtypes.half)
  Tensor.realize(a, b, c)

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

  dname = Device.DEFAULT
  def asm_kernel(A, B, C):
    gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
    lidxs = [UOp.special(THREADS, "lidx0")]
    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=max(LDS_SIZE, 65536//getenv("LIMIT_OCC",2)), 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*2*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()
    c_np = c.float().numpy()
    a_np, b_np = a.float().numpy(), b.float().numpy()
    ref = a_np @ b_np
    err = np.sqrt(np.mean((c_np - ref)**2)) / np.sqrt(np.mean(ref**2))
    print(f"relative RMSE {err:.6f}")
    if err != err or err > 0.05: raise RuntimeError(f"matmul is wrong! RMSE={err}")

if __name__ == "__main__":
  test_matmul()