From 2f9c802b3143d7cc6817a0fe0c824e6eec54286f Mon Sep 17 00:00:00 2001 From: Leon Riesebos Date: Fri, 12 Feb 2021 13:19:47 -0500 Subject: [PATCH] added max_kernel_size argument to core driver Signed-off-by: Leon Riesebos --- artiq/coredevice/core.py | 75 +++++++++++++++++++++++- artiq/test/coredevice/test_core.py | 92 ++++++++++++++++++++++++++++++ 2 files changed, 166 insertions(+), 1 deletion(-) create mode 100644 artiq/test/coredevice/test_core.py diff --git a/artiq/coredevice/core.py b/artiq/coredevice/core.py index d150df5962..fb5582b886 100644 --- a/artiq/coredevice/core.py +++ b/artiq/coredevice/core.py @@ -1,5 +1,6 @@ import os, sys import numpy +import logging from pythonparser import diagnostic @@ -18,17 +19,70 @@ from artiq.coredevice import exceptions +logger = logging.getLogger(__name__) +"""The logger for this file.""" + +_unit_to_bytes = { + 'b': 1000 ** 0, + 'kb': 1000 ** 1, + 'mb': 1000 ** 2, + 'gb': 1000 ** 3, + 'kib': 1024 ** 1, + 'mib': 1024 ** 2, + 'gib': 1024 ** 3, +} +"""Dict to map a string data size unit to bytes.""" + + +def _str_to_bytes(string): + """Convert a string to bytes.""" + assert isinstance(string, str) or string is None, 'Input must be of type str or None' + + if string: + # Format string + string = string.strip().lower() + + try: + # Extract the unit from the string + unit_str = '' + for prefix in (string[-i:] for i in [3, 2, 1] if len(string) > i): + if prefix in _unit_to_bytes: + unit_str = prefix + break + # Convert the value and the unit + unit = _unit_to_bytes[unit_str] # Do the unit first to fail early + value = int(string[:-len(unit_str)]) + # Return the product + return value * unit + except KeyError: + raise ValueError('No valid data size unit was found') + except ValueError: + raise ValueError('No valid integer was found') + else: + # Return None in case the string is None or empty + return None + + def _render_diagnostic(diagnostic, colored): def shorten_path(path): return path.replace(artiq_dir, "") lines = [shorten_path(path) for path in diagnostic.render(colored=colored)] return "\n".join(lines) + colors_supported = os.name == "posix" + + class _DiagnosticEngine(diagnostic.Engine): def render_diagnostic(self, diagnostic): sys.stderr.write(_render_diagnostic(diagnostic, colored=colors_supported) + "\n") + +class KernelSizeException(Exception): + """Raised if the maximum kernel size is exceeded.""" + pass + + class CompileError(Exception): def __init__(self, diagnostic): self.diagnostic = diagnostic @@ -65,13 +119,15 @@ class Core: :param ref_multiplier: ratio between the RTIO fine timestamp frequency and the RTIO coarse timestamp frequency (e.g. SERDES multiplication factor). + :param max_kernel_size: Maximum kernel size (e.g. ``"256 MiB"``, ``"512 kB"``). """ kernel_invariants = { "core", "ref_period", "coarse_ref_period", "ref_multiplier", } - def __init__(self, dmgr, host, ref_period, ref_multiplier=8, target="or1k"): + def __init__(self, dmgr, host, ref_period, ref_multiplier=8, target="or1k", + max_kernel_size=None): self.ref_period = ref_period self.ref_multiplier = ref_multiplier if target == "or1k": @@ -86,6 +142,9 @@ def __init__(self, dmgr, host, ref_period, ref_multiplier=8, target="or1k"): else: self.comm = CommKernel(host) + self._max_kernel_size = _str_to_bytes(max_kernel_size) + assert self._max_kernel_size is None or self._max_kernel_size >= 0 + self.first_run = True self.dmgr = dmgr self.core = self @@ -112,6 +171,20 @@ def compile(self, function, args, kwargs, set_result=None, library = target.compile_and_link([module]) stripped_library = target.strip(library) + # Obtain kernel size in bytes + kernel_size = len(stripped_library) + + if self._max_kernel_size is None: + # Report kernel size + logger.debug('Kernel size: %d bytes', kernel_size) + else: + # Check kernel size + logger.debug('Kernel size: %d/%d bytes (%.2f%%)', + kernel_size, self._max_kernel_size, + kernel_size / self._max_kernel_size) + if kernel_size > self._max_kernel_size: + raise KernelSizeException('Kernel too large') + return stitcher.embedding_map, stripped_library, \ lambda addresses: target.symbolize(library, addresses), \ lambda symbols: target.demangle(symbols) diff --git a/artiq/test/coredevice/test_core.py b/artiq/test/coredevice/test_core.py new file mode 100644 index 0000000000..6bc027a458 --- /dev/null +++ b/artiq/test/coredevice/test_core.py @@ -0,0 +1,92 @@ +import unittest +import logging +import numpy as np + +from artiq.experiment import kernel + +import artiq.coredevice.core + + +class _Core(artiq.coredevice.core.Core): + def __init__(self, *args, **kwargs): + super(_Core, self).__init__(*args, **kwargs) + self.dmgr['core'] = self + self.kernel_size = 0 + + def compile(self, *args, **kwargs): + embedding_map, kernel_library, symbolizer, demangler = \ + super(_Core, self).compile(*args, **kwargs) + self.kernel_size = len(kernel_library) + return embedding_map, kernel_library, symbolizer, demangler + + +class CoreTestCase(unittest.TestCase): + + def test_string_to_bytes(self): + data = [ + ('0b', 0), + ('15b', 15), + ('1kb', 1 * 1000 ** 1), + ('1 kb', 1 * 1000 ** 1), + (' 1 kb ', 1 * 1000 ** 1), + ('1KB', 1 * 1000 ** 1), + ('1kB', 1 * 1000 ** 1), + ('2kb', 2 * 1000 ** 1), + ('323 kb', 323 * 1000 ** 1), + ('3 mb', 3 * 1000 ** 2), + ('77gb', 77 * 1000 ** 3), + ('7kib', 7 * 1024 ** 1), + ('7mib', 7 * 1024 ** 2), + ('7gib', 7 * 1024 ** 3), + ('7 gib', 7 * 1024 ** 3), + ('7 GiB', 7 * 1024 ** 3), + ] + + for s, r in data: + self.assertEqual(artiq.coredevice.core._str_to_bytes(s), r) + + +class CoreCompilingTestCase(unittest.TestCase): + + def _set_core(self, **kwargs): + kwargs.setdefault('host', None) + kwargs.setdefault('ref_period', 1e-9) + self._update_kernel_invariants('core', clear=True) + self.core = _Core({}, **kwargs) + + def _update_kernel_invariants(self, *args: str, clear: bool = False) -> None: + kernel_invariants = set() if clear else getattr(self, 'kernel_invariants', set()) + self.kernel_invariants = kernel_invariants | set(args) + + def test_max_kernel_size(self): + # Calculate reference size + self._set_core() + self.data = np.empty(1, dtype=np.int32) + self._update_kernel_invariants('data') + self._kernel() + ref_size = self.core.kernel_size - 4 # Minus the one element used for testing the size + + sizes = [100, 256, 300, 512] + + for size in sizes: + self._set_core(max_kernel_size='{}b'.format(size * 4 + ref_size)) + self.data = np.empty(size, dtype=np.int32) + self._update_kernel_invariants('data') + self._kernel() # Fits exactly, does not raise + with self.assertRaises(artiq.coredevice.core.KernelSizeException, + msg='Kernel size exception did not raise'): + self.data = np.empty(size + 1, dtype=np.int32) + self._kernel() # One byte too big, should raise + + def test_log(self): + self._set_core() + self.data = np.empty(1, dtype=np.int32) + self._update_kernel_invariants('data') + with self.assertLogs(artiq.coredevice.core.logger, logging.DEBUG): + self._kernel() + + @kernel + def _kernel(self): + # Address `data` to get it compiled + for d in self.data: + print(d)