in openai_gemm.py [0:0]
def _get_gemm_kernel(prefix, op, cda, cdb, cdc, m, n, k):
if op[0] == 'T':
vec4A = (cda & 3) == 0 and (m & 3) == 0
vec8A = (cda & 7) == 0 and (m & 7) == 0
dimA = (k,cda)
else:
vec4A = (cda & 3) == 0 and (k & 3) == 0
vec8A = (cda & 7) == 0 and (k & 7) == 0
dimA = (m,cda)
if op[1] == 'T':
vec4B = (cdb & 3) == 0 and (k & 3) == 0
vec8B = (cdb & 7) == 0 and (k & 7) == 0
dimB = (n,cdb)
else:
vec4B = (cdb & 3) == 0 and (n & 3) == 0
vec8B = (cdb & 7) == 0 and (n & 7) == 0
dimB = (k,cdb)
vec4C = (cdc & 3) == 0 and (n & 3) == 0
dimC = (m,cdc)
dtype = np.dtype(np.float32 if prefix == 's' else np.float16)
A = drv.mem_alloc(mul(*dimA) * dtype.itemsize)
B = drv.mem_alloc(mul(*dimB) * dtype.itemsize)
C = drv.mem_alloc(mul(*dimC) * dtype.itemsize)
# TODO: use curand
dataA = np.random.uniform(-1.0, 1.0, dimA).astype(dtype)
dataB = np.random.uniform(-1.0, 1.0, dimB).astype(dtype)
drv.memcpy_htod(int(A), dataA)
drv.memcpy_htod(int(B), dataB)
# Using random data gets you more accurate autotune results
# drv.memset_d8(int(A), 0, mul(*dimA) * dtype.itemsize)
# drv.memset_d8(int(B), 0, mul(*dimB) * dtype.itemsize)
timings = []
cache = []
# scale the repeat count to amount of work
repeat = min(max(int(5e11 * 28 / (m*n*k * 2.0 * _get_sm_count()) ), 10), 5000)
warmup = repeat
#print repeat
start, end = _get_events()
flops = m * n * k * 2.0
for tileM, tileN, tileK, vecA, vecB, vecC, div, base_op, dyn_shared in selections[prefix][op]:
vecA = (vecA == 4 and vec4A) or (vecA == 8 and vec8A)
vecB = (vecB == 4 and vec4B) or (vecB == 8 and vec8B)
vecC = vecC == 1 or vec4C
vec = vecA and vecB and vecC
if base_op:
# The op is part of the base kernel name
base = "%sgemm_%dx%dx%d_%s" % (prefix, tileM, tileN, tileK, op)
opts = ( "vec", ) if vec else ()
else:
# The op is an option passed to a more generic kernel
base = "%sgemm_%dx%dx%d" % (prefix, tileM, tileN, tileK)
opts = ( op, "vec" ) if vec else (op,)
kernel = get_kernel(base, opts)
blk_A = _ceil_div(m, tileM)
blk_B = _ceil_div(n, tileN)
# TODO: perhaps autotune all possible small divisors
blk_a, blk_A = _closest_divisor(blk_A, div)
blk_b, blk_B = _closest_divisor(blk_B, div)
if blk_a == 1:
blk_a, blk_A = (blk_A, 1)
for dynamic_shared in dyn_shared:
params = [
(blk_a * blk_b, blk_B, blk_A), (kernel.threads, 1, 1), None,
C, A, B, 1.0, 0.0,
cda, cdb, cdc, m, n, k, blk_a, blk_b ]
#print kernel.name, params, dynamic_shared
# Warmup (once per config)
for r in range(warmup):
kernel.prepared_async_call(*params)
warmup = 0
# Benchmark
start.record()
for r in range(repeat):
kernel.prepared_async_call(*params, shared_size=dynamic_shared)
end.record()
end.synchronize()
msecs = end.time_since(start) / float(repeat)
gflops = flops / (msecs * 1000000.0)
params[3:8] = (None,) * 5
timings.append((msecs, gflops, kernel, params, dynamic_shared))
cache.append((msecs, gflops, kernel.name, dynamic_shared))
major, minor = _get_compute_capability()
if prefix == "h" and major == 6 and minor == 0:
cublas_gemm = cublasXgemm["h2"]
else:
cublas_gemm = cublasXgemm[prefix]
# record a cublas time for reference
cublas_handle = _get_cublas()
start.record()
for r in range(repeat):
# convert row order to col order
cublas_gemm(cublas_handle, op[1], op[0], n, m, k, 1.0, B, cdb, A, cda, 0.0, C, cdc)
end.record()
end.synchronize()
msecs = end.time_since(start) / float(repeat)
gflops = flops / (msecs * 1000000.0)
cache.append( (msecs, gflops, "cuBLAS", 0) )
# cache complete timing data for benchmark comparisons
# this data could be cached to disk for quicker autotuning on future runs
_get_bench_data()[(prefix, op, cda, cdb, cdc, m, n, k)] = cache
# return the fastest kernel
return tuple(sorted(timings)[0][2:5])