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

disabling interleaving

Browse Source
This commit is contained in:
Philippe Tillet
2019-06-13 17:16:00 -07:00
parent d487cf31ce
commit 21a9b92c87
4 changed files with 16 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -23,7 +23,7 @@ const char* src =
R"(
const tunable int32 TM = {64, 128};
const tunable int32 TN = {64, 128};
const tunable int32 TK = {16};
const tunable int32 TK = {32};
const tunable int32 GZ = {1};
void matmul(restrict read_only fp16 *A, restrict read_only fp16 *B,
@@ -117,10 +117,10 @@ class BlockSparseGemmOp : public OpKernel {
return 2.*M*N*K / ts * 1e-3;
};
// just-in-time compile source-code
jit.autotune("matmul", src, benchmark);
// 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, 16, 4, 128, 2, 2, 2, 2, 8, 32, 8, 1});
// jit.add_module("matmul", src, {8, 8, 128, 16, 8, 128, 2, 2, 2, 2, 16, 32, 8, 1 });
triton::driver::kernel* kernel = jit.get_function("matmul");
triton::jit::launch_information info = jit.get_launch_info("matmul");
std::cout << benchmark(kernel, info) << std::endl;;