tfhe-rs icon indicating copy to clipboard operation
tfhe-rs copied to clipboard

Prepare multi gpu support

Open agnesLeroy opened this issue 10 months ago • 5 comments

closes: please link all relevant issues

PR content/description

This PR prepares the code base for multi-gpu support. At the moment all functions fall back to using only GPU 0, but we'll be able to add multi-GPU support little by little on top of this.

Check-list:

  • [ ] Tests for the changes have been added (for bug fixes / features)
  • [ ] Docs have been added / updated (for bug fixes / features)
  • [ ] Relevant issues are marked as resolved/closed, related issues are linked in the description
  • [ ] Check for breaking changes (including serialization changes) and add them to commit message following the conventional commit specification

agnesLeroy avatar Apr 18 '24 07:04 agnesLeroy

@tmontaigu I've done what you suggested and created a MultiGpuFunctionExecutor for the tests, but actually I'm thinking a better way of doing this would be to pass a &[CudaStream] to all entry points, and they would make sure GPU 0 is only ever used when there is no multi-GPU support. Otherwise in the code as it is now, users could try to execute things on GPU 1, but that would conflict with the multi-GPU implementation of add, which supposes the keys, inputs & outputs are on GPU 0.

agnesLeroy avatar Apr 18 '24 07:04 agnesLeroy

If you think that is better you can do it

tmontaigu avatar Apr 18 '24 10:04 tmontaigu

Very interesting!!

mironov788 avatar Apr 20 '24 20:04 mironov788

@tmontaigu I gave a tour of the changes to @pdroalves earlier, so he has a good view of what's in the PR normally. I can go over it with you as well if you want, as this is a lot of changes...

agnesLeroy avatar Apr 24 '24 16:04 agnesLeroy

@tmontaigu thanks a lot for your feedback! I pushed a separate commit where I took your comments into account.

agnesLeroy avatar Apr 25 '24 15:04 agnesLeroy