cutile-python/test/test_list.py at main · NVIDIA/cutile-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

import torch
import pytest
import math
from unittest.mock import patch

import cuda.tile
import cuda.tile as ct
from cuda.tile._bytecode import BytecodeVersion
from typing import Annotated
from util import assert_equal
from conftest import requires_tileiras


@ct.kernel
def add_arrays(arrays, out):
    res = ct.zeros((16, 16), dtype=out.dtype)
    for i in range(len(arrays)):
        t = ct.load(arrays[i], (0, 0), (16, 16))
        res += t
    ct.store(out, (0, 0), res)


@ct.kernel
def add_arrays_with_const_index(arrays, out):
    tx = ct.load(arrays[0], (0, 0), (16, 16))
    ty = ct.load(arrays[1], (0, 0), (16, 16))
    tz = ct.load(arrays[2], (0, 0), (16, 16))
    res = tx + ty + tz
    ct.store(out, (0, 0), res)


@ct.kernel
def add_arrays_with_0d_tile_index(arrays, out):
    bid = ct.full((), 0, dtype=ct.int32)
    tx = ct.load(arrays[bid], (0, 0), (16, 16))
    ty = ct.load(arrays[bid + 1], (0, 0), (16, 16))
    tz = ct.load(arrays[bid + 2], (0, 0), (16, 16))
    res = tx + ty + tz
    ct.store(out, (0, 0), res)


@pytest.mark.parametrize("kernel", [
    add_arrays,
    add_arrays_with_const_index,
    add_arrays_with_0d_tile_index
    ])
def test_add_list_of_arrays(kernel):
    arrays = [torch.randint(0, 100, (16, 16), dtype=torch.int32, device="cuda") for _ in range(3)]
    out = torch.zeros(16, 16, dtype=torch.int32, device="cuda")
    ref = sum(arrays)

    ct.launch(torch.cuda.current_stream(), (1,), kernel, (arrays, out))
    assert_equal(out, ref)


ListOfArrayIndexedWithInt64 = Annotated[
    list, ct.ListAnnotation(element=ct.IndexedWithInt64)
]

ListWithStaticShape = Annotated[
    list, ct.ListAnnotation(element=ct.ArrayAnnotation(static_shape_dims=(0, 1)))
]


@ct.kernel
def add_int64_index_arrays(
    arrays: ListOfArrayIndexedWithInt64,
    out: ct.IndexedWithInt64,
    TILE: ct.Constant[int]
):
    bid = ct.bid(0)
    res = ct.zeros((TILE, 1), dtype=out.dtype)
    for i in range(len(arrays)):
        t = ct.load(arrays[i], (bid, 0), (TILE, 1))
        res += t
    ct.store(out, (bid, 0), res)


@ct.kernel
def add_static_shape_arrays(arrays: ListWithStaticShape, out):
    res = ct.zeros((16, 16), dtype=out.dtype)
    for i in range(len(arrays)):
        t = ct.load(arrays[i], (0, 0), (16, 16))
        res += t
    ct.store(out, (0, 0), res)


def test_add_list_static_shape():
    k = cuda.tile.kernel(add_static_shape_arrays._pyfunc)

    # (16,16) → 1st compile; (32,32) → 2nd compile; (16,16) again → cache hit.
    shapes = [(16, 16), (32, 32), (16, 16)]
    with patch('cuda.tile._compile.compile_tile',
               side_effect=cuda.tile._compile.compile_tile) as mock_compile:
        for shape in shapes:
            arrays = [torch.randint(0, 100, shape, dtype=torch.int32, device="cuda")
                      for _ in range(3)]
            out = torch.zeros((16, 16), dtype=torch.int32, device="cuda")
            ct.launch(torch.cuda.current_stream(), (1,), k, (arrays, out))
            assert_equal(out, sum(a[:16, :16] for a in arrays))

    assert mock_compile.call_count == 2


def test_add_list_static_shape_mismatch():
    k = cuda.tile.kernel(add_static_shape_arrays._pyfunc)

    arrays = [
        torch.zeros((16, 16), dtype=torch.int32, device="cuda"),
        torch.zeros((32, 16), dtype=torch.int32, device="cuda"),
    ]
    out = torch.zeros((16, 16), dtype=torch.int32, device="cuda")
    with pytest.raises(ValueError, match="vary in static shape at axis 0"):
        ct.launch(torch.cuda.current_stream(), (1,), k, (arrays, out))


@requires_tileiras(BytecodeVersion.V_13_3)
def test_add_list_of_int64_index_arrays():
    """
    Sum a list of large 2D arrays whose stride[0] exceeds INT32_MAX.

    This test may be excluded from selected CI jobs with
    ``-k "not int64_index"`` because it requires a very large allocation.
    Keep ``int64_index`` in the test name unless those CI filters are updated.
    """
    TILE = 2048
    n = (1 << 32) + TILE  # shape[0] > UINT32_MAX
    arrays = [torch.full((n, 1), i + 1, device='cuda', dtype=torch.int8) for i in range(3)]
    out = torch.zeros(n, 1, device='cuda', dtype=torch.int8)

    grid = (math.ceil(n / TILE), 1, 1)
    ct.launch(torch.cuda.current_stream(), grid, add_int64_index_arrays, (arrays, out, TILE))
    assert (out == 6).all().item()