WIP: Emit a warning in case of possible memory leaks

.github/workflows/rbc_test.yml

@@ -35,7 +35,7 @@ jobs:
     name: ${{ matrix.os }} - Python v${{ matrix.python-version }} - Numba v${{ matrix.numba-version }}
     runs-on: ${{ matrix.os }}
     strategy:
-      fail-fast: true
+      #fail-fast: true
       matrix:
         os: [ubuntu-latest, windows-latest, macos-latest]
         python-version: ['3.9', '3.8', '3.7']
@@ -104,7 +104,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     timeout-minutes: 20
     strategy:
-      fail-fast: true
+      #fail-fast: true
       matrix:
         os: [ubuntu-latest]
         python-version: ['3.9', '3.8', '3.7']

rbc/irtools.py

@@ -263,19 +263,25 @@ def compile_instance(func, sig,
                      target_context,
                      pipeline_class,
                      main_library,
-                     debug=False):
+                     debug=False,
+                     on_missing_free='warn'):
     """Compile a function with given signature. Return function name when
     succesful.
     """
+    # individual functions can override the global settig
+    if hasattr(func, '__on_missing_free__'):
+        on_missing_free = func.__on_missing_free__
     flags = compiler.Flags()
     if get_version('numba') >= (0, 54):
         flags.no_compile = True
         flags.no_cpython_wrapper = True
         flags.no_cfunc_wrapper = True
+        flags.on_missing_free = on_missing_free
     else:
         flags.set('no_compile')
         flags.set('no_cpython_wrapper')
         flags.set('no_cfunc_wrapper')
+        flags.set('on_missing_free', on_missing_free)
 
     fname = func.__name__ + sig.mangling()
     args, return_type = sigutils.normalize_signature(
@@ -347,7 +353,8 @@ def compile_to_LLVM(functions_and_signatures,
                     target_info: TargetInfo,
                     pipeline_class=compiler.Compiler,
                     user_defined_llvm_ir=None,
-                    debug=False):
+                    debug=False,
+                    on_missing_free='warn'):
     """Compile functions with given signatures to target specific LLVM IR.
 
     Parameters
@@ -360,6 +367,8 @@ def compile_to_LLVM(functions_and_signatures,
       Specify user-defined LLVM IR module that is linked in to the
       returned module.
     debug : bool
+    on_missing_free: str
+      'warn', 'fail' or 'ignore'.
 
     Returns
     -------
@@ -393,7 +402,8 @@ def compile_to_LLVM(functions_and_signatures,
                 fname = compile_instance(func, sig, target_info, typing_context,
                                          target_context, pipeline_class,
                                          main_library,
-                                         debug=debug)
+                                         debug=debug,
+                                         on_missing_free=on_missing_free)
                 if fname is not None:
                     succesful_fids.append(fid)
                     function_names.append(fname)

rbc/omnisci_backend/__init__.py

@@ -11,3 +11,4 @@
 from .python_operators import *  # noqa: F401, F403
 from .omnisci_column_list import *  # noqa: F401, F403
 from .omnisci_table_function_manager import *  # noqa: F401, F403
+from .omnisci_buffer import BufferMeta  # noqa: F401

rbc/omnisci_backend/omnisci_buffer.py

@@ -193,54 +193,6 @@ def omnisci_buffer_constructor(context, builder, sig, args):
     return fa._getpointer()
 
 
-@extending.intrinsic
-def free_all_other_buffers(typingctx, value_to_keep_alive):
-    """
-    Black magic function which automatically frees all the buffers which were
-    allocated in the function apart the given one.
-
-    value_to_keep_alive can be of any type:
-      - if it's of a Buffer type, it will be kept alive and not freed
-      - if it's of any other type, all buffers will be freed unconditionally
-
-    The end user should never call this function explicitly: it is
-    automatically inserted by omnisci_pipeline.AutoFreeBuffers.
-    """
-
-    sig = types.void(value_to_keep_alive)
-
-    def codegen(context, builder, signature, args):
-        buffers = builder_buffers[builder]
-
-        # TODO: using stdlib `free` that works only for CPU. For CUDA
-        # devices, we need to use omniscidb provided deallocator.
-        target_info = TargetInfo()
-        try:
-            free_buffer_fn_name = target_info.info['fn_free_buffer']
-        except KeyError as msg:
-            raise UnsupportedError(f'{target_info} does not provide {msg}')
-        free_buffer_fnty = llvm_ir.FunctionType(void_t, [int8_t.as_pointer()])
-        free_buffer_fn = irutils.get_or_insert_function(
-            builder.module, free_buffer_fnty, free_buffer_fn_name)
-
-        if isinstance(value_to_keep_alive, BufferPointer):
-            # free all the buffers apart value_to_keep_alive
-            [keep_alive] = args
-            keep_alive_ptr = builder.load(builder.gep(keep_alive, [int32_t(0), int32_t(0)]))
-            keep_alive_ptr = builder.bitcast(keep_alive_ptr, int8_t.as_pointer())
-            for ptr8 in buffers:
-                with builder.if_then(builder.icmp_signed('!=', keep_alive_ptr, ptr8)):
-                    builder.call(free_buffer_fn, [ptr8])
-        else:
-            # free all the buffers unconditionally
-            for ptr8 in buffers:
-                builder.call(free_buffer_fn, [ptr8])
-
-        del builder_buffers[builder]
-
-    return sig, codegen
-
-
 @extending.intrinsic
 def free_buffer(typingctx, buf):
     """

