Fix #6444: pruner issues with reference stealing functions

docs/source/developer/numba-runtime.rst

@@ -48,21 +48,39 @@ Optimizations
 -------------
 
 The compiler is allowed to emit incref/decref operations naively.  It relies
-on an optimization pass that to remove the redundant reference count
-operations.
+on an optimization pass to remove redundant reference count operations.
 
-The optimization pass runs on block level to avoid control flow analysis.
+A new optimization pass is implemented in version 0.52.0 to remove reference
+count operations that fall into the following four categories of control-flow
+structure---per basic-block, diamond, fanout, fanout+raise. See the documentation
+for :envvar:`NUMBA_LLVM_REFPRUNE_FLAGS` for their descriptions.
+
+The old optimization pass runs on block level to avoid control flow analysis.
 It depends on LLVM function optimization pass to simplify the control flow,
 stack-to-register, and simplify instructions.  It works by matching and
-removing incref and decref pairs within each block.
+removing incref and decref pairs within each block.  The old pass can be
+enabled by setting :envvar:`NUMBA_LLVM_REFPRUNE_PASS` to `0`.
+
+Important assumptions
+---------------------
+
+Both the old (pre-0.52.0) and the new (post-0.52.0) optimization passes assume
+that the only function that can consume a reference is ``NRT_decref``.
+It is important that there are no other functions that will consume references.
+Since the passes operate on LLVM IR, the "functions" here are referring to any
+callee in a LLVM call instruction.
+
+To summarize, all functions exposed to the refcount optimization pass
+**must not** consume counted references unless done so via ``NRT_decref``.
 
 
-Quirks
-------
+Quirks of legacy optimization pass
+----------------------------------
 
-Since the `refcount optimization pass <nrt-refct-opt-pass_>`_ requires LLVM
-function optimization pass, the pass works on the LLVM IR as text.  The
-optimized IR is then materialized again as a new LLVM in-memory bitcode object.
+Since the pre-0.52.0 `refcount optimization pass <nrt-refct-opt-pass_>`_
+requires LLVM function optimization pass, the pass works on the LLVM IR as
+text. The optimized IR is then materialized again as a new LLVM in-memory
+bitcode object.
 
 
 Debugging Leaks

numba/core/boxing.py

@@ -396,8 +396,9 @@ def box_array(typ, val, c):
     if c.context.enable_nrt:
         np_dtype = numpy_support.as_dtype(typ.dtype)
         dtypeptr = c.env_manager.read_const(c.env_manager.add_const(np_dtype))
-        # Steals NRT ref
         newary = c.pyapi.nrt_adapt_ndarray_to_python(typ, val, dtypeptr)
+        # Steals NRT ref
+        c.context.nrt.decref(c.builder, typ, val)
         return newary
     else:
         parent = nativeary.parent

numba/core/pythonapi.py

@@ -1189,7 +1189,7 @@ def nrt_adapt_ndarray_to_python(self, aryty, ary, dtypeptr):
         intty = ir.IntType(32)
         fnty = Type.function(self.pyobj,
                              [self.voidptr, intty, intty, self.pyobj])
-        fn = self._get_function(fnty, name="NRT_adapt_ndarray_to_python")
+        fn = self._get_function(fnty, name="NRT_adapt_ndarray_to_python_acqref")
         fn.args[0].add_attribute(lc.ATTR_NO_CAPTURE)
 
         ndim = self.context.get_constant(types.int32, aryty.ndim)

numba/core/runtime/_nrt_python.c

@@ -285,8 +285,13 @@ PyObject* try_to_return_parent(arystruct_t *arystruct, int ndim,
     return NULL;
 }
 
+/**
+ * This function was renamed in 0.52.0 to specify that it acquires references.
+ * It used to steal the reference of the arystruct.
+ * Refers to https://github.com/numba/numba/pull/6446
+ */
 NUMBA_EXPORT_FUNC(PyObject *)
-NRT_adapt_ndarray_to_python(arystruct_t* arystruct, int ndim,
+NRT_adapt_ndarray_to_python_acqref(arystruct_t* arystruct, int ndim,
                             int writeable, PyArray_Descr *descr)
 {
     PyArrayObject *array;
@@ -311,9 +316,6 @@ NRT_adapt_ndarray_to_python(arystruct_t* arystruct, int ndim,
     if (arystruct->parent) {
         PyObject *obj = try_to_return_parent(arystruct, ndim, descr);
         if (obj) {
-            /* Release NRT reference to the numpy array */
-            if (arystruct->meminfo)
-                NRT_MemInfo_release(arystruct->meminfo);
             return obj;
         }
     }
@@ -325,8 +327,9 @@ NRT_adapt_ndarray_to_python(arystruct_t* arystruct, int ndim,
         /* SETITEM steals reference */
         PyTuple_SET_ITEM(args, 0, PyLong_FromVoidPtr(arystruct->meminfo));
         /*  Note: MemInfo_init() does not incref.  This function steals the
-         *        NRT reference.
+         *        NRT reference, which we need to acquire.
          */
+        NRT_MemInfo_acquire(arystruct->meminfo);
         if (MemInfo_init(miobj, args, NULL)) {
             return NULL;
         }

numba/core/runtime/_nrt_pythonmod.c

@@ -154,7 +154,7 @@ build_c_helpers_dict(void)
 #define declmethod(func) _declpointer(#func, &NRT_##func)
 
 declmethod(adapt_ndarray_from_python);
-declmethod(adapt_ndarray_to_python);
+declmethod(adapt_ndarray_to_python_acqref);
 declmethod(adapt_buffer_from_python);
 declmethod(meminfo_new_from_pyobject);
 declmethod(meminfo_as_pyobject);