225 lines
6.2 KiB
Python
225 lines
6.2 KiB
Python
# copying the kernels from https://github.com/microsoft/ArchProbe into Python
|
|
import numpy as np
|
|
import pickle
|
|
from tinygrad.runtime.ops_gpu import CLProgram, CLBuffer
|
|
from tinygrad.helpers import dtypes
|
|
from tqdm import trange, tqdm
|
|
from matplotlib import pyplot as plt
|
|
|
|
tests = {}
|
|
|
|
|
|
def register_test(fxn):
|
|
tests[fxn.__name__] = fxn
|
|
|
|
|
|
def warp_size2(nthread):
|
|
prg = """__kernel void warp_size2(
|
|
__global float* src,
|
|
__global int* dst,
|
|
const int niter,
|
|
const int prime_number
|
|
) {
|
|
int drain = 0;
|
|
for (int j = 0; j < niter; ++j) {
|
|
drain += j / prime_number;
|
|
barrier(0);
|
|
}
|
|
dst[get_local_id(0)] = drain;
|
|
}"""
|
|
src_buf = CLBuffer(1, dtypes.float32)
|
|
dst_buf = CLBuffer(1, dtypes.int32)
|
|
cl = CLProgram("warp_size2", prg, argdtypes=[None, None, np.int32, np.int32])
|
|
return (
|
|
min(
|
|
[
|
|
cl(
|
|
[nthread, 1024, 1],
|
|
[nthread, 1, 1],
|
|
src_buf,
|
|
dst_buf,
|
|
10,
|
|
3,
|
|
wait=True,
|
|
)
|
|
for _ in range(5)
|
|
]
|
|
)
|
|
* 1e9
|
|
)
|
|
|
|
|
|
@register_test
|
|
def test_warp_size():
|
|
return [(nthread, warp_size2(nthread)) for nthread in trange(1, 256)]
|
|
|
|
|
|
def reg_count(nthread, ngrp, nreg):
|
|
reg_declr = "".join(
|
|
[f"float reg_data{i} = (float)niter + {i};\n" for i in range(nreg)]
|
|
)
|
|
reg_comp = "".join([f"reg_data{i} *= {(i-1)%nreg};\n" for i in range(nreg)])
|
|
reg_reduce = "".join([f"out_buf[{i}] = reg_data{i};\n" for i in range(nreg)])
|
|
prg = f"""__kernel void reg_count(
|
|
__global float* out_buf,
|
|
__private const int niter
|
|
) {{
|
|
{reg_declr}
|
|
int i = 0;
|
|
for (; i < niter; ++i) {{
|
|
{reg_comp}
|
|
}}
|
|
i = i >> 31;
|
|
{reg_reduce}
|
|
}}"""
|
|
out_buf = CLBuffer(1, dtypes.float32)
|
|
cl = CLProgram("reg_count", prg, argdtypes=[None, np.int32])
|
|
return (
|
|
min(
|
|
[
|
|
cl([nthread, ngrp, 1], [nthread, 1, 1], out_buf, 20, wait=True)
|
|
for _ in range(10)
|
|
]
|
|
)
|
|
* 1e9
|
|
)
|
|
|
|
|
|
@register_test
|
|
def test_reg_count(nthread=1, ngrp=1):
|
|
base = reg_count(nthread, ngrp, 1)
|
|
return [
|
|
(nreg, (reg_count(nthread, ngrp, nreg) - base) / nreg)
|
|
for nreg in trange(4, 513, 4)
|
|
]
|
|
|
|
|
|
def buf_cache_hierarchy_pchase(ndata, stride=1, NCOMP=1, steps=65536):
|
|
ndata //= NCOMP * 4 # ptr size
|
|
prg = f"""__kernel void buf_cache_hierarchy_pchase(
|
|
__global int{str(NCOMP) if NCOMP > 1 else ''}* src,
|
|
__global int* dst,
|
|
const int niter
|
|
) {{
|
|
int idx = 0;
|
|
for (int i = 0; i < niter; ++i) {{
|
|
idx = src[idx]{'.x' if NCOMP > 1 else ''};
|
|
}}
|
|
*dst = idx;
|
|
}}"""
|
|
idx_buf = np.zeros(ndata * NCOMP, dtype=np.int32)
|
|
for i in range(ndata):
|
|
idx_buf[i * NCOMP] = (i + stride) % ndata
|
|
in_buf = CLBuffer.fromCPU(idx_buf)
|
|
out_buf = CLBuffer(1, dtypes.int32)
|
|
cl = CLProgram("buf_cache_hierarchy_pchase", prg, argdtypes=[None, None, np.int32])
|
|
return (
|
|
min(
|
|
[
|
|
cl([1, 1, 1], [1, 1, 1], in_buf, out_buf, steps, wait=True) / steps
|
|
for _ in range(5)
|
|
]
|
|
)
|
|
* 1e9
|
|
)
|
|
|
|
|
|
@register_test
|
|
def test_memory_latency():
|
|
# requires cacheline < 16
|
|
szs = [int(1.3**x) for x in range(20, 70)]
|
|
return [
|
|
(ndata, buf_cache_hierarchy_pchase(ndata, NCOMP=16, steps=128 * 1024))
|
|
for ndata in tqdm(szs)
|
|
]
|
|
|
|
|
|
@register_test
|
|
def test_cacheline_size():
|
|
# TODO: this buffer must be at least 2x the L1 cache for this test to work
|
|
return [
|
|
(stride, buf_cache_hierarchy_pchase(4 * 65536, stride, steps=65536))
|
|
for stride in trange(1, 64)
|
|
]
|
|
|
|
|
|
def cl_read(sz, niter=1):
|
|
prg = f"""__kernel void copy(
|
|
__global float4* src,
|
|
__global float* dst) {{
|
|
int gid = get_global_id(0);
|
|
if (src[gid].x == 99+get_global_id(1)) *dst = 1;
|
|
}}"""
|
|
|
|
in_buf = CLBuffer(sz // 4, dtypes.float32)
|
|
out_buf = CLBuffer(1, dtypes.float32)
|
|
cl = CLProgram("copy", prg)
|
|
# NOTE: if nay of the niters form a local group, this is wrong
|
|
return (
|
|
min(
|
|
[
|
|
cl([sz // 16, niter, 1], [1, 1, 1], in_buf, out_buf, wait=True)
|
|
for _ in range(10)
|
|
]
|
|
)
|
|
* 1e9
|
|
)
|
|
|
|
|
|
@register_test
|
|
def test_read_bandwidth():
|
|
szs = list(range(128 * 1024, 20 * 1024 * 1024, 128 * 1024))
|
|
NITER = 8
|
|
base = cl_read(16, niter=NITER)
|
|
return [(sz, (sz * NITER) / (cl_read(sz, niter=NITER) - base)) for sz in tqdm(szs)]
|
|
|
|
|
|
def gflops(niter=4, nroll=4, ngroups=4096):
|
|
NCOMP = 8
|
|
prg = f"""__kernel void gflops(
|
|
__global float* out_buf
|
|
) {{
|
|
float{NCOMP} x = (float{NCOMP})({",".join(f"get_local_id(0)+{i}" for i in range(NCOMP))});
|
|
float{NCOMP} y = (float{NCOMP})({",".join(f"get_local_id(1)+{i}" for i in range(NCOMP))});
|
|
|
|
for (int i = 0; i < {niter}; i++) {{
|
|
{''.join(['x = mad(y, y, x); y = mad(x, x, y);'+chr(10)]*nroll)}
|
|
}}
|
|
|
|
out_buf[get_global_id(0) >> 31] = {'+'.join(f"y.s{'0123456789abcdef'[i]}" for i in range(NCOMP))};
|
|
}}"""
|
|
out_buf = CLBuffer(1, dtypes.float32)
|
|
cl = CLProgram("gflops", prg, options="-cl-mad-enable -cl-fast-relaxed-math")
|
|
FLOPS = NCOMP * 2 * 2 * niter * nroll * ngroups * 32
|
|
# NOTE: if nay of the niters form a local group, this is wrong
|
|
return FLOPS / (
|
|
min([cl([32, ngroups, 1], [32, 1, 1], out_buf, wait=True) for _ in range(10)])
|
|
* 1e9
|
|
)
|
|
|
|
|
|
@register_test
|
|
def test_gflops():
|
|
return [(niter, gflops(niter=niter, nroll=32)) for niter in trange(1, 32, 1)]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cache = {}
|
|
# cache = pickle.load(open("/tmp/cache.pkl", "rb"))
|
|
# tests = {"test_cacheline_size": tests["test_cacheline_size"]}
|
|
plt.figure(figsize=(16, 9))
|
|
for i, (k, test) in enumerate(tests.items()):
|
|
print(f"running {k}")
|
|
plt.subplot(2, (len(tests) + 1) // 2, i + 1)
|
|
plt.title(k)
|
|
if k == "test_memory_latency":
|
|
plt.xscale("log")
|
|
if k not in cache:
|
|
cache[k] = test()
|
|
plt.plot(*zip(*cache[k]))
|
|
# pickle.dump(cache, open("/tmp/cache.pkl", "wb"))
|
|
|
|
plt.tight_layout(pad=0.5)
|
|
plt.savefig("/tmp/results.png")
|
|
plt.show()
|