feat(hlapi): add gpu selection #1990

tmontaigu · 2025-01-17T16:05:27Z

closes: please link all relevant issues

PR content/description

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

This change is

tfhe/src/core_crypto/gpu/vec.rs

tfhe/src/high_level_api/tests/gpu_selection.rs

IceTDrinker

Thanks, code looks fine, don't know about the general logic and per thread streams, I'm not familiar enough with the GPU I believe

tfhe/src/core_crypto/gpu/vec.rs

tfhe/src/high_level_api/integers/unsigned/inner.rs

tfhe/src/high_level_api/integers/unsigned/base.rs

IceTDrinker · 2025-01-24T11:05:26Z

tfhe/src/high_level_api/booleans/inner.rs

+                if ct.gpu_indexes() == streams.gpu_indexes() {
+                    MaybeCloned::Borrowed(ct.as_ref())
+                } else {
+                    MaybeCloned::Cloned(ct.duplicate(streams).0)
+                }


seen this pattern a few times, could it be a function without it being annoying to write ?

This one is indeed a bit annoying to rewrite as in each case ct is a different type, so some kind of trait would be needed, so its not that interesting

Ok no worries

tfhe/src/high_level_api/booleans/inner.rs

tfhe/src/high_level_api/global_state.rs

tfhe/src/high_level_api/tests/gpu_selection.rs

tfhe/src/core_crypto/gpu/vec.rs

mayeul-zama · 2025-01-24T13:50:53Z

tfhe/src/high_level_api/keys/server.rs

