This document describes how to use Rayon for parallel processing in TFHE-rs, detailing configurations for single and multi-client applications with code examples.
Rayon is a popular Rust crate that simplifies writing multi-threaded code. You can use Rayon to write multi-threaded TFHE-rs code. However, due to the specifications of Rayon and TFHE-rs, certain setups are necessary.
Single-client application
The problem
The high-level API requires to call set_server_key on each thread where computations need to be done. So a first attempt to use Rayon with TFHE-rs might look like this:
use tfhe::prelude::*;use tfhe::{ConfigBuilder, set_server_key, FheUint8, generate_keys};fnmain() {let (cks, sks) =generate_keys(ConfigBuilder::default());let xs = [FheUint8::encrypt(1u8, &cks),FheUint8::encrypt(2u8, &cks), ];let ys = [FheUint8::encrypt(3u8, &cks),FheUint8::encrypt(4u8, &cks), ];// set_server_key in each closure as they might be// running in different threadslet (a, b) = rayon::join(|| {set_server_key(sks.clone());&xs[0] +&ys[0] },|| {set_server_key(sks.clone());&xs[1] +&ys[1] } );}
However, due to Rayon's work-stealing mechanism and TFHE-rs' internals, this may create BorrowMutError.
Working example
The correct way is to call rayon::broadcast as follows:
use tfhe::prelude::*;use tfhe::{ConfigBuilder, set_server_key, FheUint8, generate_keys};fnmain() {let (cks, sks) =generate_keys(ConfigBuilder::default());// set the server key in all of the rayon's threads so that// we won't need to do it later rayon::broadcast(|_|set_server_key(sks.clone()));// Set the server key in the main threadset_server_key(sks);let xs = [FheUint8::encrypt(1u8, &cks),FheUint8::encrypt(2u8, &cks), ];let ys = [FheUint8::encrypt(3u8, &cks),FheUint8::encrypt(4u8, &cks), ];let (a, b) = rayon::join(|| {&xs[0] +&ys[0] },|| {&xs[1] +&ys[1] } );let a:u8= a.decrypt(&cks);let b:u8= b.decrypt(&cks);assert_eq!(a, 4u8);assert_eq!(b, 6u8);}
Multi-client applications
For applications that need to operate concurrently on data from different clients and require each client to use multiple threads, you need to create separate Rayon thread pools: