feat: register external buffers

[mertalev]

This PR allows registering external buffers with CubeCL, which is very useful as it enables zero-copy transfer. It internally uses a BindingMemory enum to distinguish external vs managed buffers as suggested in #291. I tried to keep this PR as minimal and non-invasive as possible - DLPack integration and such is left for future work.

Fixes #291

Validate your PR with burn.

It is important that you make sure that you don't introduce any bugs in burn.

Instructions

• Create a new branch or fork of the burn repo
• Update the main Cargo.toml with this PR hash.
• Fix any broken tests or compilation errors in burn.
•  Submit a PR in burn with your fixes and link it here.

[nathanielsimard]

The PR is good overall, but I have to think if it would be better to have a user managed memory pool instead. So all memory handles/bindings would be the same, but the underlying storage would be user managed. We already have different memory pools for different usage, so that could also work. Do you have an opinion about this?

[mertalev]

The PR is good overall, but I have to think if it would be better to have a user managed memory pool instead. So all memory handles/bindings would be the same, but the underlying storage would be user managed. We already have different memory pools for different usage, so that could also work. Do you have an opinion about this?

I think either approach works, but the pool approach has the benefit that bindings and handles are simple and don't need branching whenever they're used. Let me give it a shot and see how it looks.

[mertalev]

Yup, I think the pool version ends up being cleaner! One thing to note: I made it take ownership of the buffer since that's what the pool kind of leans toward with its cleanup etc., but I suppose dlpack integration would want borrow semantics for more flexibility. What do you think?

[nathanielsimard]

I think I would also allow manually deallocating a buffer on the user managed pool. That could be a lot faster than calling cleanup on all memory pools. It would be another way to dealloc a buffer, with fined grain control instead of relying on a "GC" like cleanup. Cleanup is not called often, since most of our pools don't need to deallocate, pretty much only when switching models to reset the memory pools allocations.

[mertalev]

True, I added an unregister API to allow this to let the caller either use or drop the underlying buffer.