DGEMM in Mojo: close to C++ for large matrices, behind for small ones

[michael-lehn]

Hi guys,

I tried to write a high performance DGEMM (double precision general matrix-matrix multiply) in Mojo: https://github.com/michael-lehn/mojo-gemm

Some background: I wrote ulmBLAS, a BLAS (Basic Linear Algebra Subprograms) implementation in C++, and FLENS-LAPACK. The Mojo code is a first experiment, not a full port. It is roughly what my students implement in my HPC (High Performance Computing) course: packing, a macro kernel and a micro kernel. The micro kernels come from ulmBLAS. NEON is plain Mojo SIMD (single instruction, multiple data) code, SME2 (Scalable Matrix Extension 2) is inline assembly. Row major, column major and arbitrary strides all work.

Results on an Apple M5, one thread:

• large matrices: about 2% behind ulmBLAS with NEON, about 5% with SME2
• small matrices: SME2 falls behind clearly, about 63% of ulmBLAS at n = 250

The gap at small sizes isnt Mojo's fault. ulmBLAS packs in streaming mode and handles partial tiles with predicates. The Mojo version dont do that yet.

What I would like to know: would NuMojo use a Mojo BLAS/LAPACK (Linear Algebra PACKage) backend? I saw the "TODO: Use LAPACK" in linalg. I am thinking about blocked Cholesky, LU and the solvers first (the set ulmBLAS has), later QR, SVD (singular value decomposition) and eig. How would that fit with mojoBLAS? I dont want to duplicate work that is allready going on.

[shivasankarka]

@michael-lehn I looked at your results, That's pretty cool! I would like to sit and go through the algorithms sometime since I work on optimization too :)

Regarding your question, One of our plans for NuMojo has been to introduce a BLAS, LAPACK backend, especially as we've been working on streamlining the CPU backend with HostExecutor.

I've previously worked on porting BLAS routines to Mojo in mojoBLAS, with promising results, although I haven't spent much time optimizing them recently. The idea was to use mojoBLAS as an experimental playground for developing and optimizing BLAS routines, and eventually upstream the mature implementations into NuMojo.

It would be great to have you contribute to either mojoBLAS or NuMojo directly! We can either add all mojoBLAS and LAPACK routines together as one library and then upstream it to NuMojo or just make PRs directly to NuMojo. I'd be interested to hear which approach you think makes more sense for you. We can discuss it and coordinate our efforts. For context, our current focus in NuMojo is on GPU NDArray support and filling in some missing NumPy functionality, so the optimization domain isn't being worked on actively.

If you'd like to submit a PR, just let me know! There's a small onboarding requirement: the first PR needs to be explicitly created by an existing contributor. After that, you'll be able to submit PRs independently. Looking forward to collaborating on this!

[michael-lehn]

Thanks a lot!

The repo now goes one step further: https://github.com/michael-lehn/mojo-gemm

• The SME2 (Scalable Matrix Extension 2) GEMM (general matrix-matrix multiply) got faster for small matrices: 77 to 89% of ulmBLAS at n = 250, before it was 63%. At n = 100 there is still work to do. That mainly a port of what I did on an experimental ulmBLAS branch. Still not where I want it.
• On top of it there is now a Cholesky factorization (potrf). In place with incRow and incCol so row major, column major and arbitrary strides all work. On my MacBook Air with M5 it's on par with Accelerate's dpotrf and faster for large matrices and for row major.
• The README also explains where I want to go with this: from the kernels over LAPACK (Linear Algebra PACKage) up to the linear algebra mapping problem. That is, a compiler that turns formulas into BLAS (Basic Linear Algebra Subprograms) and LAPACK calls.

About your question: I think one library first makes more sense. LAPACK and BLAS level 3 belong together. potrf needs gemm, trsm and syrk with strides and the tuning (which algorithm, how many threads, which tile size) depends on the kernels. In mojoBLAS we can iterate fast. Once the interface is stable we upstream it into NuMojo. Does that work for you?

One design point I would like to keep: everything in place and with strides. Then NuMojo can call it on its arrays directly, whatever the layout, without copys.

And yes, I would be happy to submit a PR (pull request). A first candidate: gemm with the NEON and SME2 micro kernels and potrf. Just let me know how you want to do the onboarding.

