Repository: numba/numba

Issue:
Minsize argument for Dict.empty()
## Feature request

<!--

Please include details of the feature you would like to see, why you would
like to see it/the use case.

-->

Being able to allocate memory for a Dict based on problem-specific information would minimize resizes and increase performance. A Dict is initialized by a call to the function `numba_dict_new_minsize` in `dictobject.c`. The starting number of buckets is fixed  at `D_MINSIZE = 8`, which only allows for 5 entries without a resize.

The code comments state that this is suitable for the common case of a small dictionary used for passing keyword arguments. I believe an optional size argument would give users the ability to achieve significantly better performance with large hash tables. This would be in line with Numba's focus on fast numerical computation.

I've tried to modify Numba so that I could pass a size argument to `Dict.empty()` and have it call `numba_dict_new` instead of `numba_dict_new_minsize`. Unfortunately, I haven't been successful so far. Please consider adding such an option in a future release. Thank you.



Discussion:
Many thanks for the feature request, and for identifying the required changes. It sounds like you've made a good attempt at making the change so far - would you like to push your branch up somewhere and perhaps I can suggest how to complete the implementation?
Noting from the triage meeting that a similar request is also made for typed lists regularly.
Apologies for the delay in responding. I'm working on uploading the (unsuccessful) changes I made to a fork.

I believe typed lists have this feature under the `empty_list` method. The argument is called `allocated`.
Hi @gmarkall,

I've pushed my changes to a branch at https://github.com/stefanfed/numba/tree/dict_empty_allocated

This implementation of `Dict.empty()` works in both the object and nopython modes only when the optional argument `allocated` is set to the default `None` (or omitted). It appears to have the same behavior as before the changes. When an integer is provided, it shows the following error message and crashes the interpreter:

```
>>> Dict.empty(nb.int64, nb.int64, allocated=8)
LLVM ERROR: Symbol not found: numba_dict_new
```

I haven't figured out how these functions are linked to LLVM from `dictobject.h`. It seems `numba_dict_new` has been exported in the same way as `numba_dict_new_minsize` based on the preceding `NUMBA_EXPORT_FUNC(int)` macro. I'm not sure what that does exactly, to be honest. Likely, changes have to be made to another file to have LLVM recognize `numba_dict_new`. My understanding of the internal workings of numba with LLVM is lacking. That's my biggest problem.

The other problem is that `allocated` must be a power of 2 according to the notes in `dictobject.h`. I'm unsure what the proper behavior would be if a number that isn't a power of 2 is passed. I'd suggest rounding it up to the nearest power of 2. However, it could also be rounded down, or it could fail, leaving the check up to the user. 

I'm also wondering what the ideal method rounding to the nearest power of 2 would be and at what stage to do it. Perhaps it should be done in the C implementation?

I'd appreciate your help in completing this. Thanks.
Hi @stefanfed - many thanks for pushing your branch!

Regarding this problem:

