cutile-python/test/test_range.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
# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

import pytest
import torch

import cuda.tile as ct
from cuda.tile._exception import TileTypeError
from util import assert_equal


@ct.kernel
def copy_range(stop, out):
    idx = 0
    for i in range(stop):
        tile = ct.full((1,), fill_value=i, dtype=ct.int32)
        ct.store(out, (idx,), tile=tile)
        idx = idx + 1


@ct.kernel
def copy_range_2(start, stop, out):
    idx = 0
    for i in range(start, stop):
        tile = ct.full((1,), fill_value=i, dtype=ct.int32)
        ct.store(out, (idx,), tile=tile)
        idx = idx + 1


@ct.kernel
def copy_range_2_step_negative(start, stop, out):
    idx = 0
    for i in range(start, stop, -2):
        tile = ct.full((1,), fill_value=i, dtype=ct.int32)
        ct.store(out, (idx,), tile=tile)
        idx = idx + 1


@ct.kernel
def copy_range_3(start, stop, step, out):
    idx = 0
    for i in range(start, stop, step):
        tile = ct.full((1,), fill_value=i, dtype=ct.int32)
        ct.store(out, (idx,), tile=tile)
        idx = idx + 1


def test_range_stop():
    x = torch.zeros((10,), dtype=torch.int32, device='cuda')
    ct.launch(torch.cuda.default_stream(), (1,), copy_range, (10, x))
    assert_equal(x, torch.arange(10, device='cuda', dtype=torch.int32))


def test_range_start_stop():
    start, stop = 1, 11
    L = len(range(start, stop))
    x = torch.zeros((L,), dtype=torch.int32, device='cuda')
    ct.launch(torch.cuda.default_stream(), (1,), copy_range_2, (start, stop, x))
    assert_equal(x, torch.arange(start, stop, device='cuda', dtype=torch.int32))


def test_range_start_stop_negative_constant_step():
    start, stop = 1, 11
    L = len(range(start, stop))
    x = torch.zeros((L,), dtype=torch.int32, device='cuda')
    with pytest.raises(TileTypeError, match='Step must be positive, got -2'):
        ct.launch(torch.cuda.default_stream(), (1,), copy_range_2_step_negative, (start, stop, x))


def test_range_start_stop_positive_step():
    start, stop, step = 1, 11, 2
    L = len(range(start, stop, step))
    x = torch.zeros((L,), dtype=torch.int32, device='cuda')
    ct.launch(torch.cuda.default_stream(), (1,), copy_range_3, (start, stop, step, x))
    assert_equal(x, torch.arange(start, stop, step, device='cuda', dtype=torch.int32))


@pytest.mark.xfail(reason="Issue 314")
def test_range_negative_step():
    start, stop, step = 11, 1, -2
    L = len(range(start, stop, step))
    x = torch.zeros((L,), dtype=torch.int32, device='cuda')
    ct.launch(torch.cuda.default_stream(), (1,), copy_range_3, (start, stop, step, x))
    assert_equal(x, torch.arange(start, stop, step, device='cuda', dtype=torch.int32))


def test_range_scalar_type_error():
    x = torch.zeros((10,), dtype=torch.int32, device='cuda')
    with pytest.raises(TileTypeError, match='Expected a scalar or a 0D tile'):
        ct.launch(torch.cuda.default_stream(), (1,), copy_range, (x, x))


def test_range_integer_type_error():
    x = 1.5
    with pytest.raises(TileTypeError, match='Expected a signed integer'):
        ct.launch(torch.cuda.default_stream(), (1,), copy_range, (x, x))