Possible config object implementation

[andrewfullard]

Courtesy Numba developers. https://hackmd.io/Wmy-e2RsQCSqjJsP0VAmbg?both

from numba import jit, types
from numba.extending import overload
import functools


# Some sort of configuration object
class Config():
   def __init__(self, a, b):
       self._a = a
       self._b = b

   @property
   def a(self):
       return self._a

   @property
   def b(self):
       return self._b


# Perhaps use a cache so that the identical Config instances return the same
# jit function? This will prevent recompilation of the entry point for two
# identical config instances as the jit function passed as the argument will be
# the same.
@functools.cache
def obj2strkeydict(obj, config_name):

   # unpack object to freevars and close over them
   tmp_a = obj.a
   tmp_b = obj.b
   assert isinstance(config_name, str)
   tmp_force_heterogeneous = config_name

   @jit
   def configurator():
       d = {'a': tmp_a,
            'b': tmp_b,
            'config_name': tmp_force_heterogeneous}
       return d

   # return a configuration function that returns a string-key-dict
   # representation of the configuration object.
   return configurator


# Define some "work"
def work(pred):
   # ... elided, put python implementation here
   pass


@overload(work)
def ol_work(pred):
   assert isinstance(pred, types.Literal)
   print('Calling work with type', pred)
   return lambda pred: pred


@jit
def physics(cfig_func):
   # This is the main entry point to the application, it takes a configuration
   # function as an argument. It will specialise on each configuration
   # function.

   # call the function, config is a string-key-dict
   config = cfig_func()
   # unpack config, these types will be preserved as literals.
   a = config['a']
   b = config['b']
   # call some work to check the types.
   return work(b) + a


# demo
def demo():

   # Create two different Python based configuration objects with literal
   # entries.
   configuration1 = Config(10, True)
   configuration2 = Config(12.3, False)

   # Create corresponding converted configuration objects...
   jit_config1 = obj2strkeydict(configuration1, 'config1')
   jit_config2 = obj2strkeydict(configuration2, 'config2')
   # create "another" configuration1 instance, memoization prevents
   # duplication.
   jit_config1_again = obj2strkeydict(configuration1, 'config1')

   # Call the `physics` application, it will specialize on config.
   physics(jit_config1)
   physics(jit_config2)
   # should not trigger a 3rd compilation, config is memoized.
   physics(jit_config1_again)
   physics.inspect_types()


if __name__ == "__main__":
   demo()