rbc/omnisci_backend/omnisci_pipeline.py

@@ -1,73 +1,18 @@
 import operator
+import warnings
 
 from rbc.errors import NumbaTypeError
-from .omnisci_buffer import BufferMeta, free_all_other_buffers
+from .omnisci_buffer import BufferPointer, free_buffer
 from numba.core import ir, types
 from numba.core.compiler import CompilerBase, DefaultPassBuilder
 from numba.core.compiler_machinery import FunctionPass, register_pass
 from numba.core.untyped_passes import (IRProcessing,
                                        RewriteSemanticConstants,
                                        ReconstructSSA,
                                        DeadBranchPrune,)
-from numba.core.typed_passes import PartialTypeInference, DeadCodeElimination
-
-
-# Register this pass with the compiler framework, declare that it will not
-# mutate the control flow graph and that it is not an analysis_only pass (it
-# potentially mutates the IR).
-@register_pass(mutates_CFG=False, analysis_only=False)
-class AutoFreeBuffers(FunctionPass):
-    """
-    Black magic at work.
-
-    The goal of this pass is to "automagically" free all the buffers which
-    were allocated, apart the one which is used as a return value (if any).
-
-    NOTE: at the moment of writing there are very few tests for this and it's
-    likely that it is broken and/or does not work properly in the general
-    case. [Remove this note once we are confident that it works well]
-    """
-    _name = "auto_free_buffers"  # the common name for the pass
-
-    def __init__(self):
-        FunctionPass.__init__(self)
-
-    # implement method to do the work, "state" is the internal compiler
-    # state from the CompilerBase instance.
-    def run_pass(self, state):
-        func_ir = state.func_ir  # get the FunctionIR object
-
-        for blk in func_ir.blocks.values():
-            for stmt in blk.find_insts(ir.Assign):
-                if (
-                    isinstance(stmt.value, ir.FreeVar)
-                    and stmt.value.name in BufferMeta.class_names
-                ):
-                    break
-            else:
-                continue
-            break
-        else:
-            return False  # one does not changes the IR
-
-        for blk in func_ir.blocks.values():
-            loc = blk.loc
-            scope = blk.scope
-            for ret in blk.find_insts(ir.Return):
-
-                name = "free_all_other_buffers_fn"
-                value = ir.Global(name, free_all_other_buffers, loc)
-                target = scope.make_temp(loc)
-                stmt = ir.Assign(value, target, loc)
-                blk.insert_before_terminator(stmt)
-
-                fn_call = ir.Expr.call(func=target, args=[ret.value], kws=(), loc=loc)
-                lhs = scope.make_temp(loc)
-                var = ir.Assign(fn_call, lhs, blk.loc)
-                blk.insert_before_terminator(var)
-                break
-
-        return True  # we changed the IR
+from numba.core.typed_passes import (PartialTypeInference,
+                                     DeadCodeElimination,
+                                     NopythonTypeInference)
 
 
 @register_pass(mutates_CFG=False, analysis_only=False)
@@ -151,16 +96,107 @@ def run_pass(self, state):
         return mutated
 
 
