Skip to content
Snippets Groups Projects
Select Git revision
  • master default protected
1 result
Blame Permalink
benchmark.py 5.20 KiB
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 






from typing import Union, List
from collections import namedtuple
from pathlib import Path

from pystencils.astnodes import KernelFunction
from pystencils.enums import Backend
from pystencils.typing import get_base_type
from pystencils.sympyextensions import prod
from pystencils.transformations import get_common_field

from pystencils_benchmark.common import (_env,
                                         _kernel_source,
                                         _kernel_header,
                                         compiler_toolchain,
                                         copy_static_files,
                                         setup_directories)
from pystencils_benchmark.enums import Compiler


def _add_launch_bound(code: str, launch_bounds: tuple) -> str:
    lb_str = f"__launch_bounds__({', '.join(str(lb) for lb in launch_bounds)}) "
    splitted = code.split("void ")
    prefix = splitted[0]
    if code.startswith("void "):
        # just in case that there is nothing before the first void
        prefix = ""
    return prefix + "void " + lb_str + "void ".join(splitted[1:])


def generate_benchmark(kernel_asts: Union[KernelFunction, List[KernelFunction]],
                       path: Path = None,
                       *,
                       compiler: Compiler = Compiler.NVCC,
                       timing: bool = True,
                       launch_bounds: tuple = None,
                       ) -> None:

    src_path, include_path = setup_directories(path)

    if isinstance(kernel_asts, KernelFunction):
        kernel_asts = [kernel_asts]

    for kernel_ast in kernel_asts:
        kernel_name = kernel_ast.function_name

        header = kernel_header(kernel_ast)
        if launch_bounds:
            header = _add_launch_bound(header, launch_bounds)
        with open(include_path / f'{kernel_name}.h', 'w+') as f:
            f.write(header)

        source = kernel_source(kernel_ast)
        if launch_bounds:
            source = _add_launch_bound(source, launch_bounds)
        # TODO CUDA specific suffix
        with open(src_path / f'{kernel_name}.cu', 'w+') as f:
            f.write(source)

    with open(src_path / 'main.cu', 'w+') as f:
        f.write(kernel_main(kernel_asts, timing=timing))

    copy_static_files(path, source_file_suffix='.cu')
    compiler_toolchain(path, compiler, likwid=False)


def kernel_main(kernels_ast: List[KernelFunction], *, timing: bool = True):
    """
    Return C code of a benchmark program for the given kernel.

    Args:
        kernels_ast: A list of the pystencils AST object as returned by create_kernel for benchmarking
        timing: add timing output to the code, prints time per iteration to stdout
    Returns:
        C code as string
    """
    Kernel = namedtuple('Kernel', ['name', 'constants', 'fields', 'call_parameters',
                                   'call_argument_list', 'blocks', 'grid'])
    kernels = []
    includes = set()
    for kernel in kernels_ast:
        name = kernel.function_name
        accessed_fields = {f.name: f for f in kernel.fields_accessed}
        constants = []
        fields = []
        call_parameters = []
        block_and_thread_numbers = dict()
        for p in kernel.get_parameters():
            if not p.is_field_parameter:
                constants.append((p.symbol.name, str(p.symbol.dtype)))
                call_parameters.append(p.symbol.name)
            else:
                assert p.is_field_pointer, "Benchmark implemented only for kernels with fixed loop size"
                field = accessed_fields[p.field_name]
                dtype = str(get_base_type(p.symbol.dtype))
                elements = prod(field.shape)

                fields.append((p.field_name, dtype, elements))
                call_parameters.append(p.field_name)

            common_shape = get_common_field(kernel.fields_accessed).shape
            indexing = kernel.indexing
            block_and_thread_numbers = indexing.call_parameters(common_shape)
            block_and_thread_numbers['block'] = tuple(int(i) for i in block_and_thread_numbers['block'])
            block_and_thread_numbers['grid'] = tuple(int(i) for i in block_and_thread_numbers['grid'])

        kernels.append(Kernel(name=name, fields=fields, constants=constants, call_parameters=call_parameters,
                              call_argument_list=",".join(call_parameters),
                              blocks=block_and_thread_numbers['block'], grid=block_and_thread_numbers['grid']))
        includes.add(name)

    jinja_context = {
        'kernels': kernels,
        'includes': includes,
        'timing': timing,
    }

    main = _env.get_template('gpu/main.c').render(**jinja_context)
    return main


def kernel_header(kernel_ast: KernelFunction, dialect: Backend = Backend.C) -> str:
    return _kernel_header(kernel_ast,
                          dialect=dialect,
                          template_file='gpu/kernel.h',
                          additional_jinja_context={'target': 'gpu'})


def kernel_source(kernel_ast: KernelFunction, dialect: Backend = Backend.C) -> str:
    return _kernel_source(kernel_ast,
                          dialect=dialect,
                          template_file='gpu/kernel.cu',
                          additional_jinja_context={'target': 'gpu'})