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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
|
# Owner(s): ["module: inductor"]
import sys
import unittest
import torch
from torch.testing._internal.common_utils import IS_LINUX, skipIfXpu
from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU
try:
import triton # noqa: F401 # @manual
import triton.language as tl # @manual
except ImportError:
if __name__ == "__main__":
sys.exit(0)
raise unittest.SkipTest("requires triton") # noqa: B904
from torch._inductor import config
from torch._inductor.runtime.hints import (
AttrsDescriptorWrapper,
AutotuneHint,
DeviceProperties,
HeuristicType,
TRITON_MAX_BLOCK,
)
from torch._inductor.runtime.triton_helpers import math as tl_math
from torch._inductor.runtime.triton_heuristics import (
autotune_hints_to_configs,
CachingAutotuner,
triton_config,
)
from torch._inductor.test_case import run_tests, TestCase
class TestTritonHeuristics(TestCase):
device_type = GPU_TYPE
def test_triton_config(self):
"""
Make sure block size does not exceed the maximum defined in inductor config.
"""
cfg = triton_config({"x": 2048, "y": 2}, 64, 64)
for label in "XYZ":
key = f"{label}BLOCK"
if key not in cfg.kwargs:
continue
self.assertTrue(cfg.kwargs[key] <= TRITON_MAX_BLOCK[label])
def _test_artificial_zgrid(self):
def forward(primals_1, primals_2, primals_5):
view = torch.ops.aten.reshape.default(primals_5, [-1, 2, 4])
primals_5 = None
permute = torch.ops.aten.permute.default(view, [0, 2, 1])
clone = torch.ops.aten.clone.default(
permute, memory_format=torch.contiguous_format
)
permute = None
view_1 = torch.ops.aten.reshape.default(clone, [-1, 4])
clone = None
permute_1 = torch.ops.aten.permute.default(primals_1, [1, 0])
primals_1 = None
addmm = torch.ops.aten.addmm.default(primals_2, view_1, permute_1)
primals_2 = None
return addmm
s0 = 16777472
s1 = 8
args = [
torch.rand([2, 4], device=GPU_TYPE),
torch.rand([2], device=GPU_TYPE),
torch.rand([s0, s1], device=GPU_TYPE),
]
torch._dynamo.mark_dynamic(args[-1], 0)
foo_c = torch.compile(forward)
self.assertEqual(forward(*args), foo_c(*args))
args = [
torch.rand([2, 4], device=GPU_TYPE),
torch.rand([2], device=GPU_TYPE),
torch.rand([s0, s1], device=GPU_TYPE),
]
self.assertEqual(forward(*args), foo_c(*args))
@skipIfXpu
def test_artificial_zgrid(self):
self._test_artificial_zgrid()
@skipIfXpu
@config.patch("cpp_wrapper", True)
def test_artificial_grid_cpp_wrapper(self):
self._test_artificial_zgrid()
def _get_cos_kernel_caching_autotuner_args(self):
@triton.jit
def triton_(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
xnumel = 16
xoffset = tl.program_id(0) * XBLOCK
xindex = xoffset + tl.arange(0, XBLOCK)[:]
xmask = xindex < xnumel
x0 = xindex
tmp0 = tl.load(in_ptr0 + (x0), xmask)
tmp1 = tl_math.cos(tmp0)
tl.store(out_ptr0 + (x0), tmp1, xmask)
triton_meta = {
"signature": {"in_ptr0": "*fp32", "out_ptr0": "*fp32", "xnumel": "i32"},
"device": DeviceProperties.create(torch.device("cuda")),
"constants": {},
"configs": [
AttrsDescriptorWrapper(divisible_by_16=(0, 1, 2), equal_to_1=())
],
}
configs = [
triton_config({"x": 16}, 64),
triton_config({"x": 256}, 64),
]
inductor_meta = {}
return {
"fn": triton_,
"triton_meta": triton_meta,
"configs": configs,
"save_cache_hook": False,
"mutated_arg_names": [],
"reset_to_zero_arg_names": [],
"optimize_mem": True,
"heuristic_type": HeuristicType.POINTWISE,
"inductor_meta": inductor_meta,
}
@skipIfXpu
def test_pre_hook_assert(self):
# assert if any of the configs passed to the CachingAutotuner have pre-hooks
args = self._get_cos_kernel_caching_autotuner_args()
def pre_hook(kwargs):
if "in_ptr0" in kwargs:
kwargs["in_ptr0"].zero_()
for cfg in args["configs"]:
cfg.pre_hook = pre_hook
with self.assertRaisesRegex(AssertionError, "pre_hook"):
autotuner = CachingAutotuner(**args)
def test_autotune_hints_to_configs(self):
device_props = DeviceProperties.create(torch.device(GPU_TYPE))
device_props = device_props._replace(warp_size=8)
hints = {AutotuneHint.ONE_ELEMENT_PER_THREAD}
size_hints = (1024,)
block_size = 256
seen_num_elements_per_warp = set()
def mock_triton_config(
size_hints,
x,
y=None,
z=None,
num_stages=None,
num_elements_per_warp=None,
min_elem_per_thread=None,
):
seen_num_elements_per_warp.add(num_elements_per_warp)
return None
with unittest.mock.patch(
"torch._inductor.runtime.triton_heuristics.triton_config",
mock_triton_config,
):
_ = autotune_hints_to_configs(hints, size_hints, block_size, device_props)
self.assertTrue(8 in seen_num_elements_per_warp)
if __name__ == "__main__":
if IS_LINUX and HAS_GPU:
run_tests()
|
