Skip to content

Commit f6acec4

Browse files
pdroalvespdroalves
committed
feat(gpu): add DISPATCH_POLY_SIZE macro to replace manual switch dispatches
1 parent d36cc09 commit f6acec4

16 files changed

Lines changed: 444 additions & 2038 deletions

‎backends/tfhe-cuda-backend/cuda/src/crypto/ciphertext.cu‎

Lines changed: 15 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#include "ciphertext.cuh"
2+
#include "polynomial/dispatch.cuh"
23
#include "polynomial/parameters.cuh"
34

45
void cuda_convert_lwe_ciphertext_vector_to_gpu_64_async(
@@ -23,53 +24,13 @@ void cuda_glwe_sample_extract_64_async(
2324
uint32_t num_lwes_to_extract_per_glwe, uint32_t num_lwes_stored_per_glwe,
2425
uint32_t glwe_dimension, uint32_t polynomial_size) {
2526

26-
switch (polynomial_size) {
27-
case 256:
28-
host_sample_extract<uint64_t, AmortizedDegree<256>>(
29-
static_cast<cudaStream_t>(stream), gpu_index, (uint64_t *)lwe_array_out,
30-
(uint64_t const *)glwe_array_in, (uint32_t const *)nth_array, num_nths,
31-
num_lwes_to_extract_per_glwe, num_lwes_stored_per_glwe, glwe_dimension);
32-
break;
33-
case 512:
34-
host_sample_extract<uint64_t, AmortizedDegree<512>>(
35-
static_cast<cudaStream_t>(stream), gpu_index, (uint64_t *)lwe_array_out,
36-
(uint64_t const *)glwe_array_in, (uint32_t const *)nth_array, num_nths,
37-
num_lwes_to_extract_per_glwe, num_lwes_stored_per_glwe, glwe_dimension);
38-
break;
39-
case 1024:
40-
host_sample_extract<uint64_t, AmortizedDegree<1024>>(
41-
static_cast<cudaStream_t>(stream), gpu_index, (uint64_t *)lwe_array_out,
42-
(uint64_t const *)glwe_array_in, (uint32_t const *)nth_array, num_nths,
43-
num_lwes_to_extract_per_glwe, num_lwes_stored_per_glwe, glwe_dimension);
44-
break;
45-
case 2048:
46-
host_sample_extract<uint64_t, AmortizedDegree<2048>>(
47-
static_cast<cudaStream_t>(stream), gpu_index, (uint64_t *)lwe_array_out,
48-
(uint64_t const *)glwe_array_in, (uint32_t const *)nth_array, num_nths,
49-
num_lwes_to_extract_per_glwe, num_lwes_stored_per_glwe, glwe_dimension);
50-
break;
51-
case 4096:
52-
host_sample_extract<uint64_t, AmortizedDegree<4096>>(
53-
static_cast<cudaStream_t>(stream), gpu_index, (uint64_t *)lwe_array_out,
54-
(uint64_t const *)glwe_array_in, (uint32_t const *)nth_array, num_nths,
55-
num_lwes_to_extract_per_glwe, num_lwes_stored_per_glwe, glwe_dimension);
56-
break;
57-
case 8192:
58-
host_sample_extract<uint64_t, AmortizedDegree<8192>>(
59-
static_cast<cudaStream_t>(stream), gpu_index, (uint64_t *)lwe_array_out,
60-
(uint64_t const *)glwe_array_in, (uint32_t const *)nth_array, num_nths,
61-
num_lwes_to_extract_per_glwe, num_lwes_stored_per_glwe, glwe_dimension);
62-
break;
63-
case 16384:
64-
host_sample_extract<uint64_t, AmortizedDegree<16384>>(
65-
static_cast<cudaStream_t>(stream), gpu_index, (uint64_t *)lwe_array_out,
66-
(uint64_t const *)glwe_array_in, (uint32_t const *)nth_array, num_nths,
67-
num_lwes_to_extract_per_glwe, num_lwes_stored_per_glwe, glwe_dimension);
68-
break;
69-
default:
70-
PANIC("Cuda error: unsupported polynomial size. Supported "
71-
"N's are powers of two in the interval [256..16384].")
72-
}
27+
DISPATCH_POLY_SIZE(
28+
polynomial_size, AmortizedDegreePolicy,
29+
host_sample_extract<uint64_t, Params>(
30+
static_cast<cudaStream_t>(stream), gpu_index,
31+
(uint64_t *)lwe_array_out, (uint64_t const *)glwe_array_in,
32+
(uint32_t const *)nth_array, num_nths, num_lwes_to_extract_per_glwe,
33+
num_lwes_stored_per_glwe, glwe_dimension));
7334
}
7435

7536
void cuda_modulus_switch_inplace_64_async(void *stream, uint32_t gpu_index,
@@ -121,46 +82,13 @@ void cuda_glwe_sample_extract_128_async(
12182
uint32_t num_lwes_to_extract_per_glwe, uint32_t num_lwes_stored_per_glwe,
12283
uint32_t glwe_dimension, uint32_t polynomial_size) {
12384

124-
switch (polynomial_size) {
125-
case 256:
126-
host_sample_extract<__uint128_t, AmortizedDegree<256>>(
127-
static_cast<cudaStream_t>(stream), gpu_index,
128-
(__uint128_t *)lwe_array_out, (__uint128_t const *)glwe_array_in,
129-
(uint32_t const *)nth_array, num_nths, num_lwes_to_extract_per_glwe,
130-
num_lwes_stored_per_glwe, glwe_dimension);
131-
break;
132-
case 512:
133-
host_sample_extract<__uint128_t, AmortizedDegree<512>>(
134-
static_cast<cudaStream_t>(stream), gpu_index,
135-
(__uint128_t *)lwe_array_out, (__uint128_t const *)glwe_array_in,
136-
(uint32_t const *)nth_array, num_nths, num_lwes_to_extract_per_glwe,
137-
num_lwes_stored_per_glwe, glwe_dimension);
138-
break;
139-
case 1024:
140-
host_sample_extract<__uint128_t, AmortizedDegree<1024>>(
141-
static_cast<cudaStream_t>(stream), gpu_index,
142-
(__uint128_t *)lwe_array_out, (__uint128_t const *)glwe_array_in,
143-
(uint32_t const *)nth_array, num_nths, num_lwes_to_extract_per_glwe,
144-
num_lwes_stored_per_glwe, glwe_dimension);
145-
break;
146-
case 2048:
147-
host_sample_extract<__uint128_t, AmortizedDegree<2048>>(
148-
static_cast<cudaStream_t>(stream), gpu_index,
149-
(__uint128_t *)lwe_array_out, (__uint128_t const *)glwe_array_in,
150-
(uint32_t const *)nth_array, num_nths, num_lwes_to_extract_per_glwe,
151-
num_lwes_stored_per_glwe, glwe_dimension);
152-
break;
153-
case 4096:
154-
host_sample_extract<__uint128_t, AmortizedDegree<4096>>(
155-
static_cast<cudaStream_t>(stream), gpu_index,
156-
(__uint128_t *)lwe_array_out, (__uint128_t const *)glwe_array_in,
157-
(uint32_t const *)nth_array, num_nths, num_lwes_to_extract_per_glwe,
158-
num_lwes_stored_per_glwe, glwe_dimension);
159-
break;
160-
default:
161-
PANIC("Cuda error: unsupported polynomial size. Supported "
162-
"N's are powers of two in the interval [256..4096].")
163-
}
85+
DISPATCH_POLY_SIZE(
86+
polynomial_size, AmortizedDegreePolicy128,
87+
host_sample_extract<__uint128_t, Params>(
88+
static_cast<cudaStream_t>(stream), gpu_index,
89+
(__uint128_t *)lwe_array_out, (__uint128_t const *)glwe_array_in,
90+
(uint32_t const *)nth_array, num_nths, num_lwes_to_extract_per_glwe,
91+
num_lwes_stored_per_glwe, glwe_dimension));
16492
}
16593

16694
void cuda_modulus_switch_multi_bit_64_async(void *stream, uint32_t gpu_index,

‎backends/tfhe-cuda-backend/cuda/src/fft128/fft128.cu‎

Lines changed: 19 additions & 144 deletions
Original file line numberDiff line numberDiff line change
@@ -1,163 +1,38 @@
11
#include "fft128.cuh"
2+
#include "polynomial/dispatch.cuh"
23

34
void cuda_fourier_transform_forward_as_integer_f128_async(
45
void *stream, uint32_t gpu_index, void *re0, void *re1, void *im0,
56
void *im1, void const *standard, const uint32_t N,
67
const uint32_t number_of_samples) {
7-
switch (N) {
8-
case 64:
9-
host_fourier_transform_forward_as_integer_f128<AmortizedDegree<64>>(
10-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
11-
(double *)re1, (double *)im0, (double *)im1,
12-
(__uint128_t const *)standard, N, number_of_samples);
13-
break;
14-
case 128:
15-
host_fourier_transform_forward_as_integer_f128<AmortizedDegree<128>>(
16-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
17-
(double *)re1, (double *)im0, (double *)im1,
18-
(__uint128_t const *)standard, N, number_of_samples);
19-
break;
20-
case 256:
21-
host_fourier_transform_forward_as_integer_f128<AmortizedDegree<256>>(
22-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
23-
(double *)re1, (double *)im0, (double *)im1,
24-
(__uint128_t const *)standard, N, number_of_samples);
25-
break;
26-
case 512:
27-
host_fourier_transform_forward_as_integer_f128<AmortizedDegree<512>>(
28-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
29-
(double *)re1, (double *)im0, (double *)im1,
30-
(__uint128_t const *)standard, N, number_of_samples);
31-
break;
32-
case 1024:
33-
host_fourier_transform_forward_as_integer_f128<AmortizedDegree<1024>>(
34-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
35-
(double *)re1, (double *)im0, (double *)im1,
36-
(__uint128_t const *)standard, N, number_of_samples);
37-
break;
38-
case 2048:
39-
host_fourier_transform_forward_as_integer_f128<AmortizedDegree<2048>>(
40-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
41-
(double *)re1, (double *)im0, (double *)im1,
42-
(__uint128_t const *)standard, N, number_of_samples);
43-
break;
44-
case 4096:
45-
host_fourier_transform_forward_as_integer_f128<AmortizedDegree<4096>>(
46-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
47-
(double *)re1, (double *)im0, (double *)im1,
48-
(__uint128_t const *)standard, N, number_of_samples);
49-
break;
50-
default:
51-
PANIC("Cuda error (f128 fft): unsupported polynomial size. Supported "
52-
"N's are powers of two"
53-
" in the interval [64..4096].")
54-
}
8+
DISPATCH_POLY_SIZE(N, AmortizedDegreePolicyFFT128,
9+
host_fourier_transform_forward_as_integer_f128<Params>(
10+
static_cast<cudaStream_t>(stream), gpu_index,
11+
(double *)re0, (double *)re1, (double *)im0,
12+
(double *)im1, (__uint128_t const *)standard, N,
13+
number_of_samples));
5514
}
5615

5716
void cuda_fourier_transform_forward_as_torus_f128_async(
5817
void *stream, uint32_t gpu_index, void *re0, void *re1, void *im0,
5918
void *im1, void const *standard, const uint32_t N,
6019
const uint32_t number_of_samples) {
61-
switch (N) {
62-
case 64:
63-
host_fourier_transform_forward_as_torus_f128<AmortizedDegree<64>>(
64-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
65-
(double *)re1, (double *)im0, (double *)im1,
66-
(__uint128_t const *)standard, N, number_of_samples);
67-
break;
68-
case 128:
69-
host_fourier_transform_forward_as_torus_f128<AmortizedDegree<128>>(
70-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
71-
(double *)re1, (double *)im0, (double *)im1,
72-
(__uint128_t const *)standard, N, number_of_samples);
73-
break;
74-
case 256:
75-
host_fourier_transform_forward_as_torus_f128<AmortizedDegree<256>>(
76-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
77-
(double *)re1, (double *)im0, (double *)im1,
78-
(__uint128_t const *)standard, N, number_of_samples);
79-
break;
80-
case 512:
81-
host_fourier_transform_forward_as_torus_f128<AmortizedDegree<512>>(
82-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
83-
(double *)re1, (double *)im0, (double *)im1,
84-
(__uint128_t const *)standard, N, number_of_samples);
85-
break;
86-
case 1024:
87-
host_fourier_transform_forward_as_torus_f128<AmortizedDegree<1024>>(
88-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
89-
(double *)re1, (double *)im0, (double *)im1,
90-
(__uint128_t const *)standard, N, number_of_samples);
91-
break;
92-
case 2048:
93-
host_fourier_transform_forward_as_torus_f128<AmortizedDegree<2048>>(
94-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
95-
(double *)re1, (double *)im0, (double *)im1,
96-
(__uint128_t const *)standard, N, number_of_samples);
97-
break;
98-
case 4096:
99-
host_fourier_transform_forward_as_torus_f128<AmortizedDegree<4096>>(
100-
static_cast<cudaStream_t>(stream), gpu_index, (double *)re0,
101-
(double *)re1, (double *)im0, (double *)im1,
102-
(__uint128_t const *)standard, N, number_of_samples);
103-
break;
104-
default:
105-
PANIC("Cuda error (f128 fft): unsupported polynomial size. Supported "
106-
"N's are powers of two"
107-
" in the interval [64..4096].")
108-
}
20+
DISPATCH_POLY_SIZE(N, AmortizedDegreePolicyFFT128,
21+
host_fourier_transform_forward_as_torus_f128<Params>(
22+
static_cast<cudaStream_t>(stream), gpu_index,
23+
(double *)re0, (double *)re1, (double *)im0,
24+
(double *)im1, (__uint128_t const *)standard, N,
25+
number_of_samples));
10926
}
11027

11128
void cuda_fourier_transform_backward_as_torus_f128_async(
11229
void *stream, uint32_t gpu_index, void *standard, void const *re0,
11330
void const *re1, void const *im0, void const *im1, const uint32_t N,
11431
const uint32_t number_of_samples) {
115-
switch (N) {
116-
case 64:
117-
host_fourier_transform_backward_as_torus_f128<AmortizedDegree<64>>(
118-
static_cast<cudaStream_t>(stream), gpu_index, (__uint128_t *)standard,
119-
(double const *)re0, (double const *)re1, (double const *)im0,
120-
(double const *)im1, N, number_of_samples);
121-
break;
122-
case 128:
123-
host_fourier_transform_backward_as_torus_f128<AmortizedDegree<128>>(
124-
static_cast<cudaStream_t>(stream), gpu_index, (__uint128_t *)standard,
125-
(double const *)re0, (double const *)re1, (double const *)im0,
126-
(double const *)im1, N, number_of_samples);
127-
break;
128-
case 256:
129-
host_fourier_transform_backward_as_torus_f128<AmortizedDegree<256>>(
130-
static_cast<cudaStream_t>(stream), gpu_index, (__uint128_t *)standard,
131-
(double const *)re0, (double const *)re1, (double const *)im0,
132-
(double const *)im1, N, number_of_samples);
133-
break;
134-
case 512:
135-
host_fourier_transform_backward_as_torus_f128<AmortizedDegree<512>>(
136-
static_cast<cudaStream_t>(stream), gpu_index, (__uint128_t *)standard,
137-
(double const *)re0, (double const *)re1, (double const *)im0,
138-
(double const *)im1, N, number_of_samples);
139-
break;
140-
case 1024:
141-
host_fourier_transform_backward_as_torus_f128<AmortizedDegree<1024>>(
142-
static_cast<cudaStream_t>(stream), gpu_index, (__uint128_t *)standard,
143-
(double const *)re0, (double const *)re1, (double const *)im0,
144-
(double const *)im1, N, number_of_samples);
145-
break;
146-
case 2048:
147-
host_fourier_transform_backward_as_torus_f128<AmortizedDegree<2048>>(
148-
static_cast<cudaStream_t>(stream), gpu_index, (__uint128_t *)standard,
149-
(double const *)re0, (double const *)re1, (double const *)im0,
150-
(double const *)im1, N, number_of_samples);
151-
break;
152-
case 4096:
153-
host_fourier_transform_backward_as_torus_f128<AmortizedDegree<4096>>(
154-
static_cast<cudaStream_t>(stream), gpu_index, (__uint128_t *)standard,
155-
(double const *)re0, (double const *)re1, (double const *)im0,
156-
(double const *)im1, N, number_of_samples);
157-
break;
158-
default:
159-
PANIC("Cuda error (f128 ifft): unsupported polynomial size. Supported "
160-
"N's are powers of two"
161-
" in the interval [64..4096].")
162-
}
32+
DISPATCH_POLY_SIZE(N, AmortizedDegreePolicyFFT128,
33+
host_fourier_transform_backward_as_torus_f128<Params>(
34+
static_cast<cudaStream_t>(stream), gpu_index,
35+
(__uint128_t *)standard, (double const *)re0,
36+
(double const *)re1, (double const *)im0,
37+
(double const *)im1, N, number_of_samples));
16338
}

‎backends/tfhe-cuda-backend/cuda/src/integer/multiplication.cu‎

Lines changed: 14 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#include "integer/multiplication.cuh"
2+
#include "polynomial/dispatch.cuh"
23

34
/*
45
* when adding chunk_size times terms together, there might be some blocks
@@ -69,53 +70,12 @@ void cuda_integer_mult_inplace_64_async(
6970
// In-place variant: radix_lwe_inout *= radix_lwe_right, no aliasing check
7071
// needed
7172
PUSH_RANGE("mul_inplace")
72-
switch (polynomial_size) {
73-
case 256:
74-
host_integer_mult_radix<uint64_t, AmortizedDegree<256>>(
75-
CudaStreams(streams), radix_lwe_inout, radix_lwe_inout, is_bool_left,
76-
radix_lwe_right, is_bool_right, bsks, (uint64_t **)(ksks),
77-
(int_mul_memory<uint64_t> *)mem_ptr, num_blocks);
78-
break;
79-
case 512:
80-
host_integer_mult_radix<uint64_t, AmortizedDegree<512>>(
81-
CudaStreams(streams), radix_lwe_inout, radix_lwe_inout, is_bool_left,
82-
radix_lwe_right, is_bool_right, bsks, (uint64_t **)(ksks),
83-
(int_mul_memory<uint64_t> *)mem_ptr, num_blocks);
84-
break;
85-
case 1024:
86-
host_integer_mult_radix<uint64_t, AmortizedDegree<1024>>(
87-
CudaStreams(streams), radix_lwe_inout, radix_lwe_inout, is_bool_left,
88-
radix_lwe_right, is_bool_right, bsks, (uint64_t **)(ksks),
89-
(int_mul_memory<uint64_t> *)mem_ptr, num_blocks);
90-
break;
91-
case 2048:
92-
host_integer_mult_radix<uint64_t, AmortizedDegree<2048>>(
93-
CudaStreams(streams), radix_lwe_inout, radix_lwe_inout, is_bool_left,
94-
radix_lwe_right, is_bool_right, bsks, (uint64_t **)(ksks),
95-
(int_mul_memory<uint64_t> *)mem_ptr, num_blocks);
96-
break;
97-
case 4096:
98-
host_integer_mult_radix<uint64_t, AmortizedDegree<4096>>(
99-
CudaStreams(streams), radix_lwe_inout, radix_lwe_inout, is_bool_left,
100-
radix_lwe_right, is_bool_right, bsks, (uint64_t **)(ksks),
101-
(int_mul_memory<uint64_t> *)mem_ptr, num_blocks);
102-
break;
103-
case 8192:
104-
host_integer_mult_radix<uint64_t, AmortizedDegree<8192>>(
105-
CudaStreams(streams), radix_lwe_inout, radix_lwe_inout, is_bool_left,
106-
radix_lwe_right, is_bool_right, bsks, (uint64_t **)(ksks),
107-
(int_mul_memory<uint64_t> *)mem_ptr, num_blocks);
108-
break;
109-
case 16384:
110-
host_integer_mult_radix<uint64_t, AmortizedDegree<16384>>(
111-
CudaStreams(streams), radix_lwe_inout, radix_lwe_inout, is_bool_left,
112-
radix_lwe_right, is_bool_right, bsks, (uint64_t **)(ksks),
113-
(int_mul_memory<uint64_t> *)mem_ptr, num_blocks);
114-
break;
115-
default:
116-
PANIC("Cuda error (integer multiplication): unsupported polynomial size. "
117-
"Supported N's are powers of two in the interval [256..16384].")
118-
}
73+
DISPATCH_POLY_SIZE(polynomial_size, AmortizedDegreePolicy,
74+
host_integer_mult_radix<uint64_t, Params>(
75+
CudaStreams(streams), radix_lwe_inout, radix_lwe_inout,
76+
is_bool_left, radix_lwe_right, is_bool_right, bsks,
77+
(uint64_t **)(ksks),
78+
(int_mul_memory<uint64_t> *)mem_ptr, num_blocks));
11979
POP_RANGE()
12080
}
12181

@@ -129,22 +89,15 @@ uint64_t scratch_cuda_integer_mult_inplace_64_async(
12989
int_radix_params params(bsk_params, ksk_params, message_modulus,
13090
carry_modulus, noise_reduction_type);
13191

132-
switch (polynomial_size) {
133-
case 256:
134-
case 512:
135-
case 1024:
136-
case 2048:
137-
case 4096:
138-
case 8192:
139-
case 16384:
140-
return scratch_cuda_integer_mult_radix_ciphertext<uint64_t>(
141-
CudaStreams(streams), (int_mul_memory<uint64_t> **)mem_ptr,
142-
is_boolean_left, is_boolean_right, num_radix_blocks, params,
143-
allocate_gpu_memory);
144-
default:
92+
if (polynomial_size < 256 || polynomial_size > 16384 ||
93+
(polynomial_size & (polynomial_size - 1)) != 0)
14594
PANIC("Cuda error (integer multiplication): unsupported polynomial size. "
14695
"Supported N's are powers of two in the interval [256..16384].")
147-
}
96+
97+
return scratch_cuda_integer_mult_radix_ciphertext<uint64_t>(
98+
CudaStreams(streams), (int_mul_memory<uint64_t> **)mem_ptr,
99+
is_boolean_left, is_boolean_right, num_radix_blocks, params,
100+
allocate_gpu_memory);
148101
}
149102

150103
void cleanup_cuda_integer_mult_inplace_64(CudaStreamsFFI streams,

0 commit comments

Comments
 (0)