Skip to content

Commit

Permalink
feat(hlapi): Add initial structure of NdArray types
Browse files Browse the repository at this point in the history
  • Loading branch information
tmontaigu committed Sep 25, 2024
1 parent 8c51e22 commit 0e46b9c
Show file tree
Hide file tree
Showing 26 changed files with 4,327 additions and 45 deletions.
1 change: 1 addition & 0 deletions tfhe/docs/SUMMARY.md
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
* [Multi-threading with Rayon crate](guides/rayon\_crate.md)
* [Trivial ciphertexts](guides/trivial\_ciphertext.md)
* [PBS statistics](guides/pbs-stats.md)
* [Array](guides/array.md)

## Tutorials

Expand Down
78 changes: 78 additions & 0 deletions tfhe/docs/guides/array.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
# Array Types In High-level API

This document describes the array types provided by the High-level API.

```toml
# Cargo.toml

[dependencies]
tfhe = { version = "0.8.0", features = ["integer", "x86_64-unix"] }
```

```rust
use tfhe::{ConfigBuilder, generate_keys, set_server_key, CpuFheUint32Array, ClearArray};
use tfhe::prelude::*;

fn main() {
let config = ConfigBuilder::default().build();
let (cks, sks) = generate_keys(config);

set_server_key(sks);

let num_elems = 4 * 4;
let clear_xs = (0..num_elems as u32).collect::<Vec<_>>();
let clear_ys = vec![1u32; num_elems];

// Encrypted 2D array with values
// [[ 0, 1, 2, 3]
// [ 4, 5, 6, 7]
// [ 8, 9, 10, 11]
// [ 12, 13, 14, 15]]
let xs = CpuFheUint32Array::try_encrypt((clear_xs.as_slice(), vec![4, 4]), &cks).unwrap();
// Encrypted 2D array with values
// [[ 1, 1, 1, 1]
// [ 1, 1, 1, 1]
// [ 1, 1, 1, 1]
// [ 1, 1, 1, 1]]
let ys = CpuFheUint32Array::try_encrypt((clear_ys.as_slice(), vec![4, 4]), &cks).unwrap();

assert_eq!(xs.num_dim(), 2);
assert_eq!(xs.shape(), &[4, 4]);
assert_eq!(ys.num_dim(), 2);
assert_eq!(ys.shape(), &[4, 4]);

// Take a sub slice
// [[ 10, 11]
// [ 14, 15]]
let xss = xs.slice(&[2..4, 2..4]);
// Take a sub slice
// [[ 1, 1]
// [ 1, 1]]
let yss = ys.slice(&[2..4, 2..4]);

assert_eq!(xss.num_dim(), 2);
assert_eq!(xss.shape(), &[2, 2]);
assert_eq!(yss.num_dim(), 2);
assert_eq!(yss.shape(), &[2, 2]);

let r = &xss + &yss;

// Result is
// [[ 11, 12]
// [ 15, 16]]
let result: Vec<u32> = r.decrypt(&cks);
assert_eq!(result, vec![11, 12, 15, 16]);

// Clear 2D array with values
// [[ 10, 20]
// [ 30, 40]]
let clear_array = ClearArray::new(vec![10u32, 20u32, 30u32, 40u32], vec![2, 2]);
let r = &xss + &clear_array;

// Result is
// [[ 20, 31]
// [ 44, 55]]
let r: Vec<u32> = r.decrypt(&cks);
assert_eq!(r, vec![20, 31, 44, 55]);
}
```
42 changes: 0 additions & 42 deletions tfhe/src/high_level_api/array.rs

This file was deleted.

109 changes: 109 additions & 0 deletions tfhe/src/high_level_api/array/clear_ops.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
use crate::array::cpu::ClearContainer;
use crate::array::{ClearArray, DynDimensions};
use crate::high_level_api::array::traits::{
ArrayBackend, BackendDataContainer, ClearArithmeticArrayBackend, ClearBitwiseArrayBackend,
HasClear,
};
use crate::high_level_api::array::FheArrayBase;
use crate::FheId;
use std::ops::{Add, BitAnd};

