Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Enable ARM builds for linux via emulation #7

Merged

Conversation

benhowes
Copy link
Contributor

@benhowes benhowes commented Jan 8, 2024

Closes #6.

This PR adds support for building wheels for aarch64 (arm 64) for linux.

The only downside I can see is that due to emulation, each arm linux build takes approximately 300s, so about 1200s (20 minutes) total extra time for the build.

Testing

Demo of usage in docker on macOS with an M3 chip (arm64 based):

  1. Download artifacts from test run: https://github.com/neuronostics/rocket-fft/actions/runs/7450533530
  2. Start container: docker run -it --entrypoint bash -v ~/Downloads/artifact:/artifact python:3.11
  3. Install correct wheel pip install /artifact/rocket_fft-0.2.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
  4. Run a test file:
    root@355ff7fa867a:/# cat test.py 
    import numba as nb
    import numpy as np
    
    @nb.njit
    def jit_fft(x):
        return np.fft.fft(x)
    
    a = np.array([1, 6, 1, 8, 0, 3, 3, 9])
    print(jit_fft(a))
    
    root@355ff7fa867a:/# pip freeze
    llvmlite==0.41.1
    numba==0.58.1
    numpy==1.26.3
    rocket-fft @ file:///artifact/rocket_fft-0.2.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl#sha256=5cfa60c34480522ca3181ab3c35647630477a5280d58f8ef0dde612fd4341107
    root@355ff7fa867a:/# uname -a
    Linux 355ff7fa867a 6.5.11-linuxkit #1 SMP PREEMPT Wed Dec  6 17:08:31 UTC 2023 aarch64 GNU/Linux
    root@355ff7fa867a:/# python test.py
    [ 31.        -0.j           3.82842712+0.58578644j
      -3.        +8.j          -1.82842712-3.41421356j
     -21.        -0.j          -1.82842712+3.41421356j
      -3.        -8.j           3.82842712-0.58578644j]
    

@styfenschaer
Copy link
Owner

Hi @benhowes

Thanks a lot for looking into this, and apologize for my delayed response.

The long build time is unfortunate but nothing to worry about. Could you let me know if you've had the opportunity to test the wheel? I plan to package a new release over the coming weekend.

Best
Styfen

@benhowes
Copy link
Contributor Author

@styfenschaer I've updated the PR with some very basic testing I've just done in a docker container on mac, but happy to look at doing more if you have anything in particular you'd like to see?

@styfenschaer styfenschaer merged commit ed5f7d8 into styfenschaer:main Jan 14, 2024
@styfenschaer
Copy link
Owner

I created a new release on TestPyPI. Can you confirm that this works for you? I will then upload it to PyPI too.

pip install -i https://test.pypi.org/simple/ rocket-fft==0.2.3

@benhowes
Copy link
Contributor Author

@styfenschaer yep, seems to work fine with the same test file as above. Thanks!

@styfenschaer
Copy link
Owner

Hi @benhowes
The new release is now also on PyPI. Thanks again for your help.

@benhowes
Copy link
Contributor Author

Thanks for getting this resolved so quickly 👍

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Linux arm64 builds
2 participants