added max_kernel_size argument to core driver

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, "<artiq>")
     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)

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)