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

added fragmented axis

Browse Source
This commit is contained in:
Philippe Tillet
2019-06-07 10:32:56 -07:00
parent 781b6d377d
commit 6fce9f28ae
3 changed files with 59 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -6,14 +6,14 @@ 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 = 512, 512, 512
M, N, K = 16, 16, 16
a = tf.placeholder(tf.float16, shape=[M, K])
b = tf.placeholder(tf.float16, shape=[N, K])
locks = tf.placeholder(tf.int32, shape=[4096])
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)
ha = np.ones((M, K)).astype(np.float16)
hb = np.ones((N, K)).astype(np.float16)
hresult = np.dot(hb.T, ha)
# Run