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

[examples/python/tensorflow] fixed ABI compatibility issue in JIT

Browse Source 
destructor
This commit is contained in:
Philippe Tillet
2019-05-01 13:38:56 -04:00
parent 7b6efc0463
commit 55866f1ef6
5 changed files with 47 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
examples/python/tensorflow/blocksparse.cpp
View File

@@ -66,8 +66,26 @@ void matmul(restrict read_only fp32 *A, restrict read_only fp32 *B, fp32 *C,
fp32 b[TN, 1] = checkb ? *pb : 0;
c = dot(a, trans(b), c);
}
int32 ridx = get_range_id(0);
int32 ridy = get_range_id(1);
fp32* pc[TM, TN] = C + ryc[newaxis, :]*ldc + rxc[:, newaxis];
*pc = c;
int32 *plock = locks + ridx + ridy*grid0;
while(__atomic_cas(plock, 0, 1));
int32 *pcount = plock + grid0*grid1;
int32 count = *pcount;
int32 countp1 = select(count == GZ - 1, 0, count + 1);
int1 checkc0[TM] = rxc < M;
int1 checkc1[TN] = ryc < N;
int1 checkc[TM, TN] = checkc0[:, newaxis] && checkc1[newaxis, :];
if(count == 0) {
@checkc *pc = c;
*pcount = countp1;
}
else {
@checkc *pc = c + *pc;
*pcount = countp1;
}
__atomic_cas(plock, 1, 0);
}
)";
@@ -137,7 +155,6 @@ class BlockSparseGemmOp : public OpKernel {
kernel->setArg(10, grid[0]);
kernel->setArg(11, grid[1]);
stream->enqueue(kernel, grid, {nthreads, 1, 1});
stream->synchronize();
}
private:

2
examples/python/tensorflow/blocksparse.py
View File

@@ -3,7 +3,7 @@ import tensorflow as tf
import numpy as np
data_files_path = tf.resource_loader.get_data_files_path()
library_dir = '/home/philippe/development/triton/build/examples/python/tensorflow'
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 = 512, 512, 512