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

Merge branch 'master' of https://github.com/ptillet/triton

Browse Source
This commit is contained in:
Philippe Tillet
2019-09-09 02:29:54 -04:00
parent 4a69af08e7 96bdae25d5
commit c0fa4e0a57
1 changed files with 10 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/examples/dot.py
View File

@@ -2,7 +2,6 @@ import numpy as np
import triton import triton
def run_tf(): def run_tf():
import tensorflow as tf
M, N, K = 128, 128, 128 M, N, K = 128, 128, 128
a = tf.placeholder(tf.float32, shape=[M, K]) a = tf.placeholder(tf.float32, shape=[M, K])
b = tf.placeholder(tf.float32, shape=[N, K]) b = tf.placeholder(tf.float32, shape=[N, K])
@@ -28,7 +27,6 @@ def run_tf():
print("dif: %f" % np.max(dif)) print("dif: %f" % np.max(dif))
def run_torch(): def run_torch():
import torch
torch.manual_seed(0) torch.manual_seed(0)
M, N, K = 128, 128, 128 M, N, K = 128, 128, 128
a = torch.randn(M, K).cuda() a = torch.randn(M, K).cuda()
@@ -56,5 +54,14 @@ def run_torch():
print('Diff DA:', (torch_da - triton_da).max()) print('Diff DA:', (torch_da - triton_da).max())
print('Diff DB:', (torch_db - triton_db).max()) print('Diff DB:', (torch_db - triton_db).max())
try:
import tensorflow as tf
run_tf()
except ModuleNotFoundError:
pass
try:
import torch
run_torch() run_torch()
except ModuleNotFoundError:
pass