gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[feature] added basic tensor core support

Browse Source
This commit is contained in:
Philippe Tillet
2019-06-09 14:41:36 -07:00
parent d074a166e2
commit 06b5992509
9 changed files with 180 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
examples/python/tensorflow/dot.cpp
View File

@@ -21,8 +21,8 @@ using GPUDevice = Eigen::GpuDevice;
const char* src =
R"(
const tunable int32 TM = {16};
const tunable int32 TN = {16};
const tunable int32 TM = {8, 16, 32, 64, 128};
const tunable int32 TN = {8, 16, 32, 64, 128};
const tunable int32 TK = {8};
const tunable int32 GZ = {1};
@@ -54,11 +54,8 @@ void matmul(restrict read_only fp16 *A, restrict read_only fp16 *B,
}
int32 rxc[TM] = get_global_range[TM](0);
int32 ryc[TN] = get_global_range[TN](1);
fp32* pc[TM, TN] = C + ryc[newaxis, :]*ldc + rxc[:, newaxis];
int1 checkc0[TM] = rxc < M;
int1 checkc1[TN] = ryc < N;
int1 checkc[TM, TN] = checkc0[:, newaxis] && checkc1[newaxis, :];
@checkc *pc = c;
fp32* pc[TM, TN] = C + ryc[newaxis, :]*ldc + rxc[:, newaxis];
*pc = c;
}
)";
@@ -122,14 +119,17 @@ class BlockSparseGemmOp : public OpKernel {
stream->enqueue(kernel, grid, {nthreads, 1, 1});
stream->synchronize();
double ts = triton::tools::bench([&](){stream->enqueue(kernel, grid, {nthreads, 1, 1});},
[&](){ stream->synchronize(); }, nullptr);
[&](){ stream->synchronize(); }, ctx->device());
return 2.*M*N*K / ts * 1e-3;
};
// just-in-time compile source-code
jit.add_module("matmul", src, {4, 2, 16, 4, 2, 16, 2, 2, 1, 1, 8, 8, 8, 1});
// jit.autotune("matmul", src, benchmark);
// jit.add_module("matmul", src, {4, 2, 8, 4, 2, 32, 1, 4, 1, 1, 8, 8, 8, 1});
// jit.add_module("matmul", src, {32, 2, 128, 32, 2, 128, 2, 2, 2, 2, 4, 8, 4, 1});
jit.add_module("matmul", src, {16, 4, 128, 32, 4, 128, 2, 2, 2, 2, 8, 8, 4, 1});
triton::driver::kernel* kernel = jit.get_function("matmul");
triton::jit::launch_information info = jit.get_launch_info("matmul");
benchmark(kernel, info);
std::cout << benchmark(kernel, info) << std::endl;;
}
private:

27
examples/python/tensorflow/run.py
View File

@@ -1,28 +1,39 @@
import os
import tensorflow as tf
import numpy as np
from time import time
data_files_path = tf.resource_loader.get_data_files_path()
library_dir = '/home/philippe/development/triton/build/examples/python/tensorflow'
module = tf.load_op_library(os.path.join(library_dir, 'libtf_blocksparse.so'))
M, N, K = 256, 256, 256
M, N, K = 256,256,256
a = tf.placeholder(tf.float16, shape=[M, K])
b = tf.placeholder(tf.float16, shape=[N, K])
locks = tf.placeholder(tf.int32, shape=[4096])
# c = tf.matmul(a, b, transpose_a=True)
c = module.dot(a, b, locks)
# Reference
ha = np.random.rand(M, K).astype(np.float16)
hb = np.random.rand(N, K).astype(np.float16)
hresult = np.dot(hb.T, ha)
# Run
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
result = sess.run([c], feed_dict = {locks: np.zeros(4096),
a: ha,
b: hb})
print(result)
print(hresult)
#print(result - hresult)
print(np.max(np.abs(result - hresult)))
b: hb})[0]
#bench = tf.test.Benchmark().run_op_benchmark(sess=sess,
# op_or_tensor=c,
# feed_dict={a: ha, b: hb},
# min_iters=100)
#print(end - start)
#print(2*M*N*K / (end - start) * 1e-12)
hresult = np.dot(ha.T, hb).T
dif = np.abs(result - hresult)
print("dif: %f" % np.max(dif))
#np.savetxt("dif.txt", dif, fmt="%5.2f")
#np.savetxt("gpu.txt", result, fmt="%5.2f")
#np.savetxt("cpu.txt", hresult, fmt="%5.2f")