Nearest neighbors graph computation crashes

[manock]

Hi,

I am trying the hlca_lung_gpu notebook with my own data to test the dimensionality reduction on GPU. The notebook is running well (data loading, pca...) until the nearest neighbors computation

Here is the traceback:

sc.pp.neighbors(adata, n_neighbors=n_neighbors, n_pcs=knn_n_pcs, method='rapids')
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/errors.py in new_error_context(fmt_, *args, **kwargs)
    743     try:
--> 744         yield
    745     except NumbaError as e:

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/lowering.py in lower_block(self, block)
    229                                    loc=self.loc, errcls_=defaulterrcls):
--> 230                 self.lower_inst(inst)
    231         self.post_block(block)

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/lowering.py in lower_inst(self, inst)
    327             val = self.lower_assign(ty, inst)
--> 328             self.storevar(val, inst.target.name)
    329 

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/lowering.py in storevar(self, value, name)
   1277                                                           name=name)
-> 1278             raise AssertionError(msg)
   1279 

AssertionError: Storing i64 to ptr of i32 ('dim'). FE type int32

During handling of the above exception, another exception occurred:

LoweringError                             Traceback (most recent call last)
<timed eval> in <module>

~/miniconda3/envs/rapids/lib/python3.7/site-packages/scanpy/neighbors/__init__.py in neighbors(adata, n_neighbors, n_pcs, use_rep, knn, random_state, method, metric, metric_kwds, key_added, copy)
    120         n_neighbors=n_neighbors, knn=knn, n_pcs=n_pcs, use_rep=use_rep,
    121         method=method, metric=metric, metric_kwds=metric_kwds,
--> 122         random_state=random_state,
    123     )
    124 

~/miniconda3/envs/rapids/lib/python3.7/site-packages/scanpy/neighbors/__init__.py in compute_neighbors(self, n_neighbors, knn, n_pcs, use_rep, method, random_state, write_knn_indices, metric, metric_kwds)
    744                 knn_distances,
    745                 self._adata.shape[0],
--> 746                 self.n_neighbors,
    747             )
    748         # overwrite the umap connectivities if method is 'gauss'

~/miniconda3/envs/rapids/lib/python3.7/site-packages/scanpy/neighbors/__init__.py in _compute_connectivities_umap(knn_indices, knn_dists, n_obs, n_neighbors, set_op_mix_ratio, local_connectivity)
    345     fuzzy simplicial sets into a global one via a fuzzy union.
    346     """
--> 347     from umap.umap_ import fuzzy_simplicial_set
    348 
    349     X = coo_matrix(([], ([], [])), shape=(n_obs, 1))

~/miniconda3/envs/rapids/lib/python3.7/site-packages/umap/__init__.py in <module>
----> 1 from .umap_ import UMAP
      2 
      3 # Workaround: https://github.com/numba/numba/issues/3341
      4 import numba
      5 

~/miniconda3/envs/rapids/lib/python3.7/site-packages/umap/umap_.py in <module>
     52 from umap.spectral import spectral_layout
     53 from umap.utils import deheap_sort, submatrix
---> 54 from umap.layouts import (
     55     optimize_layout_euclidean,
     56     optimize_layout_generic,

~/miniconda3/envs/rapids/lib/python3.7/site-packages/umap/layouts.py in <module>
     34         "result": numba.types.float32,
     35         "diff": numba.types.float32,
---> 36         "dim": numba.types.int32,
     37     },
     38 )

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/decorators.py in wrapper(func)
    219             with typeinfer.register_dispatcher(disp):
    220                 for sig in sigs:
--> 221                     disp.compile(sig)
    222                 disp.disable_compile()
    223         return disp

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/dispatcher.py in compile(self, sig)
    907                 with ev.trigger_event("numba:compile", data=ev_details):
    908                     try:
--> 909                         cres = self._compiler.compile(args, return_type)
    910                     except errors.ForceLiteralArg as e:
    911                         def folded(args, kws):

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/dispatcher.py in compile(self, args, return_type)
     77 
     78     def compile(self, args, return_type):
---> 79         status, retval = self._compile_cached(args, return_type)
     80         if status:
     81             return retval

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/dispatcher.py in _compile_cached(self, args, return_type)
     91 
     92         try:
---> 93             retval = self._compile_core(args, return_type)
     94         except errors.TypingError as e:
     95             self._failed_cache[key] = e

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/dispatcher.py in _compile_core(self, args, return_type)
    109                                       args=args, return_type=return_type,
    110                                       flags=flags, locals=self.locals,
--> 111                                       pipeline_class=self.pipeline_class)
    112         # Check typing error if object mode is used
    113         if cres.typing_error is not None and not flags.enable_pyobject:

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/compiler.py in compile_extra(typingctx, targetctx, func, args, return_type, flags, locals, library, pipeline_class)
    604     pipeline = pipeline_class(typingctx, targetctx, library,
    605                               args, return_type, flags, locals)
--> 606     return pipeline.compile_extra(func)
    607 
    608 

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/compiler.py in compile_extra(self, func)
    351         self.state.lifted = ()
    352         self.state.lifted_from = None
--> 353         return self._compile_bytecode()
    354 
    355     def compile_ir(self, func_ir, lifted=(), lifted_from=None):

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/compiler.py in _compile_bytecode(self)
    413         """
    414         assert self.state.func_ir is None
