|
1 | 1 | import torch |
2 | | - |
| 2 | +import pytest |
| 3 | +import math |
3 | 4 | import triton |
4 | 5 | from triton.backends.compiler import GPUTarget |
5 | 6 | import triton.language as tl |
@@ -35,27 +36,61 @@ def test(device): |
35 | 36 | x = torch.rand([n_cols, n_rows], device=device, dtype=torch.float32) |
36 | 37 | output = torch.empty([n_cols], device=device, dtype=x.dtype) |
37 | 38 | BLOCK_SIZE = n_rows |
38 | | - grid = lambda meta: (n_cols,) |
| 39 | + grid = lambda meta: (n_cols, ) |
39 | 40 |
|
40 | 41 | reduce_kernel_2d[grid](x, output, x.stride(0), n_rows, BLOCK_SIZE=BLOCK_SIZE) |
41 | 42 | ans = torch.sum(x, dim=1) |
42 | 43 | torch.testing.assert_close(output, ans, rtol=0.001, atol=1e-5) |
43 | 44 |
|
44 | 45 | # TODO: need to check some conditions otherwise the code below does not make any difference for the test |
45 | 46 | src = triton.compiler.ASTSource( |
46 | | - fn=reduce_kernel_2d, |
47 | | - signature={"x_ptr": "*fp32", |
48 | | - "output_ptr": "*fp32", |
49 | | - "stride": "i32", |
50 | | - "n_elements": "i32", |
51 | | - "BLOCK_SIZE": "constexpr"}, |
52 | | - constexprs={"BLOCK_SIZE": 32} |
53 | | - ) |
54 | | - ret = triton.compile( |
55 | | - src, |
56 | | - target=GPUTarget(device, 0, 0) |
57 | | - ) |
| 47 | + fn=reduce_kernel_2d, signature={ |
| 48 | + "x_ptr": "*fp32", "output_ptr": "*fp32", "stride": "i32", "n_elements": "i32", "BLOCK_SIZE": "constexpr" |
| 49 | + }, constexprs={"BLOCK_SIZE": 32}) |
| 50 | + ret = triton.compile(src, target=GPUTarget(device, 0, 0)) |
58 | 51 | print(ret.asm["ttir"]) |
59 | 52 | print(ret.asm["ttsharedir"]) |
60 | 53 | print(ret.asm["llir"]) |
61 | 54 | print(ret.asm["obj"]) |
| 55 | + |
| 56 | + |
| 57 | +@pytest.mark.interpreter |
| 58 | +@pytest.mark.parametrize("dtype_str", ["int32", "float32"]) |
| 59 | +@pytest.mark.parametrize("shape", [(128, 2, 4), (64, 2, 4), (32, 2, 4), (2, 4, 32), (2, 4, 2)]) |
| 60 | +@pytest.mark.parametrize("axis", [0, 1, 2]) |
| 61 | +def test_reduce_max(dtype_str, shape, axis, device): |
| 62 | + |
| 63 | + @triton.jit |
| 64 | + def kernel( |
| 65 | + In, |
| 66 | + Out, |
| 67 | + in_shape1: tl.constexpr, |
| 68 | + in_shape2: tl.constexpr, |
| 69 | + in_shape3: tl.constexpr, |
| 70 | + ou_shape1: tl.constexpr, |
| 71 | + ou_shape2: tl.constexpr, |
| 72 | + axis: tl.constexpr, |
| 73 | + ): |
| 74 | + in_desc = tl.make_tensor_descriptor( |
| 75 | + base=In, |
| 76 | + shape=[in_shape1 * in_shape2 * in_shape3], |
| 77 | + strides=[1], |
| 78 | + block_shape=[in_shape1 * in_shape2 * in_shape3], |
| 79 | + ) |
| 80 | + out_desc = tl.make_tensor_descriptor( |
| 81 | + base=Out, |
| 82 | + shape=[ou_shape1 * ou_shape2], |
| 83 | + strides=[1], |
| 84 | + block_shape=[ou_shape1 * ou_shape2], |
| 85 | + ) |
| 86 | + val = in_desc.load([0]).reshape(in_shape1, in_shape2, in_shape3) |
| 87 | + output = tl.max(val, axis=axis) |
| 88 | + out_desc.store([0], output.reshape(out_desc.block_shape)) |
| 89 | + |
| 90 | + input = torch.arange(math.prod(shape), dtype=getattr(torch, dtype_str), |
| 91 | + device="cpu").reshape(shape).to(device=device) |
| 92 | + expected, indices = torch.max(input, dim=axis) |
| 93 | + actual = torch.zeros(expected.shape, dtype=getattr(torch, dtype_str), device=device) |
| 94 | + kernel[(1, )](input, actual, *shape, *expected.shape, axis=axis) |
| 95 | + |
| 96 | + assert torch.equal(expected, actual) |
0 commit comments