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

more cleaning

Browse Source
This commit is contained in:
Philippe Tillet
2019-06-24 12:37:13 -07:00
parent f257884eb7
commit 72867d17d4
7 changed files with 113 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -26,11 +26,12 @@ const tunable int32 TN = {64, 128};
const tunable int32 TK = {16};
const tunable int32 GZ = {1};
void matmul(restrict read_only fp16 *A, restrict read_only fp16 *B,
fp32 *C,
int32 M, int32 N, int32 K,
int32 lda, int32 ldb, int32 ldc,
int32 *locks, int32 grid0, int32 grid1) {
void matmul(restrict read_only align(4) fp16 *A,
restrict read_only align(4) fp16 *B,
align(4) fp32 *C,
int32 M, int32 N, int32 K,
int32 lda, int32 ldb, int32 ldc,
int32 *locks, int32 grid0, int32 grid1) {
int32 rxa[TM] = get_global_range[TM](0);
int32 ryb[TN] = get_global_range[TN](1);
int32 rz = get_global_range[1](2);
@@ -119,7 +120,7 @@ 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, {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 });