feat(gpu): trivium

[enzodimaria]

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 

[enzodimaria]

https://www.notion.so/zamaai/Trivium-Kreyvium-2e25a7358d5e804880acc2c1ba6c2ef1

[IceTDrinker]

haven't read the GPU code, some remarks here and there, blocking for now as I think at least one change breaks benchmarks

@IceTDrinker reviewed 15 files and all commit messages, and made 10 comments.
Reviewable status: 15 of 19 files reviewed, 9 unresolved discussions (waiting on @agnesLeroy, @enzodimaria, @soonum, and @tmontaigu).

---

tfhe-benchmark/src/utilities.rs line 8 at r1 (raw file):

use tfhe::core_crypto::gpu::{get_number_of_gpus, get_number_of_sms};
use tfhe::core_crypto::prelude::*;
#[cfg(feature = "integer")]

the removal of this feature is incorrect I think, the FheEncrypt trait, as far as I know, is part of the HL which is only available when the integer feature is enabled

---

tfhe/src/integer/server_key/radix_parallel/tests_unsigned/test_trivium.rs line 1 at r1 (raw file):

#![cfg(feature = "gpu")]

prefer to put those where the modules are declared, like in test_unsigned in that case, easier to move the cfgs if they are grouped I would think ?

also if you put the cfg I'm guessing some checks were failing on e.g. CPU with this file ?

---

tfhe/src/integer/server_key/radix_parallel/tests_unsigned/test_trivium.rs line 16 at r1 (raw file):

}

