Skip to content
This repository has been archived by the owner on Aug 16, 2024. It is now read-only.

Commit

Permalink
Merge branch 'main' into rr-public-inputs-check
Browse files Browse the repository at this point in the history
  • Loading branch information
@robik75
robik75 authored Nov 10, 2023
2 parents 3cfe6af + c4e713e commit 663d5f4
Show file tree
Hide file tree
Showing 30 changed files with 3,686 additions and 941 deletions.
4 changes: 2 additions & 2 deletions Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@ edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
# boojum = { path = "../era-boojum" }
# boojum-cuda = { path = "../era-boojum-cuda/boojum-cuda" }
# boojum = { path = "../era-boojum", package = "boojum" }
# boojum-cuda = { path = "../era-boojum-cuda/boojum-cuda", package = "boojum-cuda" }
# cudart = { path = "../era-boojum-cuda/cudart", package = "cudart" }
# circuit_definitions = { path = "../era-zkevm_test_harness/circuit_definitions", package = "circuit_definitions", optional = true }

Expand Down
2 changes: 1 addition & 1 deletion src/constraint_evaluation.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ pub fn multi_polys_as_single_slice<'a, P: PolyForm>(polys: &[Poly<'a, P>]) -> &'
if polys.is_empty() {
return &[];
};
assert_adjacent(polys);
assert_adjacent_base(polys);
unsafe {
let len = polys.len() * polys[0].domain_size();
std::slice::from_raw_parts(polys[0].storage.as_ref().as_ptr(), len)
Expand Down
165 changes: 142 additions & 23 deletions src/context.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,27 +2,109 @@ use super::*;

pub struct ProverContext;

pub const ZKSYNC_DEFAULT_TRACE_LOG_LENGTH: usize = 19;

impl ProverContext {
pub fn create() -> CudaResult<Self> {
unsafe {
assert!(_CUDA_CONTEXT.is_none());
assert!(_ALLOCATOR.is_none());
assert!(_SMALL_ALLOCATOR.is_none());
assert!(_DEVICE_ALLOCATOR.is_none());
assert!(_SMALL_DEVICE_ALLOCATOR.is_none());
assert!(_HOST_ALLOCATOR.is_none());
assert!(_SMALL_HOST_ALLOCATOR.is_none());
assert!(_EXEC_STREAM.is_none());
assert!(_H2D_STREAM.is_none());
assert!(_D2H_STREAM.is_none());
}
// size counts in field elements
let block_size = 1 << ZKSYNC_DEFAULT_TRACE_LOG_LENGTH;
let cuda_ctx = CudaContext::create(12, 12)?;

// grab small slice then consume everything
let small_device_alloc = SmallStaticDeviceAllocator::init()?;
let device_alloc = StaticDeviceAllocator::init_all(block_size)?;

let small_host_alloc = SmallStaticHostAllocator::init()?;
let host_alloc = StaticHostAllocator::init(block_size, 1 << 8)?;

unsafe {
_CUDA_CONTEXT = Some(cuda_ctx);
_DEVICE_ALLOCATOR = Some(device_alloc);
_SMALL_DEVICE_ALLOCATOR = Some(small_device_alloc);
_HOST_ALLOCATOR = Some(host_alloc);
_SMALL_HOST_ALLOCATOR = Some(small_host_alloc);
_EXEC_STREAM = Some(Stream::create()?);
_H2D_STREAM = Some(Stream::create()?);
_D2H_STREAM = Some(Stream::create()?);
}

Ok(Self {})
}

#[allow(dead_code)]
pub(crate) fn create_14gb_dev(block_size: usize) -> CudaResult<Self> {
unsafe {
assert!(_CUDA_CONTEXT.is_none());
assert!(_DEVICE_ALLOCATOR.is_none());
assert!(_SMALL_DEVICE_ALLOCATOR.is_none());
assert!(_HOST_ALLOCATOR.is_none());
assert!(_SMALL_HOST_ALLOCATOR.is_none());
assert!(_EXEC_STREAM.is_none());
assert!(_H2D_STREAM.is_none());
assert!(_D2H_STREAM.is_none());
}
// size counts in field elements
let cuda_ctx = CudaContext::create(12, 12)?;

// grab small slice then consume everything
let small_device_alloc = SmallStaticDeviceAllocator::init()?;
let device_alloc = StaticDeviceAllocator::init_14gb(block_size)?;
println!("allocated 14gb on device");
let small_host_alloc = SmallStaticHostAllocator::init()?;
let host_alloc = StaticHostAllocator::init(block_size, 1 << 8)?;

unsafe {
_CUDA_CONTEXT = Some(cuda_ctx);
_DEVICE_ALLOCATOR = Some(device_alloc);
_SMALL_DEVICE_ALLOCATOR = Some(small_device_alloc);
_HOST_ALLOCATOR = Some(host_alloc);
_SMALL_HOST_ALLOCATOR = Some(small_host_alloc);
_EXEC_STREAM = Some(Stream::create()?);
_H2D_STREAM = Some(Stream::create()?);
_D2H_STREAM = Some(Stream::create()?);
}

Ok(Self {})
}

pub fn create_14gb() -> CudaResult<Self> {
unsafe {
assert!(_CUDA_CONTEXT.is_none());
assert!(_DEVICE_ALLOCATOR.is_none());
assert!(_SMALL_DEVICE_ALLOCATOR.is_none());
assert!(_HOST_ALLOCATOR.is_none());
assert!(_SMALL_HOST_ALLOCATOR.is_none());
assert!(_EXEC_STREAM.is_none());
assert!(_H2D_STREAM.is_none());
assert!(_D2H_STREAM.is_none());
}
// size counts in field elements
let block_size = 1 << 20; // TODO:
// grab small slice then consume everything
let block_size = 1 << ZKSYNC_DEFAULT_TRACE_LOG_LENGTH;
let cuda_ctx = CudaContext::create(12, 12)?;
let small_alloc = SmallVirtualMemoryManager::init()?;
let alloc = VirtualMemoryManager::init_all(block_size)?;

// grab small slice then consume everything
let small_device_alloc = SmallStaticDeviceAllocator::init()?;
let device_alloc = StaticDeviceAllocator::init_14gb(block_size)?;
println!("allocated 14gb on device");
let small_host_alloc = SmallStaticHostAllocator::init()?;
let host_alloc = StaticHostAllocator::init(block_size, 1 << 8)?;

unsafe {
_CUDA_CONTEXT = Some(cuda_ctx);
_ALLOCATOR = Some(alloc);
_SMALL_ALLOCATOR = Some(small_alloc);
_DEVICE_ALLOCATOR = Some(device_alloc);
_SMALL_DEVICE_ALLOCATOR = Some(small_device_alloc);
_HOST_ALLOCATOR = Some(host_alloc);
_SMALL_HOST_ALLOCATOR = Some(small_host_alloc);
_EXEC_STREAM = Some(Stream::create()?);
_H2D_STREAM = Some(Stream::create()?);
_D2H_STREAM = Some(Stream::create()?);
Expand All @@ -31,19 +113,25 @@ impl ProverContext {
Ok(Self {})
}

#[cfg(test)]
#[allow(dead_code)]
pub(crate) fn dev(domain_size: usize) -> CudaResult<Self> {
assert!(domain_size.is_power_of_two());
// size counts in field elements
let block_size = domain_size;
let cuda_ctx = CudaContext::create(12, 12)?;
let small_alloc = SmallVirtualMemoryManager::init()?;
let alloc = VirtualMemoryManager::init_all(block_size)?;

let small_device_alloc = SmallStaticDeviceAllocator::init()?;
let device_alloc = StaticDeviceAllocator::init_all(block_size)?;

let small_host_alloc = SmallStaticHostAllocator::init()?;
let host_alloc = StaticHostAllocator::init(block_size, 1 << 8)?;

unsafe {
_CUDA_CONTEXT = Some(cuda_ctx);
_ALLOCATOR = Some(alloc);
_SMALL_ALLOCATOR = Some(small_alloc);
_DEVICE_ALLOCATOR = Some(device_alloc);
_SMALL_DEVICE_ALLOCATOR = Some(small_device_alloc);
_HOST_ALLOCATOR = Some(host_alloc);
_SMALL_HOST_ALLOCATOR = Some(small_host_alloc);
_EXEC_STREAM = Some(Stream::create()?);
_H2D_STREAM = Some(Stream::create()?);
_D2H_STREAM = Some(Stream::create()?);
Expand All @@ -59,8 +147,22 @@ impl Drop for ProverContext {
let cuda_ctx = _CUDA_CONTEXT.take().expect("cuda ctx");
cuda_ctx.destroy().expect("destroy cuda ctx");

_ALLOCATOR.take().unwrap().free().expect("free allocator");
_SMALL_ALLOCATOR
_DEVICE_ALLOCATOR
.take()
.unwrap()
.free()
.expect("free allocator");
_SMALL_DEVICE_ALLOCATOR
.take()
.unwrap()
.free()
.expect("free small allocator");
_HOST_ALLOCATOR
.take()
.unwrap()
.free()
.expect("free allocator");
_SMALL_HOST_ALLOCATOR
.take()
.unwrap()
.free()
Expand Down Expand Up @@ -130,21 +232,38 @@ impl Stream {
unsafe impl Send for Stream {}
unsafe impl Sync for Stream {}

pub(crate) static mut _ALLOCATOR: Option<VirtualMemoryManager> = None;
pub(crate) static mut _SMALL_ALLOCATOR: Option<SmallVirtualMemoryManager> = None;
pub(crate) static mut _DEVICE_ALLOCATOR: Option<StaticDeviceAllocator> = None;
pub(crate) static mut _SMALL_DEVICE_ALLOCATOR: Option<SmallStaticDeviceAllocator> = None;
pub(crate) static mut _HOST_ALLOCATOR: Option<StaticHostAllocator> = None;
pub(crate) static mut _SMALL_HOST_ALLOCATOR: Option<SmallStaticHostAllocator> = None;

pub(crate) fn _alloc() -> &'static StaticDeviceAllocator {
unsafe {
&_DEVICE_ALLOCATOR
.as_ref()
.expect("device allocator should be initialized")
}
}

pub(crate) fn _alloc() -> &'static VirtualMemoryManager {
pub(crate) fn _small_alloc() -> &'static SmallStaticDeviceAllocator {
unsafe {
&_SMALL_DEVICE_ALLOCATOR
.as_ref()
.expect("small device allocator should be initialized")
}
}
pub(crate) fn _host_alloc() -> &'static StaticHostAllocator {
unsafe {
&_ALLOCATOR
&_HOST_ALLOCATOR
.as_ref()
.expect("allocator should be initialized")
.expect("host allocator should be initialized")
}
}

pub(crate) fn _small_alloc() -> &'static SmallVirtualMemoryManager {
pub(crate) fn _small_host_alloc() -> &'static SmallStaticHostAllocator {
unsafe {
&_SMALL_ALLOCATOR
&_SMALL_HOST_ALLOCATOR
.as_ref()
.expect("small allocator should be initialized")
.expect("small host allocator should be initialized")
}
}
Loading

0 comments on commit 663d5f4

Please sign in to comment.