+class MissingFreeWarning(Warning):
+
+    _msg = """
+    Possible memory leak detected in function `{}` defined in {}#{}!
+
+    In RBC, arrays and buffers must be freed manually by calling the method
+    .free() or the function free_buffer(): see the relevant docs.
+    """
+
+    def __init__(self, functionname, filename, firstlineno):
+        msg = self.make_message(functionname, filename, firstlineno)
+        Warning.__init__(self, msg)
+
+    @classmethod
+    def make_message(cls, functionname, filename, firstlineno):
+        return cls._msg.format(functionname, filename, firstlineno)
+
+
+class MissingFreeError(Exception):
+
+    def __init__(self, functionname, filename, firstlineno):
+        msg = MissingFreeWarning.make_message(functionname, filename, firstlineno)
+        Exception.__init__(self, msg)
+
+
+@register_pass(mutates_CFG=False, analysis_only=True)
+class DetectMissingFree(FunctionPass):
+    _name = "DetectMissingFree"
+
+    def __init__(self):
+        FunctionPass.__init__(self)
+
+    def iter_calls(self, func_ir):
+        for block in func_ir.blocks.values():
+            for inst in block.find_insts(ir.Assign):
+                if isinstance(inst.value, ir.Expr) and inst.value.op == 'call':
+                    yield inst
+
+    def contains_buffer_constructors(self, state):
+        """
+        Check whether the func_ir contains any call which creates a buffer. This
+        could be either a direct call to e.g. xp.Array() or a call to any
+        function which returns a BufferPointer: in that case we assume that
+        the ownership is transfered upon return and that the caller is
+        responsible to free() it.
+        """
+        func_ir = state.func_ir
+        for inst in self.iter_calls(func_ir):
+            ret_type = state.typemap[inst.target.name]
+            if isinstance(ret_type, BufferPointer):
+                return True
+        return False
+
+    def is_free_buffer(self, rhs):
+        return isinstance(rhs, (ir.Global, ir.FreeVar)) and rhs.value is free_buffer
+
+    def is_BufferPoint_dot_free(self, state, expr):
+        return (isinstance(expr, ir.Expr) and
+                expr.op == 'getattr' and
+                isinstance(state.typemap[expr.value.name], BufferPointer) and
+                expr.attr == 'free')
+
+    def contains_calls_to_free(self, state):
+        """
+        Check whether there is at least an instruction which calls free_buffer()
+        or BufferPointer.free()
+        """
+        func_ir = state.func_ir
+        for inst in self.iter_calls(func_ir):
+            rhs = func_ir.get_definition(inst.value.func)
+            if self.is_free_buffer(rhs) or self.is_BufferPoint_dot_free(state, rhs):
+                return True
+        return False
+
+    def run_pass(self, state):
+        on_missing_free = state.flags.on_missing_free
+        fid = state.func_id
+        if (self.contains_buffer_constructors(state) and not self.contains_calls_to_free(state)):
+            if on_missing_free == 'warn':
+                warnings.warn(MissingFreeWarning(fid.func_name, fid.filename, fid.firstlineno))
+            elif on_missing_free == 'fail':
+                raise MissingFreeError(fid.func_name, fid.filename, fid.firstlineno)
+            else:
+                raise ValueError(
+                    f"Unexpected value for on_missing_free: "
+                    f"got {on_missing_free:r}, expected 'warn', 'fail' or 'ignore'"
+                )
+        return False  # we didn't modify the IR
+
+
 class OmnisciCompilerPipeline(CompilerBase):
     def define_pipelines(self):
         # define a new set of pipelines (just one in this case) and for ease
         # base it on an existing pipeline from the DefaultPassBuilder,
         # namely the "nopython" pipeline
         pm = DefaultPassBuilder.define_nopython_pipeline(self.state)
         # Add the new pass to run after IRProcessing
-        pm.add_pass_after(AutoFreeBuffers, IRProcessing)
         pm.add_pass_after(CheckRaiseStmts, IRProcessing)
         pm.add_pass_after(DTypeComparison, ReconstructSSA)
+        if self.state.flags.on_missing_free != 'ignore':
+            pm.add_pass_after(DetectMissingFree, NopythonTypeInference)
         # finalize
         pm.finalize()
         # return as an iterable, any number of pipelines may be defined!

rbc/omniscidb.py

@@ -1175,7 +1175,8 @@ def _register(self):
                     target_info,
                     pipeline_class=OmnisciCompilerPipeline,
                     user_defined_llvm_ir=self.user_defined_llvm_ir.get(device),
-                    debug=self.debug)
+                    debug=self.debug,
+                    on_missing_free=self.on_missing_free)
 
                 assert llvm_module.triple == target_info.triple
                 assert llvm_module.data_layout == target_info.datalayout