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

some work on conv

Browse Source
This commit is contained in:
Philippe Tillet
2019-10-31 18:08:27 -04:00
parent 91a2fd463b
commit 739a8d9061
10 changed files with 278 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/examples/conv.py Normal file
View File

@@ -0,0 +1,11 @@
import torch
import triton
N, C, K = 32, 32, 32
H, W = 32, 32
R, S = 3, 3
a = torch.randn(N, C, H, W).cuda()
b = torch.randn(C, R, S, K).cuda()
#c = torch.nn.functional.conv2d(a, b)
c = triton.ops.conv(a, b)
print(c)

4
python/examples/dot.py
View File

@@ -53,8 +53,8 @@ def run_torch():
triton_da = a.grad.clone()
triton_db = b.grad.clone()
nanosec = triton.bench_registry[triton_d]
print('TFLOPS:', 2. * M * N * K / nanosec * 1e-3)
#nanosec = triton.bench_registry[triton_d]
#print('TFLOPS:', 2. * M * N * K / nanosec * 1e-3)
print('Diff DA:', (torch_da - triton_da).max())
print('Diff DB:', (torch_db - triton_db).max())