--> 415         return self._compile_core()
    416 
    417     def _compile_ir(self):

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/compiler.py in _compile_core(self)
    393                 self.state.status.fail_reason = e
    394                 if is_final_pipeline:
--> 395                     raise e
    396         else:
    397             raise CompilerError("All available pipelines exhausted")

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/compiler.py in _compile_core(self)
    384             res = None
    385             try:
--> 386                 pm.run(self.state)
    387                 if self.state.cr is not None:
    388                     break

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/compiler_machinery.py in run(self, state)
    337                     (self.pipeline_name, pass_desc)
    338                 patched_exception = self._patch_error(msg, e)
--> 339                 raise patched_exception
    340 
    341     def dependency_analysis(self):

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/compiler_machinery.py in run(self, state)
    328                 pass_inst = _pass_registry.get(pss).pass_inst
    329                 if isinstance(pass_inst, CompilerPass):
--> 330                     self._runPass(idx, pass_inst, state)
    331                 else:
    332                     raise BaseException("Legacy pass in use")

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/compiler_lock.py in _acquire_compile_lock(*args, **kwargs)
     33         def _acquire_compile_lock(*args, **kwargs):
     34             with self:
---> 35                 return func(*args, **kwargs)
     36         return _acquire_compile_lock
     37 

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/compiler_machinery.py in _runPass(self, index, pss, internal_state)
    287             mutated |= check(pss.run_initialization, internal_state)
    288         with SimpleTimer() as pass_time:
--> 289             mutated |= check(pss.run_pass, internal_state)
    290         with SimpleTimer() as finalize_time:
    291             mutated |= check(pss.run_finalizer, internal_state)

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/compiler_machinery.py in check(func, compiler_state)
    260 
    261         def check(func, compiler_state):
--> 262             mangled = func(compiler_state)
    263             if mangled not in (True, False):
    264                 msg = ("CompilerPass implementations should return True/False. "

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/typed_passes.py in run_pass(self, state)
    461 
    462         # TODO: Pull this out into the pipeline
--> 463         NativeLowering().run_pass(state)
    464         lowered = state['cr']
    465         signature = typing.signature(state.return_type, *state.args)

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/typed_passes.py in run_pass(self, state)
    382                 lower = lowering.Lower(targetctx, library, fndesc, interp,
    383                                        metadata=metadata)
--> 384                 lower.lower()
    385                 if not flags.no_cpython_wrapper:
    386                     lower.create_cpython_wrapper(flags.release_gil)

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/lowering.py in lower(self)
    134         if self.generator_info is None:
    135             self.genlower = None
--> 136             self.lower_normal_function(self.fndesc)
    137         else:
    138             self.genlower = self.GeneratorLower(self)

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/lowering.py in lower_normal_function(self, fndesc)
    188         # Init argument values
    189         self.extract_function_arguments()
--> 190         entry_block_tail = self.lower_function_body()
    191 
    192         # Close tail of entry block

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/lowering.py in lower_function_body(self)
    214             bb = self.blkmap[offset]
    215             self.builder.position_at_end(bb)
--> 216             self.lower_block(block)
    217         self.post_lower()
    218         return entry_block_tail

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/lowering.py in lower_block(self, block)
    228             with new_error_context('lowering "{inst}" at {loc}', inst=inst,
    229                                    loc=self.loc, errcls_=defaulterrcls):
--> 230                 self.lower_inst(inst)
    231         self.post_block(block)
    232 

~/miniconda3/envs/rapids/lib/python3.7/contextlib.py in __exit__(self, type, value, traceback)
    128                 value = type()
    129             try:
--> 130                 self.gen.throw(type, value, traceback)
    131             except StopIteration as exc:
    132                 # Suppress StopIteration *unless* it's the same exception that

~/miniconda3/envs/rapids/lib/python3.7/site-packages/numba/core/errors.py in new_error_context(fmt_, *args, **kwargs)
    749         newerr = errcls(e).add_context(_format_msg(fmt_, args, kwargs))
    750         tb = sys.exc_info()[2] if numba.core.config.FULL_TRACEBACKS else None
--> 751         raise newerr.with_traceback(tb)
    752 
    753 

LoweringError: Failed in nopython mode pipeline (step: nopython mode backend)
Storing i64 to ptr of i32 ('dim'). FE type int32

File "../../../../../home/egilson/miniconda3/envs/rapids/lib/python3.7/site-packages/umap/layouts.py", line 52:
def rdist(x, y):
    <source elided>
    result = 0.0
    dim = x.shape[0]
    ^

During: lowering "dim = static_getitem(value=$8load_attr.2, index=0, index_var=$const10.3, fn=<built-in function getitem>)" at /home/egilson/miniconda3/envs/rapids/lib/python3.7/site-packages/umap/layouts.py (52)

[m0nib]

Hi, I had the same issue when using the default hlc data set.

i googled a bit found a solution that worked for me, not sure if its a solution or just a work around but might be helpful to others so here it goes.

The issue seems to be with the numba==0.53.1 and downgrading to numba==0.52 has fixed the issue for me.

pip uninstall numba
pip install numba==0.52.0

Downgrading numba can through some errors so i followed instructions as discussed here "NVIDIA/NeMo#841"

simply run
pip install llvmlite --ignore-installed
pip install numba==0.52.0 --ignore-installed

you might see some errors/warnings but i restarted the kernel in jupyter lab and the whoel workflow worked like a charm and was very quick on NVIDIA 3090.

awesome work guys.

Thank you

Best regards
Monib