feat: Add Metal/MPS support for macOS GPU acceleration#422
Open
phisanti wants to merge 3 commits intoquantumjot:mainfrom
Open
feat: Add Metal/MPS support for macOS GPU acceleration#422phisanti wants to merge 3 commits intoquantumjot:mainfrom
phisanti wants to merge 3 commits intoquantumjot:mainfrom
Conversation
This commit introduces Metal/MPS support for GPU-accelerated belief matrix updates on macOS, providing an alternative to CUDA. Key changes include: - Added `BayesianUpdates.MPS` enum to `btrack/constants.py` to represent the Metal/MPS update method. - Updated `btrack/include/defs.h` and `btrack/include/tracker.h` to include definitions and function declarations for Metal/MPS support. - Created `btrack/src/belief.mm` and `btrack/src/belief.metal` to implement the Metal compute shader for belief matrix updates. - Modified `btrack/src/tracker.cc` to incorporate the Metal/MPS cost function and select the appropriate backend for belief updates. - Added a `working_build.sh` script to set up the Metal development environment and build the tracker with Metal support. - Added `tests/test_btrack_mps.py` to test the Metal/MPS implementation. - Updated `woking_makefile.make` to include Metal compilation rules and dependencies. The Metal implementation provides significant performance improvements on macOS devices with Apple silicon. The build system now detects Metal availability and compiles the necessary files. A new test suite verifies the correctness of the Metal implementation and ensures that it produces similar results to the EXACT method.
|
Codecov ReportAll modified and coverable lines are covered by tests ✅
❗ Your organization needs to install the Codecov GitHub app to enable full functionality. Additional details and impacted files@@ Coverage Diff @@
## main #422 +/- ##
==========================================
+ Coverage 85.58% 85.71% +0.12%
==========================================
Files 31 31
Lines 2206 2219 +13
Branches 405 208 -197
==========================================
+ Hits 1888 1902 +14
- Misses 235 236 +1
+ Partials 83 81 -2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Hey there! I've been using btrack daily in my pipeline and absolutely love what you've built here. Since I work with movies containing massive numbers of objects, I thought it would be wonderful to help push the GPU acceleration forward.
This implementation brings the existing CUDA concept to Metal/MPS for macOS users. (Plus connect the CUDA backend to the main tracker code)
However, I stumbled upon an interesting insight during development: looks like the ILP solver is actually where the real bottleneck lives (unless I am doing something wrong with the config file), not in tracklet generation as we might expect. I didn't dive into that rabbit hole for this PR, but I'd be happy to collaborate on optimizing that area too if it would help the project.
I have also included a unit test file to validate that the tracking with the MPS backed mirrors what happens when using the EXACT method, and all tests are PASSED, so I'm confident this will work as intended.
Key Changes
BayesianUpdates.MPSenum tobtrack/constants.pyto represent the Metal/MPS update methodbtrack/include/defs.handbtrack/include/tracker.hto include definitions and function declarations for Metal/MPS supportbtrack/src/belief.mm- Objective-C++ wrapper for Metal integrationbtrack/src/belief.metal- Metal compute shader for belief matrix updatesbtrack/src/tracker.ccto incorporate the Metal/MPS cost function and select the appropriate backend for belief updates.Makefileto include Metal compilation rules and dependenciestests/test_btrack_mps.pyto verify Metal/MPS implementation correctness