MLIR based Numba backend

Hi, we are working on MLIR based backend for Numba python compiler https://github.com/IntelPython/numba/pull/115

We are aware of npcomp but for this project we are trying to reuse as much as possible of existing Numba infrastructure instead of doing everything from scratch. Right now we are using Numba frontend (which parses CPython bytecode) and Numba type inference to parse python code and also existing Numba llvm-based infrastructure to run the result.

It is still on very early stages, but we have some support for basic python control flow and for some (two) Numpy functions, see tests:

(more info to come…)

2 Likes

This is really nice!
I’d be interested in an overview of the flow there and how this play within Numba, let me know if you or anyone else in the project would like to present during one of our public meetings on Thursday :slight_smile:

This is really fantastic - thanks so much for sharing! (also, as someone who works on npcomp, I would just like to say that I am thrilled that numba is connecting up to MLIR in this way). Once this gets a bit further, I’d love to see how we might use numba for hacking on higher level program constructs such as sparse representations. We’re currently lacking good frontend tools for a variety of things that MLIR should in theory be good at doing, and virtuous frontend/backend cycles can really help move things forward here.

Down the road a bit, I’m interested to talk about API layerings and what might help keep us from having N different copies of LLVM for everything (i.e. making sure that the MLIR APIs provide a suitable entrypoint similar to what llvmlite does for Numba today, if I understand correctly).

I think I can give 20 min overview, when the next available slot? We are also planning more in-depth presentation but it will take some time to prepare.

What about March 4th?

Yes, March 4th should work. For this date we can prepare full (1h?) presentation, I think

2 Likes

As a reminder: this is tomorrow! :slight_smile:

Some more info on memory handling I didn’t discussed yesterday. Numba has its own reference-counted allocation scheme represented by MemInfo numba/nrt.c at mlir-develop · IntelPython/numba · GitHub and I wanted it to interoperate with MLIR memrefs without copies. So I added my own AllocOp/DeallocOp lowering (numba/lower_to_llvm.cpp at mlir-develop · IntelPython/numba · GitHub) to the corresponding numba alloc/decref calls. And I am storing numba MemInfo pointer in memref first field (which usually stores unaligned pointer).

Also I added custom Alloc op to increment reference counter (numba/PlierOps.td at mlir-develop · IntelPython/numba · GitHub) and special pass which apply it to the input memref params.
This is needed for code like this:

def foo(some_array):
    return some_array # some_array refcount must be inremented here

Some notes:

  • Can we make first element of memref struct just void* (i8* in llvm) pointer? custom allocation schemes will use it for any custom data they need (default implementation will still use it for unaligned pointer)
  • There are helper classes for allocation lowering in std-to-llvm conversion llvm-project/StandardToLLVM.cpp at main · llvm/llvm-project · GitHub but they are private and I had to copypase them into my code