Section 4 of the README has some thoughts on what Mojo would need for HPC (high performance computing) in general: an infrastructure like BLIS (BLAS-like Library Instantiation Software) with experts for micro kernels and for numerical linear algebra. I would be curious to hear your view on that.

20 years ago FLENS (Flexible Library for Efficient Numerical Solutions) and Eigen were competitors. Eigen won, and not because of the code. I could not build a community around FLENS, they had the KDE (K Desktop Environment) community behind them. Mojo is new. Compared to NumPy, PyTorch or Julia we are few. So everyone counts, and we need a strategy that makes every contribution count.

[shivasankarka]

@michael-lehn I'll go through your implementation and results! Those are some really nice improvements compared to Accelerate! The ulm-LAMP compiler is also a really interesting direction!

About your question: I think one library first makes more sense. LAPACK and BLAS level 3 belong together. potrf needs gemm, trsm and syrk with strides and the tuning (which algorithm, how many threads, which tile size) depends on the kernels. In mojoBLAS we can iterate fast. Once the interface is stable we upstream it into NuMojo. Does that work for you?

Yep, that works for me! It's much easier to iterate quickly in the mojoBLAS repo for now. It already has a basic foundation for the BLAS routines, so we can build on that and add LAPACK too. Once the implementations and interfaces are stable, we can integrate them into NuMojo.

One design point I would like to keep: everything in place and with strides. Then NuMojo can call it on its arrays directly, whatever the layout, without copys.

That makes sense! Supporting arbitrary strides and in-place operations would make integration with NuMojo's arrays much easier without unnecessary copies since NDArray in NuMojo also aims to support these. I definitely agree with keeping that flexibility.

And yes, I would be happy to submit a PR (pull request). A first candidate: gemm with the NEON and SME2 micro kernels and potrf. Just let me know how you want to do the onboarding.

For mojoBLAS, there's no onboarding process. Feel free to fork the repo, add your changes, and open a PR tagging me. I'll usually be quick to review!

For NuMojo, we have a small onboarding process for new contributors. You can fork the repo, create a branch with your changes, and let me know. I'll open the first PR from your branch, and once it's merged, I can add you as a contributor so you can submit PRs independently.

If you're on Discord, feel free to join our server too and add me there. I'm usually quicker to respond there!

Section 4 of the README has some thoughts on what Mojo would need for HPC (high performance computing) in general: an infrastructure like BLIS (BLAS-like Library Instantiation Software) with experts for micro kernels and for numerical linear algebra. I would be curious to hear your view on that.

I read through section 4, and I'm definitely interested in this direction, since one of my main interests is scientific computing and performance optimization. The idea of building a BLIS-like infrastructure around reusable microkernels, platform-specific tuning, and higher-level numerical algorithms makes a lot of sense.

I'll read through the paper and the FAMLIES material to understand the research better. I also agree with your point about community building. Having a shared infrastructure where contributions from kernel developers and numerical linear algebra experts can build on each other would be valuable, especially while the Mojo ecosystem is still growing. From my experience so far, Mojo has been great for writing high-performance kernels, and I think there's an amazing opportunity to build something useful here!

I wasn't aware of FLENS before, but I'll look into it as well. Thanks for sharing all these references and your experience! I'd be happy to discuss this further and collaborate on it!

[michael-lehn]

Great, thanks for reading all of it! Then I start with a fork of mojoBLAS and open a PR when gemm and potrf are in shape.

About FLENS: I should have mentioned that it isn't actively developed anymore, for quite a while actually. But it was a good experience after all. Mainly to understand the idea behind LAPACK and how to build numerical linear algebra packages that are both performant and modular. BLAS and LAPACK work together, but the BLAS can be optimized independently of LAPACK. And with PLASMA (Parallel Linear Algebra Software for Multicore Architectures) LAPACK became independent of the BLAS in a sense too. Not really a new idea in software engineering, but seeing how its done in a specific field was different. And I got to see what actually goes on in implementations like geev (eigenvalues) or gesdd (SVD, singular value decomposition). So it's not a library to look into today, rather the experience behind what I want to do in Mojo.

And yes, I'll join the Discord!