macro_rules! impl_other_binary_ops_variants {
(
$trait_name:ident($trait_method:ident) => $backend_trait:ident($backend_method:ident)
) => {
impl<'a, C, Clear, Id> $trait_name<ClearArray<Clear>> for FheArrayBase<C, Id>
where
Id: FheId,
Clear: Copy,
FheArrayBase<C, Id>: HasClear<Clear = Clear>,
C: BackendDataContainer,
C::Backend: $backend_trait<Clear>,
{
type Output = FheArrayBase<<C::Backend as ArrayBackend>::Owned, Id>;

fn $trait_method(self, rhs: ClearArray<Clear>) -> Self::Output {
if !self.has_same_shape(&rhs) {
panic!("Array operands do not have the same shape");
}
let lhs_slice = self.as_tensor_slice();
let rhs_slice = rhs.as_tensor_slice().map(ClearContainer::into_inner);
let inner = C::Backend::$backend_method(lhs_slice, rhs_slice);
let resulting_shape = DynDimensions::from(self.shape().to_vec());
FheArrayBase::new(inner, resulting_shape)
}
}

impl<'a, C, Clear, Id> $trait_name<ClearArray<Clear>> for &'a FheArrayBase<C, Id>
where
Id: FheId,
Clear: Copy,

FheArrayBase<C, Id>: HasClear<Clear = Clear>,
C: BackendDataContainer,
C::Backend: $backend_trait<Clear>,
{
type Output = FheArrayBase<<C::Backend as ArrayBackend>::Owned, Id>;

fn $trait_method(self, rhs: ClearArray<Clear>) -> Self::Output {
if !self.has_same_shape(&rhs) {
panic!("Array operands do not have the same shape");
}
let lhs_slice = self.as_tensor_slice();
let rhs_slice = rhs.as_tensor_slice().map(ClearContainer::into_inner);
let inner = C::Backend::$backend_method(lhs_slice, rhs_slice);
let resulting_shape = DynDimensions::from(self.shape().to_vec());
FheArrayBase::new(inner, resulting_shape)
}
}

impl<'a, C, Clear, Id> $trait_name<&'a ClearArray<Clear>> for FheArrayBase<C, Id>
where
Id: FheId,
Clear: Copy,

FheArrayBase<C, Id>: HasClear<Clear = Clear>,
C: BackendDataContainer,
C::Backend: $backend_trait<Clear>,
{
type Output = FheArrayBase<<C::Backend as ArrayBackend>::Owned, Id>;

fn $trait_method(self, rhs: &'a ClearArray<Clear>) -> Self::Output {
if !self.has_same_shape(rhs) {
panic!("Array operands do not have the same shape");
}
let lhs_slice = self.as_tensor_slice();
let rhs_slice = rhs.as_tensor_slice().map(ClearContainer::into_inner);
let inner = C::Backend::$backend_method(lhs_slice, rhs_slice);
let resulting_shape = DynDimensions::from(self.shape().to_vec());
FheArrayBase::new(inner, resulting_shape)
}
}

impl<'a, 'b, C, Clear, Id> $trait_name<&'a ClearArray<Clear>> for &'b FheArrayBase<C, Id>
where
Id: FheId,
Clear: Copy,

FheArrayBase<C, Id>: HasClear<Clear = Clear>,
C: BackendDataContainer,
C::Backend: $backend_trait<Clear>,
{
type Output = FheArrayBase<<C::Backend as ArrayBackend>::Owned, Id>;

fn $trait_method(self, rhs: &'a ClearArray<Clear>) -> Self::Output {
if !self.has_same_shape(rhs) {
panic!("Array operands do not have the same shape");
}
let lhs_slice = self.as_tensor_slice();
let rhs_slice = rhs.as_tensor_slice().map(ClearContainer::into_inner);
let inner = C::Backend::$backend_method(lhs_slice, rhs_slice);
let resulting_shape = DynDimensions::from(self.shape().to_vec());
FheArrayBase::new(inner, resulting_shape)
}
}
};
}

impl_other_binary_ops_variants!(Add(add) => ClearArithmeticArrayBackend(add_slices));
impl_other_binary_ops_variants!(BitAnd(bitand) => ClearBitwiseArrayBackend(bitand_slice));
Loading

0 comments on commit 0e46b9c

Please sign in to comment.