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

Commit

Permalink
rewrite the logic
Browse files Browse the repository at this point in the history
  • Loading branch information
@robik75
robik75 committed Dec 20, 2023
1 parent d7086ce commit 0b6dc29
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 80 deletions.
2 changes: 1 addition & 1 deletion src/prover.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1098,7 +1098,7 @@ fn gpu_prove_from_trace<
.stats
.lock()
.unwrap()
.print(true);
.print(true, true);
}

gpu_proof.public_inputs = public_inputs;
Expand Down
156 changes: 77 additions & 79 deletions src/static_allocator/device.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,109 +32,107 @@ mod stats {
use derivative::Derivative;
use std::collections::BTreeMap;

type Allocations = BTreeMap<usize, (usize, String)>;

#[derive(Derivative)]
#[derivative(Clone, Debug, Default)]
pub struct AllocationStats {
pub allocations: Allocations,
pub maximum_block_count: usize,
pub maximum_tail_index: usize,
pub maximum_block_count_at_maximum_tail_index: usize,
pub allocations_at_maximum_block_count: Allocations,
pub allocations_at_maximum_block_count_at_maximum_tail_index: Allocations,
}
pub struct Allocations(BTreeMap<usize, (usize, String)>);

impl AllocationStats {
pub fn alloc(&mut self, index: usize, size: usize, backtrace: String) {
self.allocations.insert(index, (size, backtrace));
let current_block_count = self.current_block_count();
let previous_maximum_block_count = self.maximum_block_count;
self.maximum_block_count = self.maximum_block_count.max(current_block_count);
let is_new_maximum_block_count =
self.maximum_block_count != previous_maximum_block_count;
if is_new_maximum_block_count {
self.allocations_at_maximum_block_count = self.allocations.clone();
}
let current_tail_index = index + size;
self.maximum_tail_index = self.maximum_tail_index.max(current_tail_index);
let is_maximum_tail_index = self.maximum_tail_index == current_tail_index;
if is_maximum_tail_index {
let previous_maximum_block_count_at_maximum_tail_index =
self.maximum_block_count_at_maximum_tail_index;
self.maximum_block_count_at_maximum_tail_index = self
.maximum_block_count_at_maximum_tail_index
.max(current_block_count);
let is_new_maximum_block_count_at_maximum_tail_index = self
.maximum_block_count_at_maximum_tail_index
!= previous_maximum_block_count_at_maximum_tail_index;
if is_new_maximum_block_count_at_maximum_tail_index {
self.allocations_at_maximum_block_count_at_maximum_tail_index =
self.allocations.clone();
}
}
impl Allocations {
fn insert(&mut self, index: usize, size: usize, backtrace: String) {
self.0.insert(index, (size, backtrace));
}

pub fn free(&mut self, index: usize) {
self.allocations.remove(&index);
fn remove(&mut self, index: &usize) {
self.0.remove(index);
}

pub fn current_block_count(&self) -> usize {
self.allocations.values().map(|&(size, _)| size).sum()
fn block_count(&self) -> usize {
self.0.values().map(|&(size, _)| size).sum()
}

pub fn current_tail_index(&self) -> usize {
self.allocations
pub fn tail_index(&self) -> usize {
self.0
.last_key_value()
.map_or(0, |(&index, &(size, _))| index + size)
}

fn print_allocations(allocations: &Allocations) {
if allocations.is_empty() {
fn print(&self, detailed: bool, with_backtrace: bool) {
assert!(detailed || !with_backtrace);
if self.0.is_empty() {
println!("no allocations");
return;
}
println!("block_count: {}", self.block_count());
println!("tail_index: {}", self.tail_index());
const SEPARATOR: &str = "================================";
println!("{SEPARATOR}");
let mut last_index = 0;
for (index, (length, trace)) in allocations {
let gap = index - last_index;
last_index = index + length;
if gap != 0 {
println!("gap: {gap}");
if detailed {
let mut last_index = 0;
for (index, (length, trace)) in &self.0 {
let gap = index - last_index;
last_index = index + length;
if gap != 0 {
println!("gap: {gap}");
println!("{SEPARATOR}");
}
println!("index: {index}");
println!("length: {length}");
if with_backtrace {
println!("backtrace: \n{trace}");
}
println!("{SEPARATOR}");
}
println!("index: {index}");
println!("length: {length}");
println!("backtrace: \n{trace}");
println!("{SEPARATOR}");
}
}
}

pub fn print(&self, detailed: bool) {
let AllocationStats {
allocations,
maximum_block_count,
maximum_tail_index,
maximum_block_count_at_maximum_tail_index,
allocations_at_maximum_block_count,
allocations_at_maximum_block_count_at_maximum_tail_index,
} = self;
let current_block_count = self.current_block_count();
let current_tail_index = self.current_tail_index();
println!("current_block_count: {current_block_count}");
println!("current_tail_index: {current_tail_index}");
println!("maximum_block_count: {maximum_block_count}");
println!("maximum_tail_index: {maximum_tail_index}");
println!("maximum_block_count_at_maximum_tail_index: {maximum_block_count_at_maximum_tail_index}");
if detailed {
println!("current_allocations:");
Self::print_allocations(allocations);
println!("allocations_at_maximum_block_count:");
Self::print_allocations(allocations_at_maximum_block_count);
println!("allocations_at_maximum_block_count_at_maximum_tail_index:");
Self::print_allocations(allocations_at_maximum_block_count_at_maximum_tail_index);
#[derive(Derivative)]
#[derivative(Clone, Debug, Default)]
pub struct AllocationStats {
pub allocations: Allocations,
pub allocations_at_maximum_block_count: Allocations,
pub allocations_at_maximum_block_count_at_maximum_tail_index: Allocations,
}

impl AllocationStats {
pub fn alloc(&mut self, index: usize, size: usize, backtrace: String) {
self.allocations.insert(index, size, backtrace);
let current_block_count = self.allocations.block_count();
let current_tail_index = self.allocations.tail_index();
let previous_maximum_block_count =
self.allocations_at_maximum_block_count.block_count();
if current_block_count > previous_maximum_block_count {
self.allocations_at_maximum_block_count = self.allocations.clone();
}
let previous_maximum_tail_index = self
.allocations_at_maximum_block_count_at_maximum_tail_index
.tail_index();
if current_tail_index > previous_maximum_tail_index {
self.allocations_at_maximum_block_count_at_maximum_tail_index =
self.allocations.clone();
} else if current_tail_index == previous_maximum_tail_index {
let previous_maximum_block_count_at_maximum_tail_index = self
.allocations_at_maximum_block_count_at_maximum_tail_index
.block_count();
if current_block_count > previous_maximum_block_count_at_maximum_tail_index {
self.allocations_at_maximum_block_count_at_maximum_tail_index =
self.allocations.clone();
}
}
}

pub fn free(&mut self, index: usize) {
self.allocations.remove(&index);
}

pub fn print(&self, detailed: bool, with_backtrace: bool) {
println!("allocations:");
self.allocations.print(detailed, with_backtrace);
println!("allocations_at_maximum_block_count:");
self.allocations_at_maximum_block_count
.print(detailed, with_backtrace);
println!("allocations_at_maximum_block_count_at_maximum_tail_index:");
self.allocations_at_maximum_block_count_at_maximum_tail_index
.print(detailed, with_backtrace);
}
}
}
Expand Down

0 comments on commit 0b6dc29

Please sign in to comment.