@@ -334,6 +344,22 @@ impl CudaServerKey {
    pub(crate) fn message_modulus(&self) -> crate::shortint::MessageModulus {
        self.key.key.message_modulus
    }
+
+    pub(crate) fn pbs_key(&self) -> &crate::integer::gpu::CudaServerKey {


Not sure about the name we could use here
Maybe key or server_key instead

Just noticed the name pbs_key is used in the cpu code.
It's not public so we could also rename it

Btw it seems that this function is used for code in HL array but not in other places
There are still many cuda_key.key.key which could use the same pattern

key feels too simple, computation_key could add more meaning, but its longer

What about server_key?

that could work yes

mayeul-zama · 2025-01-24T14:18:18Z

tfhe/src/high_level_api/global_state.rs

+    use std::cell::LazyCell;
+
+    thread_local! {
+        pub(crate) static CUDA_STREAMS: RefCell<CudaStreams> = RefCell::new(CudaStreams::new_multi_gpu());


Why is it needed when there's already POOL which has a field multi: LazyCell<CudaStreams>

Yes I hesitated, this can be removed, it would mean the with_thread_local_cuda_stream would then need to look at the key's gpu_indexes and return the corresponding stream.

But as it is now in the PR, with_thread_local_cuda_stream uses the stream returned by CudaStreams::new_multi_gpu()
The equivalent would be

pub(in crate::high_level_api) fn with_thread_local_cuda_streams< R, F: for<'a> FnOnce(&'a CudaStreams) -> R, >( func: F, ) -> R { POOL.with_borrow(|stream_pool| func(&stream_pool.multi)) }

Or am I mistaken?

No, with_thread_local_cuda_streams use the stream in CUDA_STREAMS, which is set when calling set_server_key/replace_server_key and so this stream is not always a multi-gpu, it can also be a single GPU

Actually this is something better to be done in another PR, as another that can be done with it would be to replace the need to do with_internal_keys and with_thread_local_cuda_streams with just 1 function all by making with_internal_keys also give a handle to the streams as well, but that requires some more changes

all the with_thing APIs are pub(crate) ?

Either pub(crate) or pub(in crate::hlapi)

tmontaigu

Reviewable status: 0 of 25 files reviewed, 11 unresolved discussions (waiting on @agnesLeroy, @IceTDrinker, and @mayeul-zama)

tfhe/src/core_crypto/gpu/vec.rs

IceTDrinker · 2025-01-28T09:42:09Z

is it in a state where it needs a re review ?

tmontaigu

yes

Reviewable status: 0 of 25 files reviewed, 11 unresolved discussions (waiting on @IceTDrinker and @mayeul-zama)

IceTDrinker

Mainly questions on whether move_to_stream is usable in some places (don't know if you checked them all)

Otherwise looks generally good to me

Reviewed 14 of 25 files at r1, 7 of 10 files at r2, 4 of 4 files at r3, all commit messages.
Reviewable status: all files reviewed, 9 unresolved discussions (waiting on @mayeul-zama and @tmontaigu)

tfhe/src/high_level_api/integers/signed/inner.rs line 183 at r3 (raw file):

            }
            Self::Cuda(cuda_ct) => {
                if cuda_ct.gpu_indexes() != streams.gpu_indexes() {

move_to_stream does not apply ?

tfhe/src/high_level_api/integers/signed/inner.rs line 224 at r3 (raw file):

                // We are on a GPU, but it may not be the correct one
                let new = with_thread_local_cuda_streams(|streams| {
                    if cuda_ct.gpu_indexes() == streams.gpu_indexes() {

move_to_stream does not apply ?

tfhe/src/high_level_api/booleans/inner.rs line 224 at r3 (raw file):

                // We are on a GPU, but it may not be the correct one
                let new = with_thread_local_cuda_streams(|streams| {
                    if cuda_ct.gpu_indexes() == streams.gpu_indexes() {

move_to_stream does not apply here ?

tfhe/src/high_level_api/keys/server.rs line 385 at r3 (raw file):

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Cpu(_) => f.debug_tuple("Cpu").finish(),

perhaps InteralServerKey::Cpu to match enum debug ? Don't recall what the format usually is

tfhe/src/high_level_api/integers/unsigned/inner.rs line 185 at r3 (raw file):

            }
            Self::Cuda(cuda_ct) => {
                if cuda_ct.gpu_indexes() != streams.gpu_indexes() {

move_to_stream does not apply ?

tfhe/src/high_level_api/integers/unsigned/inner.rs line 225 at r3 (raw file):

                // We are on a GPU, but it may not be the correct one
                let new = with_thread_local_cuda_streams(|streams| {
                    if cuda_ct.gpu_indexes() == streams.gpu_indexes() {

move_to_stream does not apply ?

IceTDrinker · 2025-01-28T13:38:16Z

tfhe/src/high_level_api/global_state.rs

+    use std::cell::LazyCell;
+
+    thread_local! {
+        pub(crate) static CUDA_STREAMS: RefCell<CudaStreams> = RefCell::new(CudaStreams::new_multi_gpu());


all the with_thing APIs are pub(crate) ?

IceTDrinker · 2025-01-28T13:38:16Z

tfhe/src/high_level_api/booleans/inner.rs

+                if ct.gpu_indexes() == streams.gpu_indexes() {
+                    MaybeCloned::Borrowed(ct.as_ref())
+                } else {
+                    MaybeCloned::Cloned(ct.duplicate(streams).0)
+                }


Ok no worries

mayeul-zama

Reviewed 15 of 25 files at r1, 10 of 10 files at r2, 4 of 4 files at r3, all commit messages.
Reviewable status: all files reviewed, 7 unresolved discussions (waiting on @tmontaigu)

tmontaigu

Reviewable status: all files reviewed, 7 unresolved discussions (waiting on @IceTDrinker)

tfhe/src/high_level_api/integers/signed/inner.rs line 183 at r3 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

move_to_stream does not apply ?

hum, we would need to add a move_to_stream_inplace(&mut self)

tfhe/src/high_level_api/integers/signed/inner.rs line 224 at r3 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

move_to_stream does not apply ?

Same as above, we would need to add a move_to_stream_inplace(&mut self)

tfhe/src/high_level_api/keys/server.rs line 385 at r3 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

perhaps InteralServerKey::Cpu to match enum debug ? Don't recall what the format usually is

Yes InteralServerKey::Cpu would probably be better

tmontaigu · 2025-01-29T10:07:27Z

tfhe/src/high_level_api/global_state.rs

+    use std::cell::LazyCell;
+
+    thread_local! {
+        pub(crate) static CUDA_STREAMS: RefCell<CudaStreams> = RefCell::new(CudaStreams::new_multi_gpu());


Either pub(crate) or pub(in crate::hlapi)

tmontaigu · 2025-01-29T11:12:40Z

tfhe/src/high_level_api/keys/server.rs

@@ -334,6 +344,22 @@ impl CudaServerKey {
    pub(crate) fn message_modulus(&self) -> crate::shortint::MessageModulus {
        self.key.key.message_modulus
    }
+
+    pub(crate) fn pbs_key(&self) -> &crate::integer::gpu::CudaServerKey {


that could work yes

cla-bot bot added the cla-signed label Jan 17, 2025

IceTDrinker reviewed Jan 17, 2025

View reviewed changes

tfhe/src/core_crypto/gpu/vec.rs Outdated Show resolved Hide resolved

tmontaigu force-pushed the tm/hlapi-context branch from 2dc9b44 to 67e3ed8 Compare January 20, 2025 10:38

tmontaigu commented Jan 20, 2025

View reviewed changes

tfhe/src/high_level_api/tests/gpu_selection.rs Outdated Show resolved Hide resolved

tmontaigu requested review from IceTDrinker and mayeul-zama January 21, 2025 09:43

IceTDrinker reviewed Jan 24, 2025

View reviewed changes

mayeul-zama reviewed Jan 24, 2025

View reviewed changes

tfhe/src/core_crypto/gpu/vec.rs Outdated Show resolved Hide resolved

tfhe/src/core_crypto/gpu/vec.rs Outdated Show resolved Hide resolved

tfhe/src/core_crypto/gpu/vec.rs Outdated Show resolved Hide resolved

mayeul-zama reviewed Jan 24, 2025

View reviewed changes

tmontaigu force-pushed the tm/hlapi-context branch from 67e3ed8 to 522a5b1 Compare January 27, 2025 12:44

tmontaigu commented Jan 27, 2025

View reviewed changes

tfhe/src/core_crypto/gpu/vec.rs Outdated Show resolved Hide resolved

tmontaigu force-pushed the tm/hlapi-context branch 3 times, most recently from e66991e to 59064a8 Compare January 27, 2025 14:36

feat(hlapi): add gpu selection

94d3e01

tmontaigu force-pushed the tm/hlapi-context branch from 59064a8 to 94d3e01 Compare January 27, 2025 15:03

tmontaigu requested a review from mayeul-zama January 27, 2025 15:22

tmontaigu commented Jan 28, 2025

View reviewed changes

IceTDrinker reviewed Jan 28, 2025

View reviewed changes

mayeul-zama approved these changes Jan 29, 2025

View reviewed changes

zama-bot added the approved label Jan 29, 2025

tmontaigu commented Jan 29, 2025

View reviewed changes

tmontaigu merged commit 57386a1 into main Jan 29, 2025
151 checks passed

tmontaigu deleted the tm/hlapi-context branch January 29, 2025 12:24

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(hlapi): add gpu selection #1990

feat(hlapi): add gpu selection #1990

tmontaigu commented Jan 17, 2025 •

edited

Loading

IceTDrinker left a comment

IceTDrinker Jan 24, 2025

tmontaigu Jan 27, 2025

IceTDrinker Jan 28, 2025

mayeul-zama Jan 24, 2025 •

edited

Loading

tmontaigu Jan 27, 2025

mayeul-zama Jan 27, 2025

tmontaigu Jan 29, 2025

mayeul-zama Jan 24, 2025

tmontaigu Jan 24, 2025

mayeul-zama Jan 24, 2025

tmontaigu Jan 27, 2025

tmontaigu Jan 27, 2025

IceTDrinker Jan 28, 2025

tmontaigu Jan 29, 2025

tmontaigu left a comment

IceTDrinker commented Jan 28, 2025

tmontaigu left a comment

IceTDrinker left a comment

IceTDrinker Jan 28, 2025

IceTDrinker Jan 28, 2025

mayeul-zama left a comment

tmontaigu left a comment

tmontaigu Jan 29, 2025

tmontaigu Jan 29, 2025

feat(hlapi): add gpu selection #1990

feat(hlapi): add gpu selection #1990

Conversation

tmontaigu commented Jan 17, 2025 • edited Loading

PR content/description

Check-list:

IceTDrinker left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mayeul-zama Jan 24, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tmontaigu left a comment

Choose a reason for hiding this comment

IceTDrinker commented Jan 28, 2025

tmontaigu left a comment

Choose a reason for hiding this comment

IceTDrinker left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mayeul-zama left a comment

Choose a reason for hiding this comment

tmontaigu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tmontaigu commented Jan 17, 2025 •

edited

Loading

mayeul-zama Jan 24, 2025 •

edited

Loading