chore(gpu): automatically generate rust bindings for cuda functions, …

…except device.cu
zama-ai · Oct 14, 2024 · e698d18 · e698d18
1 parent 416fb5a
commit e698d18
Show file tree

Hide file tree

Showing 104 changed files with 3,900 additions and 3,128 deletions.
diff --git a/.github/workflows/benchmark_gpu_core_crypto.yml b/.github/workflows/benchmark_gpu_core_crypto.yml
@@ -56,14 +56,15 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}
           ./bootstrap
           make -j"$(nproc)"
           sudo make install
 
+
       - name: Checkout tfhe-rs repo with tags
         uses: actions/checkout@eef61447b9ff4aafe5dcd4e0bbf5d482be7e7871
         with:

diff --git a/.github/workflows/benchmark_gpu_integer.yml b/.github/workflows/benchmark_gpu_integer.yml
@@ -59,7 +59,7 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}

diff --git a/.github/workflows/benchmark_gpu_integer_2H100_full.yml b/.github/workflows/benchmark_gpu_integer_2H100_full.yml
@@ -63,7 +63,7 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}

diff --git a/.github/workflows/benchmark_gpu_integer_full.yml b/.github/workflows/benchmark_gpu_integer_full.yml
@@ -63,7 +63,7 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}

diff --git a/.github/workflows/benchmark_gpu_integer_multi_bit.yml b/.github/workflows/benchmark_gpu_integer_multi_bit.yml
@@ -72,7 +72,7 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}

diff --git a/.github/workflows/benchmark_gpu_integer_multi_bit_multi_gpu.yml b/.github/workflows/benchmark_gpu_integer_multi_bit_multi_gpu.yml
@@ -73,7 +73,7 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}

diff --git a/.github/workflows/benchmark_gpu_integer_multi_gpu_full.yml b/.github/workflows/benchmark_gpu_integer_multi_gpu_full.yml
@@ -63,7 +63,7 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}

diff --git a/.github/workflows/benchmark_gpu_l40.yml b/.github/workflows/benchmark_gpu_l40.yml
@@ -63,7 +63,7 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}

diff --git a/.github/workflows/gpu_fast_h100_tests.yml b/.github/workflows/gpu_fast_h100_tests.yml
@@ -99,7 +99,7 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev 
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}

diff --git a/.github/workflows/gpu_fast_tests.yml b/.github/workflows/gpu_fast_tests.yml
@@ -97,7 +97,7 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev 
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}

diff --git a/.github/workflows/gpu_full_h100_tests.yml b/.github/workflows/gpu_full_h100_tests.yml
@@ -57,7 +57,7 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev 
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}

diff --git a/.github/workflows/gpu_full_multi_gpu_tests.yml b/.github/workflows/gpu_full_multi_gpu_tests.yml
@@ -99,7 +99,7 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev 
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}

diff --git a/.github/workflows/gpu_signed_integer_h100_tests.yml b/.github/workflows/gpu_signed_integer_h100_tests.yml
@@ -100,14 +100,15 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev 
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}
           ./bootstrap
           make -j"$(nproc)"
           sudo make install
 
+
       - name: Checkout tfhe-rs
         uses: actions/checkout@eef61447b9ff4aafe5dcd4e0bbf5d482be7e7871
 

diff --git a/.github/workflows/gpu_signed_integer_tests.yml b/.github/workflows/gpu_signed_integer_tests.yml
@@ -107,14 +107,15 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev 
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}
           ./bootstrap
           make -j"$(nproc)"
           sudo make install
 
+
       - name: Checkout tfhe-rs
         uses: actions/checkout@eef61447b9ff4aafe5dcd4e0bbf5d482be7e7871
         with:

diff --git a/.github/workflows/gpu_unsigned_integer_h100_tests.yml b/.github/workflows/gpu_unsigned_integer_h100_tests.yml
@@ -100,14 +100,15 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev 
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}
           ./bootstrap
           make -j"$(nproc)"
           sudo make install
 
+
       - name: Checkout tfhe-rs
         uses: actions/checkout@eef61447b9ff4aafe5dcd4e0bbf5d482be7e7871
 

diff --git a/.github/workflows/gpu_unsigned_integer_tests.yml b/.github/workflows/gpu_unsigned_integer_tests.yml
@@ -107,14 +107,15 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt update
-          sudo apt install -y checkinstall zlib1g-dev libssl-dev 
+          sudo apt install -y checkinstall zlib1g-dev libssl-dev libclang-dev
           wget https://github.com/Kitware/CMake/releases/download/v${{ env.CMAKE_VERSION }}/cmake-${{ env.CMAKE_VERSION }}.tar.gz
           tar -zxvf cmake-${{ env.CMAKE_VERSION }}.tar.gz
           cd cmake-${{ env.CMAKE_VERSION }}
           ./bootstrap
           make -j"$(nproc)"
           sudo make install
 
+
       - name: Checkout tfhe-rs
         uses: actions/checkout@eef61447b9ff4aafe5dcd4e0bbf5d482be7e7871
 