impl TriviumRef {

maybe instantiate trivium ref tests to make sure we don't break the implem ? Plus it's going to be way faster than FHE so ok to have in tests at all times I think (may also allow you to get rid of the gpu feature at the top of the file)

---

tfhe-benchmark/benches/integer/trivium.rs line 65 at r1 (raw file):

        }

        {

those two blocks could be factorized with a loop over the number of steps

for steps_count in [64, 512] {
// Run bench, adapt bench naming and logging
}

---

tfhe/src/integer/gpu/server_key/radix/trivium.rs line 11 at r1 (raw file):

impl RadixClientKey {
    pub fn encrypt_bits_from_u64_slice_to_radix_ciphertext(&self, bits: &[u64]) -> RadixCiphertext {

I think there is something to encrypt iterators, @tmontaigu may recall better than me

---

tfhe/src/integer/gpu/server_key/radix/trivium.rs line 21 at r1 (raw file):

    }

    pub fn decrypt_bits_from_radix_ciphertext_to_vec_u8(

same thing, there have to be primitives that already do what you need

---

tfhe/src/integer/gpu/server_key/radix/trivium.rs line 53 at r1 (raw file):

        let num_iv_bits = 80;

        assert_eq!(

asserting is a good reflex, however we are trying to move towards returning Results to better manage errors, so here if you can try to return a Result and return the errors if preconditions aren't met

---

tfhe/src/integer/gpu/server_key/radix/trivium.rs line 79 at r1 (raw file):

            match &self.bootstrapping_key {
                CudaBootstrappingKey::Classic(d_bsk) => {
                    cuda_backend_trivium_generate_keystream(

haven't checked each parameter is where it should be

---

backends/tfhe-cuda-backend/src/bindings.rs line 2514 at r1 (raw file):

    );
}
unsafe extern "C" {

any concerns with KS32 ? I'm asking since I know Andrei had been working on it

[enzodimaria]

@IceTDrinker
In utilities.rs, the removal of the feature was a mistake but I have to manually add "use tfhe::prelude::FheEncrypt;" because if I don't, I have a pcc error

[IceTDrinker]

@IceTDrinker In utilities.rs, the removal of the feature was a mistake but I have to manually add "use tfhe::prelude::FheEncrypt;" because if I don't, I have a pcc error

then something is not right, it means something is not gated properly by integer somewhere as FheEncrypt is only available in that case

[IceTDrinker]

Approving to unblock, I did not check the GPU code since I'm not an expert

remaining questions/sanity checks:

• KS32 is not managed ?
• haven't checked the parameters are properly forwarded to the binding (likely is ok but just wanted to highlight I did not specifically check those)

@IceTDrinker reviewed 5 files and all commit messages, made 1 comment, and resolved 8 discussions.
Reviewable status: 15 of 19 files reviewed, 2 unresolved discussions (waiting on @agnesLeroy, @enzodimaria, @pdroalves, @soonum, and @tmontaigu).

[agnesLeroy]

@IceTDrinker we are not integrating KS32 at this stage so no need to take it into account in this PR.

[pdroalves]

backends/tfhe-cuda-backend/cuda/include/trivium/trivium_utilities.h line 67 at r4 (raw file):

};

template <typename Torus> struct int_trivium_state_workspaces {

I'm missing comments explaining what things are in these methods. For instance, I have to idea what's the meaning of "workspaces" in here. Can you add a few lines?

[pdroalves]

backends/tfhe-cuda-backend/cuda/include/trivium/trivium_utilities.h line 88 at r4 (raw file):

  CudaRadixCiphertextFFI *packed_flush_out;

  int_trivium_state_workspaces(CudaStreams streams,

Still regarding comments: If you look this file without knowing trivium it will be just a a huge wall of text and parameters. Use comments to split them in blocks, and add somewhere a quick comment explaining each variable.

[pdroalves]

backends/tfhe-cuda-backend/cuda/src/trivium/trivium.cuh line 3 at r4 (raw file):

#ifndef TRIVIUM_CUH
#define TRIVIUM_CUH

Once again asking for comments. This file doesn't have a single line ;-)

[pdroalves]

backends/tfhe-cuda-backend/cuda/src/trivium/trivium.cuh line 37 at r4 (raw file):

template <typename Torus>
__host__ __forceinline__ void
trivium_add(CudaStreams streams, int_trivium_buffer<Torus> *mem,

Do we really need this function?

[pdroalves]

backends/tfhe-cuda-backend/cuda/include/trivium/trivium_utilities.h line 273 at r4 (raw file):

    state = nullptr;

    cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));

We don't need this synchronization. It is expected that the stream synchronizes inside release().

[pdroalves]

tfhe-benchmark/Cargo.toml line 140 at r4 (raw file):

path = "benches/integer/trivium.rs"
harness = false
required-features = ["integer", "internal-keycache"]

Shouldn't we add gpu here?

[pdroalves]

tfhe-benchmark/Cargo.toml line 146 at r4 (raw file):

path = "benches/integer/aes256.rs"
harness = false
required-features = ["integer", "internal-keycache"]

Same question here. Is there a reason we don't have gpu in this list @agnesLeroy ?

[pdroalves]

backends/tfhe-cuda-backend/cuda/src/trivium/trivium.cuh line 281 at r4 (raw file):

  trivium_init(streams, mem, key_bitsliced, iv_bitsliced, bsks, ksks);

  uint32_t num_batches = num_steps / 64;

If num_steps is not divisible by 64 we may have a problem here, no?
We should mitigate it or at least add a PANIC in case num_steps % 64 != 0

[pdroalves]

@pdroalves reviewed 4 files and all commit messages, and resolved 1 discussion.
Reviewable status: all files reviewed, 9 unresolved discussions (waiting on @agnesLeroy, @enzodimaria, @soonum, and @tmontaigu).

[enzodimaria]

tfhe-benchmark/Cargo.toml line 146 at r4 (raw file):

path = "benches/integer/aes256.rs"
harness = false
required-features = ["integer", "internal-keycache"]

Same question here. Is there a reason we don't have gpu in this list @agnesLeroy ?

Hi Pedro! We didn't put the "gpu" feature in the aes benches neither, but maybe it was a mistake?

[agnesLeroy]

tfhe-benchmark/Cargo.toml line 146 at r4 (raw file):

Previously, enzodimaria wrote…

Hi Pedro! We didn't put the "gpu" feature in the aes benches neither, but maybe it was a mistake?

None of the bench has the "gpu" feature in that Cargo, it's only in the bench file that if the GPU feature is enabled, the bench runs, otherwise if it's available on CPU it runs on CPU. This bench only executes something for GPU but I think it's ok.

[enzodimaria]

Thank you @pdroalves for your comments, I updated the code with more comments, some additional checks about the size of the inputs and the inlining of additions (I wrote the function to debug the logic, I did add + flush into it, but it's not neccessary anymore)

[tmontaigu]

tfhe/src/integer/gpu/server_key/radix/trivium.rs line 11 at r1 (raw file):

Previously, IceTDrinker wrote…

I think there is something to encrypt iterators, @tmontaigu may recall better than me

I' m late to the party, but I don't see the encrypt\_bits\_from\_u64\_slice\_to\_radix\_ciphertext code anymore ?

[IceTDrinker]

tfhe/src/integer/gpu/server_key/radix/trivium.rs line 11 at r1 (raw file):

Previously, tmontaigu (tmontaigu) wrote…

I' m late to the party, but I don't see the encrypt\_bits\_from\_u64\_slice\_to\_radix\_ciphertext code anymore ?

yes replaced by utils functions directly in the test file, preferred avoiding adding encryption APIs without thinking them through

[tmontaigu]

tfhe/src/integer/gpu/server_key/radix/trivium.rs line 11 at r1 (raw file):

Previously, IceTDrinker wrote…

yes replaced by utils functions directly in the test file, preferred avoiding adding encryption APIs without thinking them through

Ok

[pdroalves]

Thanks @enzodimaria ! It is good enough for me.

@pdroalves reviewed all commit messages and made 1 comment.
Reviewable status: 16 of 19 files reviewed, 9 unresolved discussions (waiting on @agnesLeroy, @enzodimaria, @IceTDrinker, @soonum, and @tmontaigu).

[enzodimaria]

Thank you @pdroalves !