kerncraft_interface.py 7.09 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
from tempfile import TemporaryDirectory
Martin Bauer's avatar
Martin Bauer committed
2

Martin Bauer's avatar
Martin Bauer committed
3
import sympy as sp
Martin Bauer's avatar
Martin Bauer committed
4
import os
Martin Bauer's avatar
Martin Bauer committed
5
from collections import defaultdict
Martin Bauer's avatar
Martin Bauer committed
6
7
import subprocess
import kerncraft
Martin Bauer's avatar
Martin Bauer committed
8
import kerncraft.kernel
9
from kerncraft.iaca import iaca_analyse_instrumented_binary, iaca_instrumentation
Martin Bauer's avatar
Martin Bauer committed
10
from pystencils.kerncraft_coupling.generate_benchmark import generate_benchmark
Martin Bauer's avatar
Martin Bauer committed
11
from pystencils.astnodes import LoopOverCoordinate, SympyAssignment, ResolvedFieldAccess
Martin Bauer's avatar
Martin Bauer committed
12
from pystencils.field import get_layout_from_strides
Martin Bauer's avatar
Martin Bauer committed
13
from pystencils.sympyextensions import count_operations_in_ast
Martin Bauer's avatar
Martin Bauer committed
14
15
16
17
18
19
20
21
from pystencils.utils import DotDict


class PyStencilsKerncraftKernel(kerncraft.kernel.Kernel):
    """
    Implementation of kerncraft's kernel interface for pystencils CPU kernels.
    Analyses a list of equations assuming they will be executed on a CPU
    """
Martin Bauer's avatar
Martin Bauer committed
22
23
    LIKWID_BASE = '/usr/local/likwid'

Jan Hönig's avatar
Jan Hönig committed
24
25
    def __init__(self, ast, machine=None):
        super(PyStencilsKerncraftKernel, self).__init__(machine)
Martin Bauer's avatar
Martin Bauer committed
26
27

        self.ast = ast
Martin Bauer's avatar
Martin Bauer committed
28
        self.temporary_dir = TemporaryDirectory()
Martin Bauer's avatar
Martin Bauer committed
29
30

        # Loops
Martin Bauer's avatar
Martin Bauer committed
31
32
        inner_loops = [l for l in ast.atoms(LoopOverCoordinate) if l.is_innermost_loop]
        if len(inner_loops) == 0:
Martin Bauer's avatar
Martin Bauer committed
33
            raise ValueError("No loop found in pystencils AST")
Martin Bauer's avatar
Martin Bauer committed
34
        elif len(inner_loops) > 1:
Martin Bauer's avatar
Martin Bauer committed
35
36
            raise ValueError("pystencils AST contains multiple inner loops - only one can be analyzed")
        else:
Martin Bauer's avatar
Martin Bauer committed
37
            inner_loop = inner_loops[0]
Martin Bauer's avatar
Martin Bauer committed
38
39

        self._loop_stack = []
Martin Bauer's avatar
Martin Bauer committed
40
41
42
        cur_node = inner_loop
        while cur_node is not None:
            if isinstance(cur_node, LoopOverCoordinate):
Martin Bauer's avatar
Martin Bauer committed
43
44
45
                loop_counter_sym = cur_node.loop_counter_symbol
                loop_info = (loop_counter_sym.name, cur_node.start, cur_node.stop, cur_node.step)
                self._loop_stack.append(loop_info)
Martin Bauer's avatar
Martin Bauer committed
46
            cur_node = cur_node.parent
Martin Bauer's avatar
Martin Bauer committed
47
48
49
        self._loop_stack = list(reversed(self._loop_stack))

        # Data sources & destinations
50
51
        self.sources = defaultdict(list)
        self.destinations = defaultdict(list)
Martin Bauer's avatar
Martin Bauer committed
52

Martin Bauer's avatar
Martin Bauer committed
53
        reads, writes = search_resolved_field_accesses_in_ast(inner_loop)
Martin Bauer's avatar
Martin Bauer committed
54
        for accesses, target_dict in [(reads, self.sources), (writes, self.destinations)]:
Martin Bauer's avatar
Martin Bauer committed
55
            for fa in accesses:
Martin Bauer's avatar
Martin Bauer committed
56
                coord = [sp.Symbol(LoopOverCoordinate.get_loop_counter_name(i), positive=True, integer=True) + off
Martin Bauer's avatar
Martin Bauer committed
57
                         for i, off in enumerate(fa.offsets)]