>```
> >>> Dict.empty(nb.int64, nb.int64, allocated=8)
>LLVM ERROR: Symbol not found: numba_dict_new

The `numba_dict_new` function needs to be added to LLVM's symbol table. This is done using [`llvmlite.binding.add_symbol()`](https://llvmlite.readthedocs.io/en/latest/user-guide/binding/dynamic-libraries.html#llvmlite.binding.add_symbol). One place in Numba where a collection of these are added is in the [`_load_global_helpers()`](https://github.com/numba/numba/blob/2f89d45807b55c676646338bb4d730b69658d2c5/numba/core/base.py#L147) function:

https://github.com/numba/numba/blob/2f89d45807b55c676646338bb4d730b69658d2c5/numba/core/base.py#L155-L158

The `c_helpers` dicts of these modules are constructed using a function with a macro called `declmethod` (for `_helperlib`, for example):

https://github.com/numba/numba/blob/2f89d45807b55c676646338bb4d730b69658d2c5/numba/_helpermod.c#L28-L49

In order to make this work for the `numba_dict_new` method, I would suggest adding a similar construct that builds a helpers dict for typed dict methods containing the `numba_dict_new` method, and adding it to the list of `c_helpers` dicts that get initialized in `_load_global_helpers()`.

>  It seems `numba_dict_new` has been exported in the same way as `numba_dict_new_minsize` based on the preceding `NUMBA_EXPORT_FUNC(int)` macro. I'm not sure what that does exactly, to be honest.

Basically it just adds `__attribute__ ((visibility("hidden")))` on Unix-like platforms so that the symbols are defined but not exported by default (I believe this is already the default on Windows anyway).

> The other problem is that `allocated` must be a power of 2 according to the notes in `dictobject.h`. I'm unsure what the proper behavior would be if a number that isn't a power of 2 is passed. I'd suggest rounding it up to the nearest power of 2. However, it could also be rounded down, or it could fail, leaving the check up to the user.

I think rounding up and documenting somewhere that the behaviour is to round up would be acceptable.

> I'm also wondering what the ideal method rounding to the nearest power of 2 would be and at what stage to do it. Perhaps it should be done in the C implementation?

My guess would be that this would be the easiest place to ensure that only power-of-2 sizes are allocated - I think having no check / adjustment in the C side and relying on something in the Python side to round things up to power-of-2 sizes leaves open more opportunity for a future change to accidentally bypass the rounding-up, and I'd imagine it will also be easier to reason about the correctness of your changes at the point of review as well.

> I'd appreciate your help in completing this. Thanks.

Many thanks for your efforts so far! I hope we can get this to completion - please do let us know if we have provide further guidance.
Hi @gmarkall,

I was able to fix the LLVM error by adding a single line to `_helpermod.c`. Thanks for pointing me there.

I added a function for rounding up to the next power of 2 in `dictobject.c`. To ensure no resizes are performed, I round up from 1.5 times the requested size because the load factor is set to 2/3 (after which the dictionary is resized). So the new parameter `allocated` is the number of dictionary entries that can be added without requiring a resize.

The implementation appears to work correctly now. I tested it with various inputs where the number of keys to be added is known ahead of time. The performance gain is over 2x for some ideal inputs (like a monotonically increasing range of ints), and it grows in steps with the input size as expected. When the ints are shuffled, the ratio goes down but the amount of time saved appears to stay the same. I saw no performance regressions.

I have a concern about the way I round up to the next power of 2:
https://github.com/stefanfed/numba/blob/078c70148de5f10041cb11f57c4f5e8ecdecba33/numba/cext/dictobject.c#L282-L296

Because `Py_ssize_t` is signed, it overflows when the next size is 2^63 on 64-bit, for example. I really doubt anyone has anywhere near enough memory to get to there, so this won't be a problem on 64-bit (though it could cause a crash if someone tried it). However, it could be a substantial problem on 32-bit. I wasn't sure how to resolve this - I'd appreciate your advice.

I'm not sure why `Py_ssize_t` is used for the size parameter as opposed to `size_t`. The size can never be negative. Could I change that, or would it cause problems elsewhere? Also, I'll note that the method may not be the most efficient because it doesn't take advantage of certain instructions that exist in most modern desktop processors. This likely isn't a priority, but it could be optimized.

Another change I hesitated to make involves `numba_dict_new`:
https://github.com/stefanfed/numba/blob/078c70148de5f10041cb11f57c4f5e8ecdecba33/numba/cext/dictobject.c#L560-L585

Currently, the only functionality of `numba_dict_new_minsize` is to call `numba_dict_new` with `size=D_MINSIZE`. Since we can guarantee `D_MINSIZE` is a power of 2, we don't need to round it up. Therefore, in `numba_dict_new`, I only attempt to round up if `size != D_MINSIZE`. An alternative is to copy all the code from `numba_dict_new` to `numba_dict_new_minsize` without the rounding up part. Would you say that's a good idea?

Should I open a pull request? I noticed that you require unit tests. I'm not 100% sure on this part. What kind of tests would you like me to write?
On second thought, `numba_dict_new` needs to be changed because the size parameter takes on 2 different meanings. In the minsize case, it's literally the allocation size, while in other cases it's the number of keys that can be inserted without a resize. So in the single case where someone passes `allocated=8`, it won't behave as I intend it to. I'll update it with the changes shortly.

Resolving patch:
diff --git a/numba/typed/dictobject.py b/numba/typed/dictobject.py
--- a/numba/typed/dictobject.py
+++ b/numba/typed/dictobject.py
@@ -70,13 +70,16 @@ class Status(IntEnum):
     ERR_CMP_FAILED = -5
 
 
-def new_dict(key, value):
-    """Construct a new dict.
+def new_dict(key, value, n_keys=0):
+    """Construct a new dict with enough space for *n_keys* without a resize.
 
     Parameters
     ----------
     key, value : TypeRef
         Key type and value type of the new dict.
+    n_keys : int, default 0
+        The number of keys to insert without needing a resize.
+        A value of 0 creates a dict with minimum size.
     """
     # With JIT disabled, ignore all arguments and return a Python dict.
     return dict()
@@ -213,41 +216,44 @@ def _imp_dtor(context, module):
 
 
 @intrinsic
-def _dict_new_minsize(typingctx, keyty, valty):
-    """Wrap numba_dict_new_minsize.
+def _dict_new_sized(typingctx, n_keys, keyty, valty):
+    """Wrap numba_dict_new_sized.
 
-    Allocate a new dictionary object with the minimum capacity.
+    Allocate a new dictionary object with enough space to hold
+    *n_keys* keys without needing a resize.
 
     Parameters
     ----------
     keyty, valty: Type
         Type of the key and value, respectively.
-
+    n_keys: int
+        The number of keys to insert without needing a resize.
+        A value of 0 creates a dict with minimum size.
     """
     resty = types.voidptr
-    sig = resty(keyty, valty)
+    sig = resty(n_keys, keyty, valty)
 
     def codegen(context, builder, sig, args):
-        fnty = ir.FunctionType(
-            ll_status,
-            [ll_dict_type.as_pointer(), ll_ssize_t, ll_ssize_t],
-        )
-        fn = cgutils.get_or_insert_function(builder.module, fnty,
-                                            'numba_dict_new_minsize')
+        n_keys = builder.bitcast(args[0], ll_ssize_t)
+
         # Determine sizeof key and value types
         ll_key = context.get_data_type(keyty.instance_type)
         ll_val = context.get_data_type(valty.instance_type)
         sz_key = context.get_abi_sizeof(ll_key)
         sz_val = context.get_abi_sizeof(ll_val)
+
         refdp = cgutils.alloca_once(builder, ll_dict_type, zfill=True)
-        status = builder.call(
-            fn,
-            [refdp, ll_ssize_t(sz_key), ll_ssize_t(sz_val)],
-        )
-        _raise_if_error(
-            context, builder, status,
-            msg="Failed to allocate dictionary",
-        )
+
+        argtys = [ll_dict_type.as_pointer(), ll_ssize_t, ll_ssize_t, ll_ssize_t]
+        fnty = ir.FunctionType(ll_status, argtys)
+        fn = ir.Function(builder.module, fnty, 'numba_dict_new_sized')
+
+        args = [refdp, n_keys, ll_ssize_t(sz_key), ll_ssize_t(sz_val)]
+        status = builder.call(fn, args)
+
+        allocated_failed_msg = "Failed to allocate dictionary"
+        _raise_if_error(context, builder, status, msg=allocated_failed_msg)
+
         dp = builder.load(refdp)
         return dp
 
@@ -645,9 +651,11 @@ def codegen(context, builder, signature, args):
 
 
 @overload(new_dict)
-def impl_new_dict(key, value):
+def impl_new_dict(key, value, n_keys=0):
     """Creates a new dictionary with *key* and *value* as the type
-    of the dictionary key and value, respectively.
+    of the dictionary key and value, respectively. *n_keys* is the
+    number of keys to insert without requiring a resize, where a
+    value of 0 creates a dictionary with minimum size.
     """
     if any([
         not isinstance(key, Type),
@@ -657,8 +665,10 @@ def impl_new_dict(key, value):
 
     keyty, valty = key, value
 
-    def imp(key, value):
-        dp = _dict_new_minsize(keyty, valty)
+    def imp(key, value, n_keys=0):
+        if n_keys < 0:
+            raise RuntimeError("expecting *n_keys* to be >= 0")
+        dp = _dict_new_sized(n_keys, keyty, valty)
         _dict_set_method_table(dp, keyty, valty)
         d = _make_dict(keyty, valty, dp)
         return d
@@ -813,7 +823,7 @@ def impl(dct, key, default=None):
         elif ix < DKIX.EMPTY:
             raise AssertionError("internal dict error during lookup")
         else:
-            status = _dict_delitem(dct,hashed, ix)
+            status = _dict_delitem(dct, hashed, ix)
             if status != Status.OK:
                 raise AssertionError("internal dict error during delitem")
             return val
@@ -865,7 +875,7 @@ def impl_copy(d):
     key_type, val_type = d.key_type, d.value_type
 
     def impl(d):
-        newd = new_dict(key_type, val_type)
+        newd = new_dict(key_type, val_type, n_keys=len(d))
         for k, v in d.items():
             newd[k] = v
         return newd
@@ -1335,7 +1345,7 @@ def cast_LiteralStrKeyDict_LiteralStrKeyDict(context, builder, fromty, toty,
                                   toty.literal_value.items()):
         # these checks are just guards, typing should have picked up any
         # problems
-        if k1 != k2: # keys must be same
+        if k1 != k2:  # keys must be same
             msg = "LiteralDictionary keys are not the same {} != {}"
             raise LoweringError(msg.format(k1, k2))
         # values must be same ty
diff --git a/numba/typed/typeddict.py b/numba/typed/typeddict.py
--- a/numba/typed/typeddict.py
+++ b/numba/typed/typeddict.py
@@ -19,8 +19,9 @@
 
 
 @njit
-def _make_dict(keyty, valty):
-    return dictobject._as_meminfo(dictobject.new_dict(keyty, valty))
+def _make_dict(keyty, valty, n_keys=0):
+    return dictobject._as_meminfo(dictobject.new_dict(keyty, valty,
+                                                      n_keys=n_keys))
 
 
 @njit
@@ -84,21 +85,24 @@ class Dict(MutableMapping):
     Implements the MutableMapping interface.
     """
 
-    def __new__(cls, dcttype=None, meminfo=None):
+    def __new__(cls, dcttype=None, meminfo=None, n_keys=0):
         if config.DISABLE_JIT:
             return dict.__new__(dict)
         else:
             return object.__new__(cls)
 
     @classmethod
-    def empty(cls, key_type, value_type):
+    def empty(cls, key_type, value_type, n_keys=0):
         """Create a new empty Dict with *key_type* and *value_type*
         as the types for the keys and values of the dictionary respectively.
+
+        Optionally, allocate enough memory to hold *n_keys* without requiring
+        resizes. The default value of 0 returns a dict with minimum size.
         """
         if config.DISABLE_JIT:
             return dict()
         else:
-            return cls(dcttype=DictType(key_type, value_type))
+            return cls(dcttype=DictType(key_type, value_type), n_keys=n_keys)
 
     def __init__(self, **kwargs):
         """
@@ -117,14 +121,15 @@ def __init__(self, **kwargs):
         else:
             self._dict_type = None
 
-    def _parse_arg(self, dcttype, meminfo=None):
+    def _parse_arg(self, dcttype, meminfo=None, n_keys=0):
         if not isinstance(dcttype, DictType):
             raise TypeError('*dcttype* must be a DictType')
 
         if meminfo is not None:
             opaque = meminfo
         else:
-            opaque = _make_dict(dcttype.key_type, dcttype.value_type)
+            opaque = _make_dict(dcttype.key_type, dcttype.value_type,
+                                n_keys=n_keys)
         return dcttype, opaque
 
     @property
@@ -209,12 +214,12 @@ def copy(self):
 
 
 @overload_classmethod(types.DictType, 'empty')
-def typeddict_empty(cls, key_type, value_type):
+def typeddict_empty(cls, key_type, value_type, n_keys=0):
     if cls.instance_type is not DictType:
         return
 
-    def impl(cls, key_type, value_type):
-        return dictobject.new_dict(key_type, value_type)
+    def impl(cls, key_type, value_type, n_keys=0):
+        return dictobject.new_dict(key_type, value_type, n_keys=n_keys)
 
     return impl
 
@@ -280,7 +285,7 @@ def convert(mi, typ):
             sig = signature(typ, *argtypes)
             nil_typeref = context.get_constant_null(argtypes[1])
             args = (mi, nil_typeref)
-            is_error, dctobj = c.pyapi.call_jit_code(convert , sig, args)
+            is_error, dctobj = c.pyapi.call_jit_code(convert, sig, args)
             # decref here because we are stealing a reference.
             c.context.nrt.decref(c.builder, typ, dctobj)
 


Test patch:
diff --git a/numba/tests/test_dictimpl.py b/numba/tests/test_dictimpl.py
--- a/numba/tests/test_dictimpl.py
+++ b/numba/tests/test_dictimpl.py
@@ -84,8 +84,8 @@ def popitem(self):
 
     def dict_new_minsize(self, key_size, val_size):
         dp = ctypes.c_void_p()
-        status = self.tc.numba_dict_new_minsize(
-            ctypes.byref(dp), key_size, val_size,
+        status = self.tc.numba_dict_new_sized(
+            ctypes.byref(dp), 0, key_size, val_size,
         )
         self.tc.assertEqual(status, 0)
         return dp
@@ -224,16 +224,18 @@ def wrap(name, restype, argtypes=()):
             'test_dict',
             ctypes.c_int,
         )
-        # numba_dict_new_minsize(
+        # numba_dict_new_sized(
         #    NB_Dict **out,
+        #    Py_ssize_t n_keys,
         #    Py_ssize_t key_size,
         #    Py_ssize_t val_size
         # )
-        self.numba_dict_new_minsize = wrap(
-            'dict_new_minsize',
+        self.numba_dict_new_sized = wrap(
+            'dict_new_sized',
             ctypes.c_int,
             [
                 ctypes.POINTER(dict_t),  # out
+                ctypes.c_ssize_t,        # n_keys
                 ctypes.c_ssize_t,        # key_size
                 ctypes.c_ssize_t,        # val_size
             ],