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
|
# mypy: allow-untyped-defs
import functools
import itertools
import logging
from typing import List, Optional
from unittest.mock import patch
from ...autotune_process import TensorMeta
from ...ir import Buffer, IRNode, Layout
from ...utils import IndentedBuffer, unique
from ...virtualized import V
from ..common import KernelTemplate
from .rocm_benchmark_request import ROCmBenchmarkRequest
from .rocm_kernel import ROCmTemplateCaller, ROCmTemplateKernel
from .rocm_template_buffer import ROCmTemplateBuffer
log = logging.getLogger(__name__)
class ROCmTemplate(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 ROCm C++ Templates, derived from KernelTemplate. Not to be instantiated directly.
Args:
name (str): The name of the ROCmTemplate 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,
**kwargs,
) -> ROCmTemplateCaller:
"""
Generates the ROCm template caller object for the given GEMM template and operation. This ROCmTemplateCaller
may be used to call and benchmark the generated ROCm kernel in a standalone manner to enable Autotuning.
Args:
kwargs: Additional keyword arguments.
Returns:
A ROCmTemplateCaller object representing the generated ROCm template caller.
"""
kernel_name = f"rocm_{self.name}"
kernel_hash_name = f"rocm_{self.name}_{next(self.index_counter)}"
with patch.object(
V.graph, "get_dtype", self._fake_get_dtype(self.output_node)
), ROCmTemplateKernel(
kernel_name=kernel_name,
) as kernel:
code = self.render(kernel=kernel, **kwargs)
_, call_args, _, _ = kernel.args.python_argdefs()
log.debug("Autotune key: %s, Generated Code:\n%s", kernel_hash_name, 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,
)
size_args = (
self.size_args() if hasattr(self, "size_args") else ()
) # subclass should define def size_args()
size_args_ints = [
V.graph.sizevars.size_hint(arg) for arg in size_args
] # resolve to ints for benchmarking
bmreq = ROCmBenchmarkRequest(
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_ints,
source_code=code,
)
def make_kernel_render(
template_node: ROCmTemplateBuffer,
epilogue_nodes: Optional[List[IRNode]] = None,
):
kernel = ROCmTemplateKernel(
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 ROCmTemplateCaller(
kernel_hash_name,
self.name,
self.input_nodes,
self.output_node.get_layout(),
make_kernel_render,
bmreq,
self,
kwargs,
)
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
// as long as there is no custom arithmetic it's fine
using bfloat16 = uint16_t;
using float8_e4m3fnuz = uint8_t;
using float8_e5m2fnuz = uint8_t;
"""
)
return res
def render(self, **kwargs) -> str:
raise NotImplementedError
|