Martin Bauer's avatar
Martin Bauer committed
58
                coord += list(fa.idx_coordinate_values)
Martin Bauer's avatar
Martin Bauer committed
59
60
                layout = get_layout_from_strides(fa.field.strides)
                permuted_coord = [coord[i] for i in layout]
Martin Bauer's avatar
Martin Bauer committed
61
                target_dict[fa.field.name].append(permuted_coord)
Martin Bauer's avatar
Martin Bauer committed
62
63

        # Variables (arrays)
Martin Bauer's avatar
Martin Bauer committed
64
65
66
67
68
        fields_accessed = ast.fields_accessed
        for field in fields_accessed:
            layout = get_layout_from_strides(field.strides)
            permuted_shape = list(field.shape[i] for i in layout)
            self.set_variable(field.name, str(field.dtype), tuple(permuted_shape))
Martin Bauer's avatar
Martin Bauer committed
69

70
71
72
73
        for param in ast.get_parameters():
            if not param.is_field_parameter:
                self.set_variable(param.symbol.name, str(param.symbol.dtype), None)
                self.sources[param.symbol.name] = [None]
Martin Bauer's avatar
Martin Bauer committed
74
75
76
77
78

        # data type
        self.datatype = list(self.variables.values())[0][0]

        # flops
Martin Bauer's avatar
Martin Bauer committed
79
        operation_count = count_operations_in_ast(inner_loop)
Martin Bauer's avatar
Martin Bauer committed
80
        self._flops = {
Martin Bauer's avatar
Martin Bauer committed
81
82
83
            '+': operation_count['adds'],
            '*': operation_count['muls'],
            '/': operation_count['divs'],
Martin Bauer's avatar
Martin Bauer committed
84
        }
Jan Hönig's avatar
Jan Hönig committed
85
86
        for k in [k for k, v in self._flops.items() if v == 0]:
            del self._flops[k]
Martin Bauer's avatar
Martin Bauer committed
87
88
        self.check()

Jan Hönig's avatar
Jan Hönig committed
89
90
91
    def iaca_analysis(self, micro_architecture, asm_block='auto',
                      pointer_increment='auto_with_manual_fallback', verbose=False):
        compiler, compiler_args = self._machine.get_compiler()
Martin Bauer's avatar
Martin Bauer committed
92
93
        if '-std=c99' not in compiler_args:
            compiler_args += ['-std=c99']
Martin Bauer's avatar
Martin Bauer committed
94
        header_path = kerncraft.get_header_path()
Martin Bauer's avatar
Martin Bauer committed
95

Martin Bauer's avatar
Martin Bauer committed
96
        compiler_cmd = [compiler] + compiler_args + ['-I' + header_path]
Martin Bauer's avatar
Martin Bauer committed
97

Martin Bauer's avatar
Martin Bauer committed
98
99
100
        src_file = os.path.join(self.temporary_dir.name, "source.c")
        asm_file = os.path.join(self.temporary_dir.name, "source.s")
        iaca_asm_file = os.path.join(self.temporary_dir.name, "source.iaca.s")
Martin Bauer's avatar
Martin Bauer committed
101
        dummy_src_file = os.path.join(header_path, "dummy.c")
Martin Bauer's avatar
Martin Bauer committed
102
103
        dummy_asm_file = os.path.join(self.temporary_dir.name, "dummy.s")
        binary_file = os.path.join(self.temporary_dir.name, "binary")
Martin Bauer's avatar
Martin Bauer committed
104
105

        # write source code to file
Martin Bauer's avatar
Martin Bauer committed
106
107
        with open(src_file, 'w') as f:
            f.write(generate_benchmark(self.ast, likwid=False))
Martin Bauer's avatar
Martin Bauer committed
108
109

        # compile to asm files
Martin Bauer's avatar
Martin Bauer committed
110
        subprocess.check_output(compiler_cmd + [src_file, '-S', '-o', asm_file])
Martin Bauer's avatar
Martin Bauer committed
111
        subprocess.check_output(compiler_cmd + [dummy_src_file, '-S', '-o', dummy_asm_file])
Martin Bauer's avatar
Martin Bauer committed
112

Martin Bauer's avatar
Martin Bauer committed
113
114
        with open(asm_file) as read, open(iaca_asm_file, 'w') as write:
            instrumented_asm_block = iaca_instrumentation(read, write)
