burn
burn copied to clipboard
tracel-ai
Problem with no_std on thumbv6m-none-eabi
As discussed here, there seems to be a problem because AtomicBool is not available for the thumbv6m-none-eabi target and Burn uses it for locking (multithreading).
As suggested by @antimora, it could be possible to remove that dependency on the lock.
@spadarian. Just to give a background. We use locking to update a random seed. By default when a op module (e.g. Conv) is is created, we initialize random weights. One solution is to introduce a feature flag that would disable the locking and assume a single thread. This is safe for inference mode since the weights won't be initialized randomly.
Another option that I just looked up is to use spin's portable_atomic but it seems it not straight forward to do it on Burn.
The RP2040 has 2 cores but I think the single thread assumption is fair for this particular case.
I've successfully removed a lock on the random's rng by using thread_rng for std and bypassing the lock for no_std. However, I've noted that no_std assumes single-threaded execution. This may be valid if someone builds with no_std but runs the library in a multi-threaded environment. Simply put, even though the threading API isn't accessible to the crate, it doesn't mean the code isn't running with multiple threads. To account for this, we might have to introduce a new feature flag, potentially named single_thread. This would be a significant change and challenging to enforce across all downstream crates.
Even if we were to implement a single_thread feature flag, further refactoring in BathNorm would still be required. We'd need to replace the RunningState struct with an alternative that can operate without locks.
In summary, this task represents a substantial shift and is more complex than initially anticipated.
CCing @nathanielsimard
Thanks for looking into that.
My use case is quite simple at the moment and I just need to perform inference... but it does look complex to integrate this into a large crate.
BatchNorm just for inference would be easy to implement but not sure if the use case is important enough to change things.
@spadarian we will be looking into this some more. It feels like it should be doable but at the moment it's not straight forward.
I don't think having a single_thread feature flag is the right approach here, Having a lock that can run on thumbv6m-none-eabi seems easier to do and more correct. A spin lock might actually be a good alternative here, but I'm not sure if the implementations can run on your platform (maybe they use AtomicBool).
OK. I figured out how to do this by using portable atomic lib. Here is the PR in WIP: https://github.com/burn-rs/burn/pull/374
My work is still blocked till we make changes in the upstream lib: Ndarray. We are using ArcArray and references system Arc. I found a crate that can replace it with portable atomic implementation. I think it will take a while to get this, since their project is moving very slow. But I'll work on this more and I'll submit a PR.
Someone submitted a PR to make ndarray thumbv6m-none-eabi compatible:
https://github.com/rust-ndarray/ndarray/pull/1384
This PR has been merged, all that is needed is a new release by ndarray, it seems like that will happen relatively soon.
