diff --git a/pystencils/kerncraft_coupling/kerncraft_interface.py b/pystencils/kerncraft_coupling/kerncraft_interface.py index 0aa4dcd24185cbb6064fdc3a1023178b30fea9af..b58ccb96a28af75994837183b400807fbde9acb1 100644 --- a/pystencils/kerncraft_coupling/kerncraft_interface.py +++ b/pystencils/kerncraft_coupling/kerncraft_interface.py @@ -3,7 +3,7 @@ from tempfile import TemporaryDirectory import sympy as sp from collections import defaultdict import kerncraft -import kerncraft.kernel +from kerncraft.kerncraft import KernelCode from typing import Optional from kerncraft.machinemodel import MachineModel @@ -16,7 +16,7 @@ from pystencils.utils import DotDict import warnings -class PyStencilsKerncraftKernel(kerncraft.kernel.KernelCode): +class PyStencilsKerncraftKernel(KernelCode): """ Implementation of kerncraft's kernel interface for pystencils CPU kernels. Analyses a list of equations assuming they will be executed on a CPU @@ -42,6 +42,7 @@ class PyStencilsKerncraftKernel(kerncraft.kernel.KernelCode): self.kernel_ast = ast self.temporary_dir = TemporaryDirectory() + self._keep_intermediates = debug_print # Loops inner_loops = [l for l in filtered_tree_iteration(ast, LoopOverCoordinate, stop_type=SympyAssignment) @@ -127,14 +128,24 @@ class PyStencilsKerncraftKernel(kerncraft.kernel.KernelCode): print("----------------------------- FLOPS -------------------------------") pprint(self._flops) - def as_code(self, type_='iaca', openmp=False): + def as_code(self, type_='iaca', openmp=False, as_filename=False): """ Generate and return compilable source code. - :param type: can be iaca or likwid. - :param openmp: if true, openmp code will be generated + Args: + type_: can be iaca or likwid. + openmp: if true, openmp code will be generated + as_filename: """ - return generate_benchmark(self.kernel_ast, likwid=type_ == 'likwid', openmp=openmp) + code = generate_benchmark(self.kernel_ast, likwid=type_ == 'likwid', openmp=openmp) + if as_filename: + fp, already_available = self._get_intermediate_file('kernel_{}.c'.format(type_), + machine_and_compiler_dependent=False) + if not already_available: + fp.write(code) + return fp.name + else: + return code class KerncraftParameters(DotDict): diff --git a/pystencils_tests/test_kerncraft_coupling.py b/pystencils_tests/test_kerncraft_coupling.py index 1ff9ab04c70d089a20f0aef0adc4802abcb78418..5d3485e4ab3f26e5d8d728d30ec06dcbaa5a4d2a 100644 --- a/pystencils_tests/test_kerncraft_coupling.py +++ b/pystencils_tests/test_kerncraft_coupling.py @@ -18,9 +18,12 @@ INPUT_FOLDER = os.path.join(SCRIPT_FOLDER, "kerncraft_inputs") @pytest.mark.kernkraft def test_compilation(): + machine_file_path = os.path.join(INPUT_FOLDER, "default_machine_file.yaml") + machine = MachineModel(path_to_yaml=machine_file_path) + kernel_file_path = os.path.join(INPUT_FOLDER, "2d-5pt.c") with open(kernel_file_path) as kernel_file: - reference_kernel = KernelCode(kernel_file.read(), machine=None, filename=kernel_file_path) + reference_kernel = KernelCode(kernel_file.read(), machine=machine, filename=kernel_file_path) reference_kernel.as_code('likwid') size = [30, 50, 3]