diff --git a/Makefile b/Makefile
@@ -418,6 +418,14 @@ clippy_cuda_backend: install_rs_check_toolchain
 	RUSTFLAGS="$(RUSTFLAGS)" cargo "$(CARGO_RS_CHECK_TOOLCHAIN)" clippy --all-targets \
 		-p tfhe-cuda-backend -- --no-deps -D warnings
 
+.PHONY: check_rust_bindings_did_not_change # Check rust bindings are up to date for tfhe-cuda-backend
+check_rust_bindings_did_not_change: 
+	cargo build -p tfhe-cuda-backend && \
+	git diff --quiet HEAD -- backends/tfhe-cuda-backend/src/bindings.rs || \
+	( echo "Generated bindings have changed! Please run 'git add backends/tfhe-cuda-backend/src/bindings.rs' \
+	and commit the changes." && exit 1 ) 
+
+
 .PHONY: tfhe_lints # Run custom tfhe-rs lints
 tfhe_lints: install_tfhe_lints
 	cd tfhe && RUSTFLAGS="$(RUSTFLAGS)" cargo tfhe-lints \
@@ -1257,7 +1265,7 @@ pcc: no_tfhe_typo no_dbg_log check_fmt check_typos lint_doc check_md_docs_are_te
 clippy_all tfhe_lints check_compile_tests
 
 .PHONY: pcc_gpu # pcc stands for pre commit checks for GPU compilation
-pcc_gpu: clippy_gpu clippy_cuda_backend check_compile_tests_benches_gpu
+pcc_gpu: clippy_gpu clippy_cuda_backend check_compile_tests_benches_gpu check_rust_bindings_did_not_change
 
 .PHONY: fpcc # pcc stands for pre commit checks, the f stands for fast
 fpcc: no_tfhe_typo no_dbg_log check_fmt check_typos lint_doc check_md_docs_are_tested clippy_fast \

diff --git a/backends/tfhe-cuda-backend/Cargo.toml b/backends/tfhe-cuda-backend/Cargo.toml
@@ -14,3 +14,4 @@ keywords = ["fully", "homomorphic", "encryption", "fhe", "cryptography"]
 [build-dependencies]
 cmake = { version = "0.1" }
 pkg-config = { version = "0.3" }
+bindgen = "0.70.1"
diff --git a/backends/tfhe-cuda-backend/build.rs b/backends/tfhe-cuda-backend/build.rs
@@ -1,5 +1,6 @@
-use std::env;
+use std::path::PathBuf;
 use std::process::Command;
