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
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
|
# mypy: allow-untyped-defs
import functools
import itertools
import logging
from typing import List, Optional
from unittest.mock import patch
import sympy
import torch
from ...autotune_process import CUDABenchmarkRequest, TensorMeta
from ...ir import Buffer, CUDATemplateBuffer, IRNode, Layout
from ...utils import IndentedBuffer, unique
from ...virtualized import V
from ..common import KernelTemplate
from .cuda_kernel import CUDATemplateCaller, CUDATemplateKernel
log = logging.getLogger(__name__)
class CUDATemplate(KernelTemplate):
index_counter = itertools.count()
def __init__(
self,
name: str,
input_nodes: List[Buffer],
layout: Layout,
input_reorder: Optional[List[int]] = None,
) -> None:
"""
Baseclass for CUDA C++ Templates, derived from KernelTemplate. Not to be instantiated directly.
Args:
name (str): The name of the CUDATemplate object.
input_nodes (List[IRNode]): A list of input IRNodes.
layout (Layout): The layout of the output buffer / tensor.
input_reorder (Optional[List[int]]): An optional list that specifies the order of the input nodes.
"""
super().__init__(name)
self.input_nodes = input_nodes
self.output_node: Buffer = Buffer(name="buf_out", layout=layout)
self.input_reorder = input_reorder
self.layout = layout
def generate( # type: ignore[override]
self,
description,
**kwargs,
) -> CUDATemplateCaller:
"""
Generates the CUDA template caller object for the given GEMM template and operation. This CUDATemplateCaller
may be used to call and benchmark the generated CUDA kernel in a standalone manner to enable Autotuning.
Args:
kwargs: Additional keyword arguments.
Returns:
A CUDATemplateCaller object representing the generated CUDA template caller.
"""
kernel_name = f"cuda_{self.name}"
with patch.object(
V.graph, "get_dtype", self._fake_get_dtype(self.output_node)
), CUDATemplateKernel(
kernel_name=kernel_name,
) as kernel:
code = self.render(kernel=kernel, **kwargs)
_, call_args, _, _ = kernel.args.python_argdefs()
log.debug("Generated Code:\n%s", code)
log.debug(
"Args: cpp_argdefs: %s, python_argdefs: %s",
kernel.args.cpp_argdefs(),
kernel.args.python_argdefs(),
)
input_reorder = (
self.input_reorder
if self.input_reorder is not None
else list(range(len(self.input_nodes)))
)
expected_args = list(
unique(self.input_nodes[idx].get_name() for idx in input_reorder)
)
expected_args.extend([self.output_node.get_name()])
assert list(call_args)[: len(expected_args)] == expected_args, (
call_args,
expected_args,
)
extra_args = V.graph.sizevars.size_hints(
map(sympy.expand, call_args[len(expected_args) :])
)
size_args = V.graph.sizevars.size_hints(kernel.get_layout_args())
kernel_hash_name = f"cuda_{self.name}_{next(self.index_counter)}"
# create the BenchmarkRequest
bmreq = CUDABenchmarkRequest(
kernel_name=kernel_name,
input_tensor_meta=TensorMeta.from_irnodes(self.input_nodes),
output_tensor_meta=TensorMeta.from_irnodes(self.output_node),
extra_args=size_args,
source_code=code,
)
def make_kernel_render(
template_node: CUDATemplateBuffer,
epilogue_nodes: Optional[List[IRNode]] = None,
):
kernel = CUDATemplateKernel(
kernel_name="KERNEL_NAME",
)
render = functools.partial(
self.render,
kernel=kernel,
template_buffer_node=template_node,
epilogue_nodes=epilogue_nodes,
**kwargs, # includes "op" argument in case of CUTLASSGemmTemplate
)
return kernel, render
return CUDATemplateCaller(
kernel_hash_name,
self.name,
self.input_nodes,
self.output_node.get_layout(),
make_kernel_render,
bmreq,
self,
kwargs,
description,
)
def header(self) -> IndentedBuffer:
res = IndentedBuffer()
res.splice(
"""
#include <exception>
#include <iostream>
#include <memory>
#include <random>
#include <vector>
"""
)
return res
def globals(self) -> IndentedBuffer:
res = IndentedBuffer()
res.splice(
"""
// We compile all models with -fvisibility=hidden. Any symbols that need to be
// exposed in the final shared library must be declared with PT_EXPORT to make
// them visible.
#ifdef __GNUC__ // Applies to any compiler with GNU extensions (clang and g++)
#define PT_EXPORT __attribute__((__visibility__("default")))
#else
#ifdef _WIN32
#define PT_EXPORT __declspec(dllexport)
#else
#define PT_EXPORT
#endif
#endif
using bfloat16 = nv_bfloat16;
"""
)
return res
def render(self, **kwargs) -> str:
raise NotImplementedError
class CUTLASSTemplate(CUDATemplate):
"""
CUTLASSTemplate is a class that provides a template for generating CUTLASS Templates. Used as a baseclass for the
CUTLASSGemmTemplate, providing functionality that might also be relevant for non-GEMM CUTLASS Kernels.
"""
def header(self) -> IndentedBuffer:
res = super().header()
res.splice(
"""
#include "cute/tensor.hpp"
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/device/tensor_fill.h"
#include "cutlass/util/device_memory.h"
"""
)
return res
def globals(self) -> IndentedBuffer:
res = super().globals()
res.splice(
"""
using namespace cute;
#define CUTLASS_CHECK(status) \\
{ \\
cutlass::Status error = status; \\
if (error != cutlass::Status::kSuccess) { \\
auto msg = std::string("[") + __FILE__ + "] Got cutlass error: " + \\
cutlassGetStatusString(error) + " at: " + std::to_string(__LINE__); \\
throw std::runtime_error(msg); \\
} \\
}
// Used as pass-through functor in EVT just for type casting / rounding
template <typename T>
struct identity_op {
CUTLASS_HOST_DEVICE
T operator()(T val) const { return val; }
};
"""
)
return res
def cute_int(self, int_str: str, var_name: str) -> str:
res = ""
if int_str in {"1", "1L"}:
res = "cute::Int<1>{}"
else:
res = int_str
return f"{res} /* {var_name} */"
_DTYPE_TO_CUTLASS = {
torch.float32: "float",
torch.float64: "double",
torch.float16: "cutlass::half_t",
torch.int32: "int32_t",
torch.int16: "int16_t",
torch.int8: "int8_t",
torch.uint8: "uint8_t",
torch.bool: "bool",
torch.bfloat16: "cutlass::bfloat16_t",
}
_DTYPE_TO_CUTLASS_SPARSE_META = {
torch.int32: "uint32_t",
torch.int16: "uint16_t",
}
def cutlass_type_cast(self, node: IRNode, ptr: str) -> str:
if node is None:
return ptr
else:
return f"({self._DTYPE_TO_CUTLASS.get(node.get_dtype())}*)({ptr})"
def cutlass_sparse_meta_type_cast(self, node: IRNode, ptr: str) -> str:
if node is None:
return ptr
else:
return (
f"({self._DTYPE_TO_CUTLASS_SPARSE_META.get(node.get_dtype())}*)({ptr})"
)
|
