An Error Correction (De)Coding library in pure Mojo 🔥
MoCodes is a stand-alone Error Correction (De)Coding framework that leverages the power of Mojo.
As discussed by Modular, Mojo is a language for the future of AI development. Built on top of MLIR technology, rather than existing GCC and LLVM approaches, Mojo looks and feels like Python code, yet performs much closer to languages like Rust or C++.
Error Correction Codes are being used in domains such as Wireless Communication and Quantum Computing. They are known to be very compute intensive, so much so that until recently they were implemented in dedicated ASIC silicon or programmed on FPGA accelerators.
In recent years with the advent of wide-vector SIMD CPU architectures and affordable yet powerful GPU cores they have been also implemented in C++ with a lot of vendor-specific intrinsics on CPU or with CUDA on GPU as well.
About time then to take a stab at how well Mojo lives up to the challenge (and how well the authors live up to the challenge of understanding how Mojo is meant to be used)
We've been on a voyage of exploration to find out how platform-independent frameworks originally meant for machine learning can be repurposed for error correction decoding. We have had reasonable results on GPU and TPU but not until Mojo came along we've reached decent throughputs on CPU. Results ofc vary with platform specifics, we have tried Intel, AMD and Macbook M3.
While that looks awesome we estimate that there is a performance gap of about a factor 2 yet to be closed wrt C++ code that uses vendor specific intrinsics eg from the avx instruction set on Intel. Game on !
For now we only support generic (ir)regular LDPC codes with embarrassingly parallel flooding batch decoding. We have committed just one example (1512 x 1872) LDPC Parity Check Matrix to this repo. This sparse parity check matrix has 7092 non-zero elements and is shown hereafter.
For now, this parity check matrix gets translated to look-up tables by an offline scipy script that takes an .npz file as input. The look-up tables get stored in the /codebook/ subdirectory.
Try out the LDPC benchmark for yourself, on your own platform:
mojo build ldpcdec.mojo
./ldpcdec
You can tweak the following parameters in the main() function of types.mojo: intra_codeword_parallellism_factor, ncodewordperthread, nthread. Currently committed defaults seem to be close to optimal regardless the platform we have tried.
- support for irregular LDPC decoding
- support for (batch) early stopping
- Improve throughput (target: factor 2) by community expertise injection
- Add a serving functionality (preferably gRPC based, ideally leveraging MAX serving)
- Add profiling and proper benchmarking tests
- incorporate generation of Look-Up Tables in the mojo code, such that the .npz file becomes the only configuration input that defines the code
- add an LDPC encoder
- add a script to simulate and visualise BER and BLER curves
- Autotuning and related features
- Add polar codes
- Add Reed-Solomon codes
- Add CRC check codes
- Add layered LDPC decoding
The way we set this repo up should allow Mojo experts to contribute without necessarily being Error Correction Coding specialists. Notably, the LDPC heavy lifting is done by a handful of functions in types.mojo, i.e. fn all_Lqij() and fn cnpu(). Memory load and store at this point seem to determine the throughput so all tips and tricks to speed up that memory access would much appreciated, along with any other improvements that can be spotted by expert Mojicans.
If you are considering larger contributions, feel free to contact us for a smoother communication channel on Discord. If you find a bug or have an idea for a feature, please use our issue tracker. Before creating a new issue, please:
- Check if the issue already exists. If an issue is already reported, you can contribute by commenting on the existing issue.
- If not, create a new issue and include all the necessary details to understand/recreate the problem or feature request.
- Fork the Project
- Create your Feature Branch
- Commit your Changes
- Push to the Branch
- Open a Pull Request
Once your changes are pushed, navigate to your fork on GitHub. And create a pull request against the original repository.
- Before creating a PR make sure the functional text output of ./ldpcdec is the same as the one on the main branch, i.e the batch syndrome on the all-zeros codewords should stay zero and the logit output values should all still equal 127.
- In the pull request, provide a detailed description of the changes and why they're needed. Link any relevant issues.
Distributed under the Apache 2.0 License.