Martin Bauer's avatar
Martin Bauer committed
115
116

        # assemble asm files to executable
Martin Bauer's avatar
Martin Bauer committed
117
        subprocess.check_output(compiler_cmd + [iaca_asm_file, dummy_asm_file, '-o', binary_file])
Martin Bauer's avatar
Martin Bauer committed
118

Martin Bauer's avatar
Martin Bauer committed
119
        result = iaca_analyse_instrumented_binary(binary_file, micro_architecture)
Martin Bauer's avatar
Martin Bauer committed
120
    
Martin Bauer's avatar
Martin Bauer committed
121
        return result, instrumented_asm_block
Martin Bauer's avatar
Martin Bauer committed
122

Jan Hönig's avatar
Jan Hönig committed
123
124
    def build(self, lflags=None, verbose=False, openmp=False):
        # TODO do we use openmp or not???
Jan Hönig's avatar
Jan Hönig committed
125
        compiler, compiler_args = self._machine.get_compiler()
Martin Bauer's avatar
Martin Bauer committed
126
127
        if '-std=c99' not in compiler_args:
            compiler_args.append('-std=c99')
Martin Bauer's avatar
Martin Bauer committed
128
        header_path = kerncraft.get_header_path()
Martin Bauer's avatar
Martin Bauer committed
129
130
131
132

        cmd = [compiler] + compiler_args + [
            '-I' + os.path.join(self.LIKWID_BASE, 'include'),
            '-L' + os.path.join(self.LIKWID_BASE, 'lib'),
Martin Bauer's avatar
Martin Bauer committed
133
            '-I' + header_path,
Martin Bauer's avatar
Martin Bauer committed
134
135
136
            '-Wl,-rpath=' + os.path.join(self.LIKWID_BASE, 'lib'),
        ]

Martin Bauer's avatar
Martin Bauer committed
137
        dummy_src_file = os.path.join(header_path, 'dummy.c')
Martin Bauer's avatar
Martin Bauer committed
138
139
        src_file = os.path.join(self.temporary_dir.name, "source_likwid.c")
        bin_file = os.path.join(self.temporary_dir.name, "benchmark")
Martin Bauer's avatar
Martin Bauer committed
140

Martin Bauer's avatar
Martin Bauer committed
141
142
        with open(src_file, 'w') as f:
            f.write(generate_benchmark(self.ast, likwid=True))
Martin Bauer's avatar
Martin Bauer committed
143

Martin Bauer's avatar
Martin Bauer committed
144
145
        subprocess.check_output(cmd + [src_file, dummy_src_file, '-pthread', '-llikwid', '-o', bin_file])
        return bin_file
146

Martin Bauer's avatar
Martin Bauer committed
147
148

class KerncraftParameters(DotDict):
Martin Bauer's avatar
Martin Bauer committed
149
150
    def __init__(self, **kwargs):
        super(KerncraftParameters, self).__init__(**kwargs)
Martin Bauer's avatar
Martin Bauer committed
151
152
153
154
155
        self['asm_block'] = 'auto'
        self['asm_increment'] = 0
        self['cores'] = 1
        self['cache_predictor'] = 'SIM'
        self['verbose'] = 0
Jan Hönig's avatar
Jan Hönig committed
156
        self['pointer_increment'] = 'auto'
Jan Hönig's avatar
Jan Hönig committed
157
158
        self['iterations'] = 10

Martin Bauer's avatar
Martin Bauer committed
159
160
161
162

# ------------------------------------------- Helper functions ---------------------------------------------------------


Martin Bauer's avatar
Martin Bauer committed
163
def search_resolved_field_accesses_in_ast(ast):
Martin Bauer's avatar
Martin Bauer committed
164
165
166
167
168
169
170
171
172
    def visit(node, reads, writes):
        if not isinstance(node, SympyAssignment):
            for a in node.args:
                visit(a, reads, writes)
            return

        for expr, accesses in [(node.lhs, writes), (node.rhs, reads)]:
            accesses.update(expr.atoms(ResolvedFieldAccess))

Martin Bauer's avatar
Martin Bauer committed
173
174
175
    read_accesses = set()
    write_accesses = set()
    visit(ast, read_accesses, write_accesses)
Martin Bauer's avatar
Martin Bauer committed
176
    return read_accesses, write_accesses