+use std::{env, fs};
 
 fn main() {
     if let Ok(val) = env::var("DOCS_RS") {
@@ -26,6 +27,7 @@ fn main() {
     println!("cargo::rerun-if-changed=cuda/tests_and_benchmarks");
     println!("cargo::rerun-if-changed=cuda/CMakeLists.txt");
     println!("cargo::rerun-if-changed=src");
+
     if env::consts::OS == "linux" {
         let output = Command::new("./get_os_name.sh").output().unwrap();
         let distribution = String::from_utf8(output.stdout).unwrap();
@@ -35,6 +37,7 @@ fn main() {
                 Only Ubuntu is supported by tfhe-cuda-backend at this time. Build may fail\n"
             );
         }
+
         let dest = cmake::build("cuda");
         println!("cargo:rustc-link-search=native={}", dest.display());
         println!("cargo:rustc-link-lib=static=tfhe_cuda_backend");
@@ -51,6 +54,37 @@ fn main() {
         println!("cargo:rustc-link-lib=cudart");
         println!("cargo:rustc-link-search=native=/usr/lib/x86_64-linux-gnu/");
         println!("cargo:rustc-link-lib=stdc++");
+
+        let header_path = "wrapper.h";
+        println!("cargo:rerun-if-changed={}", header_path);
+
+        let out_path = PathBuf::from("src").join("bindings.rs");
+
+        // Check modification times
+        let header_modified = fs::metadata(header_path).unwrap().modified().unwrap();
+        let bindings_modified = if out_path.exists() {
+            fs::metadata(&out_path).unwrap().modified().unwrap()
+        } else {
+            std::time::SystemTime::UNIX_EPOCH // If bindings file doesn't exist, consider it older
+        };
+        // Regenerate bindings only if header has been modified
+        if header_modified > bindings_modified {
+            let bindings = bindgen::Builder::default()
+                .header(header_path)
+                .clang_arg("-x")
+                .clang_arg("c++")
+                .clang_arg("-std=c++17")
+                .clang_arg("-I/usr/include")
+                .clang_arg("-I/usr/local/include")
+                .ctypes_prefix("ffi")
+                .raw_line("use crate::ffi;")
+                .generate()
+                .expect("Unable to generate bindings");
+
+            bindings
+                .write_to_file(&out_path)
+                .expect("Couldn't write bindings!");
+        }
     } else {
         panic!(
             "Error: platform not supported, tfhe-cuda-backend not built (only Linux is supported)"

diff --git a/backends/tfhe-cuda-backend/cuda/include/ciphertext.h b/backends/tfhe-cuda-backend/cuda/include/ciphertext.h
@@ -1,25 +1,24 @@
 #ifndef CUDA_CIPHERTEXT_H
 #define CUDA_CIPHERTEXT_H
 
-#include "device.h"
-#include <cstdint>
+#include "stdint.h"
 
 extern "C" {
 void cuda_convert_lwe_ciphertext_vector_to_gpu_64(void *stream,
                                                   uint32_t gpu_index,
-                                                  void *dest, void *src,
+                                                  void *dest, void const *src,
                                                   uint32_t number_of_cts,
                                                   uint32_t lwe_dimension);
 void cuda_convert_lwe_ciphertext_vector_to_cpu_64(void *stream,
                                                   uint32_t gpu_index,
-                                                  void *dest, void *src,
+                                                  void *dest, void const *src,
                                                   uint32_t number_of_cts,
                                                   uint32_t lwe_dimension);
 
 void cuda_glwe_sample_extract_64(void *stream, uint32_t gpu_index,
-                                 void *lwe_array_out, void *glwe_array_in,
-                                 uint32_t *nth_array, uint32_t num_nths,
+                                 void *lwe_array_out, void const *glwe_array_in,
+                                 uint32_t const *nth_array, uint32_t num_nths,
                                  uint32_t glwe_dimension,
                                  uint32_t polynomial_size);
-};
+}
 #endif
diff --git a/backends/tfhe-cuda-backend/cuda/include/device.h b/backends/tfhe-cuda-backend/cuda/include/device.h
@@ -42,7 +42,7 @@ void cuda_check_valid_malloc(uint64_t size, uint32_t gpu_index);
 void cuda_memcpy_async_to_gpu(void *dest, void *src, uint64_t size,
                               cudaStream_t stream, uint32_t gpu_index);
 
-void cuda_memcpy_async_gpu_to_gpu(void *dest, void *src, uint64_t size,
+void cuda_memcpy_async_gpu_to_gpu(void *dest, void const *src, uint64_t size,
                                   cudaStream_t stream, uint32_t gpu_index);
 
 void cuda_memcpy_gpu_to_gpu(void *dest, void *src, uint64_t size,

diff --git a/backends/tfhe-cuda-backend/cuda/include/integer/compression/compression.h b/backends/tfhe-cuda-backend/cuda/include/integer/compression/compression.h
@@ -0,0 +1,45 @@
+#ifndef CUDA_INTEGER_COMPRESSION_H
+#define CUDA_INTEGER_COMPRESSION_H
+
+#include "../../pbs/pbs_enums.h"
+
+extern "C" {
+void scratch_cuda_integer_compress_radix_ciphertext_64(
+    void *const *streams, uint32_t const *gpu_indexes, uint32_t gpu_count,
+    int8_t **mem_ptr, uint32_t compression_glwe_dimension,
+    uint32_t compression_polynomial_size, uint32_t lwe_dimension,
+    uint32_t ks_level, uint32_t ks_base_log, uint32_t num_radix_blocks,
+    uint32_t message_modulus, uint32_t carry_modulus, PBS_TYPE pbs_type,
+    uint32_t lwe_per_glwe, uint32_t storage_log_modulus,
+    bool allocate_gpu_memory);
+
+void scratch_cuda_integer_decompress_radix_ciphertext_64(
+    void *const *streams, uint32_t const *gpu_indexes, uint32_t gpu_count,
+    int8_t **mem_ptr, uint32_t encryption_glwe_dimension,
+    uint32_t encryption_polynomial_size, uint32_t compression_glwe_dimension,
+    uint32_t compression_polynomial_size, uint32_t lwe_dimension,
+    uint32_t pbs_level, uint32_t pbs_base_log, uint32_t num_radix_blocks,
+    uint32_t message_modulus, uint32_t carry_modulus, PBS_TYPE pbs_type,
+    uint32_t storage_log_modulus, uint32_t body_count,
+    bool allocate_gpu_memory);
+
+void cuda_integer_compress_radix_ciphertext_64(
+    void *const *streams, uint32_t const *gpu_indexes, uint32_t gpu_count,
+    void *glwe_array_out, void const *lwe_array_in, void *const *fp_ksk,
+    uint32_t num_nths, int8_t *mem_ptr);
+
+void cuda_integer_decompress_radix_ciphertext_64(
+    void *const *streams, uint32_t const *gpu_indexes, uint32_t gpu_count,
+    void *lwe_array_out, void const *glwe_in, uint32_t const *indexes_array,
+    uint32_t indexes_array_size, void *const *bsks, int8_t *mem_ptr);
+
+void cleanup_cuda_integer_compress_radix_ciphertext_64(
+    void *const *streams, uint32_t const *gpu_indexes, uint32_t gpu_count,
+    int8_t **mem_ptr_void);
+
+void cleanup_cuda_integer_decompress_radix_ciphertext_64(
+    void *const *streams, uint32_t const *gpu_indexes, uint32_t gpu_count,
+    int8_t **mem_ptr_void);
+}
+
+#endif