I have defined StructRef type S with a single member x being 1d array of integers.
As we know, in the NRT context, arrays have their own MemInfo structure associated with it that counts references to the memory region managed by the array object.
Accordingly, creating array in a jitted scope wraps it in a MemInfo that starts off with the reference count of 1. We can observe that happening, see the reproducer below.
Creating a structure S instance with the array a as its member x should subsequently bump up the reference count of the a array’s MemInfo to 2. The reproducer below agrees with this expectation as well.
However, mutating the same structure by re-assigning a different array b to its member does not appear to result in a unit drop of the reference count of a. This appears to be counter-intuitive. (The rest of the behavior - in regards to the reference counts of b - agrees with expectations.)
Reproducer (two auxiliary utilities are needed, feel free to skip straight to demo on the first read):
from numba import njit
from numba.core import types
from numba.experimental import structref
from numba.extending import intrinsic
from numpy import array, int64
@intrinsic
def _deref_int64_intp(typingctx, p_int_ty):
"""
Auxiliary utility # 1.
From address `p` given as int, read `int64`.
"""
def codegen(context, builder, signature, args):
p_ty_ll = context.get_value_type(p_int_ty).as_pointer()
ptr = builder.inttoptr(args[0], p_ty_ll)
return builder.load(ptr)
return types.int64(types.intp), codegen
@intrinsic
def _structref_meminfo(typingctx, s_type):
"""
Auxiliary utility # 2.
Get meminfo for the NRT-managed data `s`, its layout begins with `int64`-valued reference count.
"""
def codegen(context, builder, signature, args):
s = args[0]
meminfo = context.nrt.get_meminfos(builder, s_type, s)[0]
type_data, meminfo_p = meminfo
meminfo_p_as_int = builder.ptrtoint(meminfo_p, context.get_data_type(types.intp))
return meminfo_p_as_int
return types.intp(s_type), codegen
@structref.register
class STypeClass(types.StructRef):
pass
class S(structref.StructRefProxy):
def __new__(cls, x):
raise NotImplementedError("not intended to be called from python")
@property
@njit
def x(self):
return self.x
structref.define_proxy(S, STypeClass, ["x"])
@njit
def demo():
""" This is jitted because we want `a` and `b` to get wrapped in `MemInfo` - which requires NRT """
a_ = array([34, 56], dtype=int64)
a_meminfo_p = _structref_meminfo(a_)
a_ref_ct_1_ = _deref_int64_intp(a_meminfo_p)
s_ = S(a_)
a_ref_ct_2_ = _deref_int64_intp(a_meminfo_p)
b_ = array([-14, 17], dtype=int64)
b_meminfo_p = _structref_meminfo(b_)
b_ref_ct_1_ = _deref_int64_intp(b_meminfo_p)
s_.x = b_
a_ref_ct_3_ = _deref_int64_intp(a_meminfo_p)
b_ref_ct_2_ = _deref_int64_intp(b_meminfo_p)
return a_, b_, s_, a_ref_ct_1_, a_ref_ct_2_, a_ref_ct_3_, b_ref_ct_1_, b_ref_ct_2_
if __name__ == "__main__":
a, b, s, a_ref_ct_1, a_ref_ct_2, a_ref_ct_3, b_ref_ct_1, b_ref_ct_2 = demo()
print(f"a_ref_ct_1 = {a_ref_ct_1}") # 1 OK, array `a_` created
print(f"a_ref_ct_2 = {a_ref_ct_2}") # 2 OK, `s_` now references `a_`'s payload too
print(f"a_ref_ct_3 = {a_ref_ct_3}") # 2 Why? Didn't it get evicted from `s_`?
print(f"b_ref_ct_1 = {b_ref_ct_1}") # 1 OK, array `b_` created
print(f"b_ref_ct_2 = {b_ref_ct_2}") # 2 OK, `s_` now references `b_`'s payload too
print(f"s.x = {s.x}") # [-14, 17], i.e., `b`