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

[LANG] Various (relatively minor) improvements (#320)

Browse Source
This commit is contained in:
Philippe Tillet
2021-10-04 18:39:40 -07:00
committed by GitHub
parent 12b6158c5c
commit 5123db0b7d
10 changed files with 59 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/test/unit/language/test_core.py
View File

@@ -515,6 +515,22 @@ def test_dot(epilogue, device='cuda'):
assert 'ld.global.v4' in ptx
assert 'st.global.v4' in ptx
# ---------------
# test arange
# ---------------
@pytest.mark.parametrize("start", [0, 1, 7, 16])
def test_arange(start, device='cuda'):
BLOCK = 128
z_tri = torch.empty(BLOCK, dtype=torch.int32, device=device)
@triton.jit
def _kernel(z, **meta):
off = tl.arange(0, meta['BLOCK'])
val = tl.arange(meta['START'], meta['END'])
tl.store(z + off, val)
_kernel[(1,)](z_tri, START=start, END=start+BLOCK, BLOCK=BLOCK)
z_ref = torch.arange(start, BLOCK+start, dtype=torch.int32, device=device)
triton.testing.assert_almost_equal(z_tri, z_ref)
# ---------------
# test load