mirror of
https://github.com/zama-ai/concrete.git
synced 2026-02-08 19:44:57 -05:00
637 lines
24 KiB
C++
637 lines
24 KiB
C++
// Part of the Concrete Compiler Project, under the BSD3 License with Zama
|
|
// Exceptions. See
|
|
// https://github.com/zama-ai/concrete-compiler-internal/blob/main/LICENSE.txt
|
|
// for license information.
|
|
|
|
#include <atomic>
|
|
#include <cstdarg>
|
|
#include <iostream>
|
|
#include <list>
|
|
#include <memory>
|
|
#include <numeric>
|
|
#include <queue>
|
|
#include <thread>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include <concretelang/ClientLib/Types.h>
|
|
#include <concretelang/Runtime/stream_emulator_api.h>
|
|
#include <concretelang/Runtime/wrappers.h>
|
|
|
|
#ifdef CONCRETELANG_CUDA_SUPPORT
|
|
#include "bootstrap.h"
|
|
#include "device.h"
|
|
#include "keyswitch.h"
|
|
#include "linear_algebra.h"
|
|
|
|
using MemRef2 = concretelang::clientlib::MemRefDescriptor<2>;
|
|
using RuntimeContext = mlir::concretelang::RuntimeContext;
|
|
|
|
namespace mlir {
|
|
namespace concretelang {
|
|
namespace gpu_dfg {
|
|
namespace {
|
|
|
|
static std::atomic<size_t> next_device = {0};
|
|
static size_t num_devices = 0;
|
|
|
|
struct Void {};
|
|
union Param {
|
|
Void _;
|
|
uint32_t val;
|
|
};
|
|
union Context {
|
|
Void _;
|
|
RuntimeContext *val;
|
|
};
|
|
static const int32_t host_location = -1;
|
|
struct Stream;
|
|
struct Dependence;
|
|
struct PBS_buffer {
|
|
PBS_buffer(void *stream, uint32_t gpu_idx, uint32_t glwe_dimension,
|
|
uint32_t polynomial_size, uint32_t input_lwe_ciphertext_count)
|
|
: max_pbs_buffer_samples(input_lwe_ciphertext_count),
|
|
glwe_dim(glwe_dimension), poly_size(polynomial_size),
|
|
gpu_stream(stream), gpu_index(gpu_idx) {
|
|
scratch_cuda_bootstrap_amortized_64(
|
|
gpu_stream, gpu_index, &pbs_buffer, glwe_dim, poly_size,
|
|
max_pbs_buffer_samples, cuda_get_max_shared_memory(gpu_index), true);
|
|
}
|
|
~PBS_buffer() {
|
|
assert(pbs_buffer != nullptr);
|
|
cleanup_cuda_bootstrap_amortized(gpu_stream, gpu_index, &pbs_buffer);
|
|
}
|
|
int8_t *get_pbs_buffer(void *stream, uint32_t gpu_idx,
|
|
uint32_t glwe_dimension, uint32_t polynomial_size,
|
|
uint32_t input_lwe_ciphertext_count) {
|
|
assert(glwe_dimension == glwe_dim);
|
|
assert(polynomial_size == poly_size);
|
|
assert(input_lwe_ciphertext_count <= max_pbs_buffer_samples);
|
|
assert(stream == gpu_stream);
|
|
assert(gpu_idx == gpu_index);
|
|
assert(pbs_buffer != nullptr);
|
|
return pbs_buffer;
|
|
}
|
|
|
|
private:
|
|
int8_t *pbs_buffer;
|
|
uint32_t max_pbs_buffer_samples;
|
|
uint32_t glwe_dim;
|
|
uint32_t poly_size;
|
|
void *gpu_stream;
|
|
uint32_t gpu_index;
|
|
};
|
|
struct GPU_DFG {
|
|
uint32_t gpu_idx;
|
|
void *gpu_stream;
|
|
GPU_DFG(uint32_t idx) : gpu_idx(idx), pbs_buffer(nullptr) {
|
|
gpu_stream = cuda_create_stream(idx);
|
|
}
|
|
~GPU_DFG() {
|
|
if (pbs_buffer != nullptr)
|
|
delete pbs_buffer;
|
|
free_streams();
|
|
cuda_destroy_stream((cudaStream_t *)gpu_stream, gpu_idx);
|
|
free_stream_order_dependent_data();
|
|
}
|
|
inline void register_stream(Stream *s) { streams.push_back(s); }
|
|
inline void register_stream_order_dependent_allocation(void *p) {
|
|
to_free_list.push_back(p);
|
|
}
|
|
inline void free_stream_order_dependent_data() {
|
|
for (auto p : to_free_list)
|
|
free(p);
|
|
to_free_list.clear();
|
|
}
|
|
inline int8_t *get_pbs_buffer(uint32_t glwe_dimension,
|
|
uint32_t polynomial_size,
|
|
uint32_t input_lwe_ciphertext_count) {
|
|
if (pbs_buffer == nullptr)
|
|
pbs_buffer = new PBS_buffer(gpu_stream, gpu_idx, glwe_dimension,
|
|
polynomial_size, input_lwe_ciphertext_count);
|
|
return pbs_buffer->get_pbs_buffer(gpu_stream, gpu_idx, glwe_dimension,
|
|
polynomial_size,
|
|
input_lwe_ciphertext_count);
|
|
}
|
|
void drop_pbs_buffer() {
|
|
delete pbs_buffer;
|
|
pbs_buffer = nullptr;
|
|
}
|
|
void free_streams();
|
|
|
|
private:
|
|
std::list<void *> to_free_list;
|
|
std::list<Stream *> streams;
|
|
PBS_buffer *pbs_buffer;
|
|
};
|
|
struct Dependence {
|
|
int32_t location;
|
|
int32_t rank;
|
|
MemRef2 host_data;
|
|
void *device_data;
|
|
bool onHostReady;
|
|
bool hostAllocated;
|
|
bool used;
|
|
bool read;
|
|
Dependence(int32_t l, int32_t r, MemRef2 &hd, void *dd, bool ohr,
|
|
bool alloc = false)
|
|
: location(l), rank(r), host_data(hd), device_data(dd), onHostReady(ohr),
|
|
hostAllocated(alloc), used(false), read(false) {}
|
|
Dependence(int32_t l, int32_t r, uint64_t val, void *dd, bool ohr,
|
|
bool alloc = false)
|
|
: location(l), rank(r), device_data(dd), onHostReady(ohr),
|
|
hostAllocated(alloc), used(false), read(false) {
|
|
*host_data.aligned = val;
|
|
}
|
|
inline void free_data(GPU_DFG *dfg) {
|
|
if (location >= 0) {
|
|
cuda_drop_async(device_data, (cudaStream_t *)dfg->gpu_stream, location);
|
|
}
|
|
if (onHostReady && host_data.allocated != nullptr && hostAllocated) {
|
|
// As streams are not synchronized aside from the GET operation,
|
|
// we cannot free host-side data until after the synchronization
|
|
// point as it could still be used by an asynchronous operation.
|
|
dfg->register_stream_order_dependent_allocation(host_data.allocated);
|
|
}
|
|
delete (this);
|
|
}
|
|
};
|
|
struct Process {
|
|
std::vector<Stream *> input_streams;
|
|
std::vector<Stream *> output_streams;
|
|
GPU_DFG *dfg;
|
|
Param level;
|
|
Param base_log;
|
|
Param input_lwe_dim;
|
|
Param output_lwe_dim;
|
|
Param poly_size;
|
|
Param glwe_dim;
|
|
Param precision;
|
|
Param output_size;
|
|
Context ctx;
|
|
void (*fun)(Process *);
|
|
char name[80];
|
|
};
|
|
|
|
static inline void schedule_kernel(Process *p) { p->fun(p); }
|
|
struct Stream {
|
|
stream_type type;
|
|
Dependence *dep;
|
|
Process *producer;
|
|
std::vector<Process *> consumers;
|
|
GPU_DFG *dfg;
|
|
Stream(stream_type t)
|
|
: type(t), dep(nullptr), producer(nullptr), dfg(nullptr) {}
|
|
~Stream() {
|
|
if (dep != nullptr)
|
|
dep->free_data(dfg);
|
|
delete producer;
|
|
}
|
|
void put(Dependence *d) {
|
|
if (type == TS_STREAM_TYPE_X86_TO_TOPO_LSAP) {
|
|
assert(d->onHostReady &&
|
|
"Host-to-device stream should have data initially on host.");
|
|
size_t data_size = sizeof(uint64_t);
|
|
for (int i = 0; i < d->rank; ++i)
|
|
data_size *= d->host_data.sizes[i];
|
|
d->device_data = cuda_malloc_async(
|
|
data_size, (cudaStream_t *)dfg->gpu_stream, dfg->gpu_idx);
|
|
cuda_memcpy_async_to_gpu(
|
|
d->device_data, d->host_data.aligned + d->host_data.offset, data_size,
|
|
(cudaStream_t *)dfg->gpu_stream, dfg->gpu_idx);
|
|
d->location = dfg->gpu_idx;
|
|
}
|
|
if (type == TS_STREAM_TYPE_TOPO_TO_TOPO_LSAP)
|
|
assert(d->location == (int32_t)dfg->gpu_idx &&
|
|
"Data transfers between GPUs not supported yet");
|
|
// TODO: in case of TS_STREAM_TYPE_TOPO_TO_X86_LSAP, we could
|
|
// initiate transfer back to host early here - but need to
|
|
// allocate memory and then copy out again. Tradeoff might be
|
|
// worth testing.
|
|
|
|
// If a dependence was already present, schedule deallocation.
|
|
if (dep != nullptr)
|
|
dep->free_data(dfg);
|
|
dep = d;
|
|
}
|
|
void schedule_work() {
|
|
// If there's no producer process for this stream, it is fed by
|
|
// the control program - nothing to do
|
|
if (producer == nullptr) {
|
|
assert(dep != nullptr && "Data missing on control program stream.");
|
|
return;
|
|
}
|
|
// Recursively go up the DFG to check if new data is available
|
|
for (auto s : producer->input_streams)
|
|
s->schedule_work();
|
|
// Check if any of the inputs have changed - and if so recompute
|
|
// this value. Do not recompute if no changes.
|
|
for (auto s : producer->input_streams)
|
|
if (dep == nullptr || s->dep->used == false) {
|
|
schedule_kernel(producer);
|
|
break;
|
|
}
|
|
}
|
|
Dependence *get_on_host(MemRef2 &out, bool has_scheduled = false) {
|
|
if (!has_scheduled)
|
|
schedule_work();
|
|
assert(dep != nullptr && "GET on empty stream not allowed.");
|
|
dep->used = true;
|
|
// If this was already copied to host, copy out
|
|
if (dep->onHostReady) {
|
|
memref_copy_one_rank(dep->host_data.allocated, dep->host_data.aligned,
|
|
dep->host_data.offset, dep->host_data.sizes[0],
|
|
dep->host_data.strides[0], out.allocated,
|
|
out.aligned, out.offset, out.sizes[0],
|
|
out.strides[0]);
|
|
} else {
|
|
cuda_memcpy_async_to_cpu(out.aligned + out.offset, dep->device_data,
|
|
out.sizes[0] * sizeof(uint64_t),
|
|
(cudaStream_t *)dfg->gpu_stream, dep->location);
|
|
cudaStreamSynchronize(*(cudaStream_t *)dfg->gpu_stream);
|
|
// After this synchronization point, all of the host-side
|
|
// allocated memory can be freed as we know all asynchronous
|
|
// operations have finished.
|
|
dfg->free_stream_order_dependent_data();
|
|
dep->onHostReady = true;
|
|
}
|
|
return dep;
|
|
}
|
|
Dependence *get(int32_t location) {
|
|
schedule_work();
|
|
assert(dep != nullptr && "Dependence could not be computed.");
|
|
dep->used = true;
|
|
if (location == host_location) {
|
|
if (dep->onHostReady)
|
|
return dep;
|
|
size_t data_size = sizeof(uint64_t);
|
|
for (int i = 0; i < dep->rank; ++i)
|
|
data_size *= dep->host_data.sizes[i];
|
|
dep->host_data.allocated = dep->host_data.aligned =
|
|
(uint64_t *)malloc(data_size);
|
|
dep->hostAllocated = true;
|
|
get_on_host(dep->host_data, true);
|
|
return dep;
|
|
}
|
|
assert(dep->location == location &&
|
|
"Multi-GPU within the same SDFG not supported");
|
|
return dep;
|
|
}
|
|
};
|
|
|
|
void GPU_DFG::free_streams() {
|
|
streams.sort();
|
|
streams.unique();
|
|
for (auto s : streams)
|
|
delete s;
|
|
}
|
|
|
|
static inline mlir::concretelang::gpu_dfg::Process *
|
|
make_process_1_1(void *dfg, void *sin1, void *sout, void (*fun)(Process *)) {
|
|
mlir::concretelang::gpu_dfg::Process *p =
|
|
new mlir::concretelang::gpu_dfg::Process;
|
|
mlir::concretelang::gpu_dfg::Stream *s1 =
|
|
(mlir::concretelang::gpu_dfg::Stream *)sin1;
|
|
mlir::concretelang::gpu_dfg::Stream *so =
|
|
(mlir::concretelang::gpu_dfg::Stream *)sout;
|
|
p->input_streams.push_back(s1);
|
|
p->dfg = (GPU_DFG *)dfg;
|
|
p->fun = fun;
|
|
p->output_streams.push_back(so);
|
|
s1->consumers.push_back(p);
|
|
so->producer = p;
|
|
so->dfg = s1->dfg = (GPU_DFG *)dfg;
|
|
p->dfg->register_stream(s1);
|
|
p->dfg->register_stream(so);
|
|
return p;
|
|
}
|
|
|
|
static inline mlir::concretelang::gpu_dfg::Process *
|
|
make_process_2_1(void *dfg, void *sin1, void *sin2, void *sout,
|
|
void (*fun)(Process *)) {
|
|
mlir::concretelang::gpu_dfg::Process *p =
|
|
new mlir::concretelang::gpu_dfg::Process;
|
|
mlir::concretelang::gpu_dfg::Stream *s1 =
|
|
(mlir::concretelang::gpu_dfg::Stream *)sin1;
|
|
mlir::concretelang::gpu_dfg::Stream *s2 =
|
|
(mlir::concretelang::gpu_dfg::Stream *)sin2;
|
|
mlir::concretelang::gpu_dfg::Stream *so =
|
|
(mlir::concretelang::gpu_dfg::Stream *)sout;
|
|
p->input_streams.push_back(s1);
|
|
p->input_streams.push_back(s2);
|
|
p->dfg = (GPU_DFG *)dfg;
|
|
p->fun = fun;
|
|
p->output_streams.push_back(so);
|
|
s1->consumers.push_back(p);
|
|
s2->consumers.push_back(p);
|
|
so->producer = p;
|
|
so->dfg = s1->dfg = s2->dfg = (GPU_DFG *)dfg;
|
|
p->dfg->register_stream(s1);
|
|
p->dfg->register_stream(s2);
|
|
p->dfg->register_stream(so);
|
|
return p;
|
|
}
|
|
|
|
// Stream emulator processes
|
|
void memref_keyswitch_lwe_u64_process(Process *p) {
|
|
Dependence *idep = p->input_streams[0]->get(p->dfg->gpu_idx);
|
|
void *ct0_gpu = idep->device_data;
|
|
void *ksk_gpu = p->ctx.val->get_ksk_gpu(
|
|
p->level.val, p->input_lwe_dim.val, p->output_lwe_dim.val,
|
|
p->dfg->gpu_idx, (cudaStream_t *)p->dfg->gpu_stream);
|
|
MemRef2 out = {0, 0, 0, {p->output_size.val}, {1}};
|
|
void *out_gpu =
|
|
cuda_malloc_async(out.sizes[0] * sizeof(uint64_t),
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx);
|
|
// Schedule the keyswitch kernel on the GPU
|
|
cuda_keyswitch_lwe_ciphertext_vector_64(
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx, out_gpu, ct0_gpu,
|
|
ksk_gpu, p->input_lwe_dim.val, p->output_lwe_dim.val, p->base_log.val,
|
|
p->level.val, 1 /*batch size*/);
|
|
Dependence *dep =
|
|
new Dependence((int32_t)p->dfg->gpu_idx, 1, out, out_gpu, false);
|
|
p->output_streams[0]->put(dep);
|
|
}
|
|
|
|
// Construct the glwe accumulator (on CPU) then put it on a stream as
|
|
// input to the bootstrap
|
|
void memref_bootstrap_glwe_accumulator_process(Process *p) {
|
|
uint64_t glwe_ct_len = p->poly_size.val * (p->glwe_dim.val + 1);
|
|
uint64_t glwe_ct_size = glwe_ct_len * sizeof(uint64_t);
|
|
uint64_t *glwe_ct = (uint64_t *)malloc(glwe_ct_size);
|
|
Dependence *idep = p->input_streams[0]->get(host_location);
|
|
MemRef2 &mtlu = idep->host_data;
|
|
auto tlu = mtlu.aligned + mtlu.offset;
|
|
|
|
// Glwe trivial encryption
|
|
for (size_t i = 0; i < p->poly_size.val * p->glwe_dim.val; i++) {
|
|
glwe_ct[i] = 0;
|
|
}
|
|
for (size_t i = 0; i < p->poly_size.val; i++) {
|
|
glwe_ct[p->poly_size.val * p->glwe_dim.val + i] = tlu[i];
|
|
}
|
|
MemRef2 m = {glwe_ct, glwe_ct, 0, {glwe_ct_len}, {1}};
|
|
Dependence *dep = new Dependence(host_location, 1, m, nullptr, true, true);
|
|
p->output_streams[0]->put(dep);
|
|
}
|
|
|
|
void memref_bootstrap_lwe_u64_process(Process *p) {
|
|
assert(p->output_size.val == p->glwe_dim.val * p->poly_size.val + 1);
|
|
uint32_t num_samples = 1; // TODO batching
|
|
void *fbsk_gpu = p->ctx.val->get_bsk_gpu(
|
|
p->input_lwe_dim.val, p->poly_size.val, p->level.val, p->glwe_dim.val,
|
|
p->dfg->gpu_idx, (cudaStream_t *)p->dfg->gpu_stream);
|
|
Dependence *idep0 = p->input_streams[0]->get(p->dfg->gpu_idx);
|
|
Dependence *idep1 = p->input_streams[1]->get(p->dfg->gpu_idx);
|
|
void *ct0_gpu = idep0->device_data;
|
|
void *glwe_ct_gpu = idep1->device_data;
|
|
|
|
MemRef2 out = {0, 0, 0, {p->output_size.val}, {1}};
|
|
uint64_t out_batch_size = out.sizes[0];
|
|
void *out_gpu =
|
|
cuda_malloc_async(out_batch_size * sizeof(uint64_t),
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx);
|
|
// Move test vector indexes to the GPU, the test vector indexes is set of 0
|
|
uint32_t num_test_vectors = 1, lwe_idx = 0,
|
|
test_vector_idxes_size = num_samples * sizeof(uint64_t);
|
|
void *test_vector_idxes = malloc(test_vector_idxes_size);
|
|
memset(test_vector_idxes, 0, test_vector_idxes_size);
|
|
void *test_vector_idxes_gpu =
|
|
cuda_malloc_async(test_vector_idxes_size,
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx);
|
|
cuda_memcpy_async_to_gpu(test_vector_idxes_gpu, test_vector_idxes,
|
|
test_vector_idxes_size,
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx);
|
|
// Schedule the bootstrap kernel on the GPU
|
|
int8_t *pbs_buffer =
|
|
p->dfg->get_pbs_buffer(p->glwe_dim.val, p->poly_size.val, num_samples);
|
|
cuda_bootstrap_amortized_lwe_ciphertext_vector_64(
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx, out_gpu, glwe_ct_gpu,
|
|
test_vector_idxes_gpu, ct0_gpu, fbsk_gpu, (int8_t *)pbs_buffer,
|
|
p->input_lwe_dim.val, p->glwe_dim.val, p->poly_size.val, p->base_log.val,
|
|
p->level.val, num_samples, num_test_vectors, lwe_idx,
|
|
cuda_get_max_shared_memory(p->dfg->gpu_idx));
|
|
p->dfg->drop_pbs_buffer();
|
|
cuda_drop_async(test_vector_idxes_gpu, (cudaStream_t *)p->dfg->gpu_stream,
|
|
p->dfg->gpu_idx);
|
|
Dependence *dep =
|
|
new Dependence((int32_t)p->dfg->gpu_idx, 1, out, out_gpu, false);
|
|
// As streams are not synchronized, we can only free this vector
|
|
// after a later synchronization point where we are guaranteed that
|
|
// this vector is no longer needed.
|
|
p->dfg->register_stream_order_dependent_allocation(test_vector_idxes);
|
|
p->output_streams[0]->put(dep);
|
|
}
|
|
|
|
void memref_add_lwe_ciphertexts_u64_process(Process *p) {
|
|
Dependence *idep0 = p->input_streams[0]->get(p->dfg->gpu_idx);
|
|
Dependence *idep1 = p->input_streams[1]->get(p->dfg->gpu_idx);
|
|
MemRef2 ct0 = idep0->host_data;
|
|
void *out_gpu =
|
|
cuda_malloc_async(ct0.sizes[0] * sizeof(uint64_t),
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx);
|
|
MemRef2 out = {0, 0, 0, {ct0.sizes[0]}, {1}};
|
|
cuda_add_lwe_ciphertext_vector_64(
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx, out_gpu,
|
|
idep0->device_data, idep1->device_data,
|
|
/*p->input_lwe_dim.val*/ ct0.sizes[0] - 1, 1 /* num_samples */);
|
|
Dependence *dep = new Dependence(p->dfg->gpu_idx, 1, out, out_gpu, false);
|
|
p->output_streams[0]->put(dep);
|
|
}
|
|
|
|
void memref_add_plaintext_lwe_ciphertext_u64_process(Process *p) {
|
|
Dependence *idep0 = p->input_streams[0]->get(p->dfg->gpu_idx);
|
|
Dependence *idep1 = p->input_streams[1]->get(p->dfg->gpu_idx);
|
|
MemRef2 ct0 = idep0->host_data;
|
|
void *out_gpu =
|
|
cuda_malloc_async(ct0.sizes[0] * sizeof(uint64_t),
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx);
|
|
MemRef2 out = {0, 0, 0, {ct0.sizes[0]}, {1}};
|
|
cuda_add_lwe_ciphertext_vector_plaintext_vector_64(
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx, out_gpu,
|
|
idep0->device_data, idep1->device_data,
|
|
/*p->input_lwe_dim.val*/ ct0.sizes[0] - 1, 1 /* num_samples */);
|
|
|
|
Dependence *dep = new Dependence(p->dfg->gpu_idx, 1, out, out_gpu, false);
|
|
p->output_streams[0]->put(dep);
|
|
}
|
|
|
|
void memref_mul_cleartext_lwe_ciphertext_u64_process(Process *p) {
|
|
Dependence *idep0 = p->input_streams[0]->get(p->dfg->gpu_idx);
|
|
Dependence *idep1 = p->input_streams[1]->get(p->dfg->gpu_idx);
|
|
MemRef2 ct0 = idep0->host_data;
|
|
void *out_gpu =
|
|
cuda_malloc_async(ct0.sizes[0] * sizeof(uint64_t),
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx);
|
|
MemRef2 out = {0, 0, 0, {ct0.sizes[0]}, {1}};
|
|
cuda_mult_lwe_ciphertext_vector_cleartext_vector_64(
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx, out_gpu,
|
|
idep0->device_data, idep1->device_data,
|
|
/*p->input_lwe_dim.val*/ ct0.sizes[0] - 1, 1 /* num_samples */);
|
|
|
|
Dependence *dep = new Dependence(p->dfg->gpu_idx, 1, out, out_gpu, false);
|
|
p->output_streams[0]->put(dep);
|
|
}
|
|
|
|
void memref_negate_lwe_ciphertext_u64_process(Process *p) {
|
|
Dependence *idep = p->input_streams[0]->get(p->dfg->gpu_idx);
|
|
MemRef2 ct0 = idep->host_data;
|
|
void *out_gpu =
|
|
cuda_malloc_async(ct0.sizes[0] * sizeof(uint64_t),
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx);
|
|
MemRef2 out = {0, 0, 0, {ct0.sizes[0]}, {1}};
|
|
cuda_negate_lwe_ciphertext_vector_64(
|
|
(cudaStream_t *)p->dfg->gpu_stream, p->dfg->gpu_idx, out_gpu,
|
|
idep->device_data,
|
|
/*p->input_lwe_dim.val*/ ct0.sizes[0] - 1, 1 /* num_samples */);
|
|
Dependence *dep = new Dependence(p->dfg->gpu_idx, 1, out, out_gpu, false);
|
|
p->output_streams[0]->put(dep);
|
|
}
|
|
|
|
} // namespace
|
|
} // namespace gpu_dfg
|
|
} // namespace concretelang
|
|
} // namespace mlir
|
|
|
|
using namespace mlir::concretelang::gpu_dfg;
|
|
|
|
// Code generation interface
|
|
void stream_emulator_make_memref_add_lwe_ciphertexts_u64_process(void *dfg,
|
|
void *sin1,
|
|
void *sin2,
|
|
void *sout) {
|
|
Process *p = make_process_2_1(dfg, sin1, sin2, sout,
|
|
memref_add_lwe_ciphertexts_u64_process);
|
|
static int count = 0;
|
|
sprintf(p->name, "add_lwe_ciphertexts_%d", count++);
|
|
}
|
|
|
|
void stream_emulator_make_memref_add_plaintext_lwe_ciphertext_u64_process(
|
|
void *dfg, void *sin1, void *sin2, void *sout) {
|
|
Process *p = make_process_2_1(
|
|
dfg, sin1, sin2, sout, memref_add_plaintext_lwe_ciphertext_u64_process);
|
|
static int count = 0;
|
|
sprintf(p->name, "add_plaintext_lwe_ciphertexts_%d", count++);
|
|
}
|
|
|
|
void stream_emulator_make_memref_mul_cleartext_lwe_ciphertext_u64_process(
|
|
void *dfg, void *sin1, void *sin2, void *sout) {
|
|
Process *p = make_process_2_1(
|
|
dfg, sin1, sin2, sout, memref_mul_cleartext_lwe_ciphertext_u64_process);
|
|
static int count = 0;
|
|
sprintf(p->name, "mul_cleartext_lwe_ciphertexts_%d", count++);
|
|
}
|
|
|
|
void stream_emulator_make_memref_negate_lwe_ciphertext_u64_process(void *dfg,
|
|
void *sin1,
|
|
void *sout) {
|
|
Process *p = make_process_1_1(dfg, sin1, sout,
|
|
memref_negate_lwe_ciphertext_u64_process);
|
|
static int count = 0;
|
|
sprintf(p->name, "negate_lwe_ciphertext_%d", count++);
|
|
}
|
|
|
|
void stream_emulator_make_memref_keyswitch_lwe_u64_process(
|
|
void *dfg, void *sin1, void *sout, uint32_t level, uint32_t base_log,
|
|
uint32_t input_lwe_dim, uint32_t output_lwe_dim, uint32_t output_size,
|
|
void *context) {
|
|
Process *p =
|
|
make_process_1_1(dfg, sin1, sout, memref_keyswitch_lwe_u64_process);
|
|
p->level.val = level;
|
|
p->base_log.val = base_log;
|
|
p->input_lwe_dim.val = input_lwe_dim;
|
|
p->output_lwe_dim.val = output_lwe_dim;
|
|
p->output_size.val = output_size;
|
|
p->ctx.val = (RuntimeContext *)context;
|
|
static int count = 0;
|
|
sprintf(p->name, "keyswitch_%d", count++);
|
|
}
|
|
|
|
void stream_emulator_make_memref_bootstrap_lwe_u64_process(
|
|
void *dfg, void *sin1, void *sin2, void *sout, uint32_t input_lwe_dim,
|
|
uint32_t poly_size, uint32_t level, uint32_t base_log, uint32_t glwe_dim,
|
|
uint32_t precision, uint32_t output_size, void *context) {
|
|
|
|
// We need to generate two processes: one for building the glwe
|
|
// accumulator and one for the bootstrap (plus a stream to connect).
|
|
void *accu_s =
|
|
stream_emulator_make_memref_stream("", TS_STREAM_TYPE_X86_TO_TOPO_LSAP);
|
|
Process *accu_p = make_process_1_1(dfg, sin2, accu_s,
|
|
memref_bootstrap_glwe_accumulator_process);
|
|
accu_p->input_lwe_dim.val = input_lwe_dim;
|
|
accu_p->poly_size.val = poly_size;
|
|
accu_p->level.val = level;
|
|
accu_p->base_log.val = base_log;
|
|
accu_p->glwe_dim.val = glwe_dim;
|
|
accu_p->precision.val = precision;
|
|
accu_p->output_size.val = output_size;
|
|
accu_p->ctx.val = (RuntimeContext *)context;
|
|
static int count_ = 0;
|
|
sprintf(accu_p->name, "glwe_accumulator_%d", count_++);
|
|
|
|
Process *p = make_process_2_1(dfg, sin1, accu_s, sout,
|
|
memref_bootstrap_lwe_u64_process);
|
|
p->input_lwe_dim.val = input_lwe_dim;
|
|
p->poly_size.val = poly_size;
|
|
p->level.val = level;
|
|
p->base_log.val = base_log;
|
|
p->glwe_dim.val = glwe_dim;
|
|
p->precision.val = precision;
|
|
p->output_size.val = output_size;
|
|
p->ctx.val = (RuntimeContext *)context;
|
|
static int count = 0;
|
|
sprintf(p->name, "bootstrap_%d", count++);
|
|
}
|
|
|
|
void *stream_emulator_make_uint64_stream(const char *name, stream_type stype) {
|
|
return (void *)new Stream(stype);
|
|
}
|
|
void stream_emulator_put_uint64(void *stream, uint64_t e) {
|
|
Stream *s = (Stream *)stream;
|
|
uint64_t *data = (uint64_t *)malloc(sizeof(uint64_t));
|
|
*data = e;
|
|
MemRef2 m = {data, data, 0, {1}, {1}};
|
|
Dependence *dep = new Dependence(host_location, 1, m, nullptr, true, true);
|
|
s->put(dep);
|
|
}
|
|
uint64_t stream_emulator_get_uint64(void *stream) {
|
|
uint64_t res;
|
|
auto s = (Stream *)stream;
|
|
MemRef2 m = {&res, &res, 0, {1}, {1}};
|
|
s->get_on_host(m);
|
|
return res;
|
|
}
|
|
|
|
void *stream_emulator_make_memref_stream(const char *name, stream_type stype) {
|
|
return (void *)new Stream(stype);
|
|
}
|
|
void stream_emulator_put_memref(void *stream, uint64_t *allocated,
|
|
uint64_t *aligned, uint64_t offset,
|
|
uint64_t size, uint64_t stride) {
|
|
Stream *s = (Stream *)stream;
|
|
MemRef2 m = {allocated, aligned, offset, {size}, {stride}};
|
|
Dependence *dep = new Dependence(host_location, 1, m, nullptr, true);
|
|
s->put(dep);
|
|
}
|
|
void stream_emulator_get_memref(void *stream, uint64_t *out_allocated,
|
|
uint64_t *out_aligned, uint64_t out_offset,
|
|
uint64_t out_size, uint64_t out_stride) {
|
|
MemRef2 mref = {
|
|
out_allocated, out_aligned, out_offset, {out_size}, {out_stride}};
|
|
auto s = (Stream *)stream;
|
|
s->get_on_host(mref);
|
|
}
|
|
|
|
void *stream_emulator_init() {
|
|
int num;
|
|
if (num_devices == 0) {
|
|
assert(cudaGetDeviceCount(&num) == cudaSuccess);
|
|
num_devices = num;
|
|
}
|
|
int device = next_device.fetch_add(1) % num_devices;
|
|
return new GPU_DFG(device);
|
|
}
|
|
void stream_emulator_run(void *dfg) {}
|
|
void stream_emulator_delete(void *dfg) { delete (GPU_DFG *)dfg; }
|
|
#endif
|