Skip to content

Commit 9c10358

Browse files
876pol876pol
committed
style: small cleanup
1 parent 90c3ff2 commit 9c10358

File tree

1 file changed

+92
-104
lines changed

1 file changed

+92
-104
lines changed

‎crates/vm/src/system/cuda/memory.rs‎

Lines changed: 92 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -230,115 +230,105 @@ impl MemoryInventoryGPU {
230230
})
231231
.collect();
232232
let in_num_records = in_records.len();
233-
if in_num_records == 0 {
234-
self.boundary
235-
.finalize_records_persistent::<DIGEST_WIDTH>(Vec::new());
236-
mem.tracing_info("merkle update");
237-
persistent.merkle_tree.finalize();
238-
None
239-
} else {
240-
let out_words = in_num_records
241-
* (std::mem::size_of::<MemoryInventoryRecord<8, 2>>()
242-
/ std::mem::size_of::<u32>());
243-
let d_in_records = in_records.to_device().unwrap().as_buffer::<u32>();
244-
let d_tmp_records = DeviceBuffer::<u32>::with_capacity(out_words);
245-
let d_out_records = DeviceBuffer::<u32>::with_capacity(out_words);
246-
let d_out_num_records = DeviceBuffer::<usize>::with_capacity(1);
247-
let d_flags = DeviceBuffer::<u32>::with_capacity(in_num_records);
248-
let d_positions = DeviceBuffer::<u32>::with_capacity(in_num_records);
249-
let d_initial_mem = match &self.boundary.fields {
250-
BoundaryFields::Persistent(fields) => {
251-
fields.initial_leaves.to_device().unwrap()
252-
}
253-
BoundaryFields::Volatile(_) => {
254-
panic!("`merge_records` requires persistent memory")
255-
}
256-
};
257-
let addr_space_offsets: Vec<u32> = {
258-
let mut offsets = Vec::new();
259-
let mut acc = 0u32;
260-
for addr_sp in persistent
261-
.merkle_tree
262-
.mem_config()
263-
.addr_spaces
264-
.iter()
265-
.skip(ADDR_SPACE_OFFSET as usize)
266-
{
267-
offsets.push(acc);
268-
acc = acc.saturating_add(addr_sp.layout.size() as u32);
269-
}
270-
offsets.push(acc);
271-
offsets
272-
};
273-
let d_addr_space_offsets = addr_space_offsets.to_device().unwrap();
274-
let mut temp_bytes = 0usize;
275-
unsafe {
276-
inventory::merge_records_get_temp_bytes(
277-
&d_flags,
278-
in_num_records,
279-
&mut temp_bytes,
280-
)
281-
.expect("merge_records_get_temp_bytes failed");
282-
}
283-
let d_temp_storage = if temp_bytes == 0 {
284-
DeviceBuffer::<u8>::new()
285-
} else {
286-
DeviceBuffer::<u8>::with_capacity(temp_bytes)
287-
};
288-
unsafe {
289-
inventory::merge_records(
290-
&d_in_records,
291-
in_num_records,
292-
&d_initial_mem,
293-
&d_addr_space_offsets,
294-
&d_tmp_records,
295-
&d_out_records,
296-
&d_flags,
297-
&d_positions,
298-
&d_temp_storage,
299-
temp_bytes,
300-
&d_out_num_records,
301-
)
302-
.expect("merge_records failed");
233+
let out_words = in_num_records
234+
* (std::mem::size_of::<MemoryInventoryRecord<8, 2>>()
235+
/ std::mem::size_of::<u32>());
236+
let d_in_records = in_records.to_device().unwrap().as_buffer::<u32>();
237+
let d_tmp_records = DeviceBuffer::<u32>::with_capacity(out_words);
238+
let d_out_records = DeviceBuffer::<u32>::with_capacity(out_words);
239+
let d_out_num_records = DeviceBuffer::<usize>::with_capacity(1);
240+
let d_flags = DeviceBuffer::<u32>::with_capacity(in_num_records);
241+
let d_positions = DeviceBuffer::<u32>::with_capacity(in_num_records);
242+
let d_initial_mem = match &self.boundary.fields {
243+
BoundaryFields::Persistent(fields) => {
244+
fields.initial_leaves.to_device().unwrap()
303245
}
304-
let out_num_records = d_out_num_records.to_host().unwrap()[0] as usize;
305-
self.boundary
306-
.finalize_records_persistent_device::<DIGEST_WIDTH>(
307-
d_out_records,
308-
out_num_records,
309-
);
310-
let d_merkle_records = DeviceBuffer::<u32>::with_capacity(out_words);
311-
unsafe {
312-
cuda_memcpy::<true, true>(
313-
d_merkle_records.as_mut_raw_ptr(),
314-
self.boundary.persistent_records().as_raw_ptr(),
315-
out_words * std::mem::size_of::<u32>(),
316-
)
317-
.expect("failed to copy boundary records for merkle update");
246+
BoundaryFields::Volatile(_) => {
247+
panic!("`merge_records` requires persistent memory")
318248
}
319-
persistent.merkle_records = Some(d_merkle_records);
320-
321-
let unpadded_merkle_height =
322-
persistent.merkle_tree.calculate_unpadded_height(&partition);
323-
#[cfg(feature = "metrics")]
249+
};
250+
let addr_space_offsets: Vec<u32> = {
251+
let mut offsets = Vec::new();
252+
let mut acc = 0u32;
253+
for addr_sp in persistent
254+
.merkle_tree
255+
.mem_config()
256+
.addr_spaces
257+
.iter()
258+
.skip(ADDR_SPACE_OFFSET as usize)
324259
{
325-
self.unpadded_merkle_height = unpadded_merkle_height;
260+
offsets.push(acc);
261+
acc = acc.saturating_add(addr_sp.layout.size() as u32);
326262
}
327-
328-
mem.tracing_info("boundary finalize");
329-
mem.tracing_info("merkle update");
330-
persistent.merkle_tree.finalize();
331-
let merkle_tree_ctx = persistent.merkle_tree.update_with_touched_blocks(
332-
unpadded_merkle_height,
333-
persistent
334-
.merkle_records
335-
.as_ref()
336-
.expect("missing merkle records"),
263+
offsets.push(acc);
264+
offsets
265+
};
266+
let d_addr_space_offsets = addr_space_offsets.to_device().unwrap();
267+
let mut temp_bytes = 0usize;
268+
unsafe {
269+
inventory::merge_records_get_temp_bytes(
270+
&d_flags,
271+
in_num_records,
272+
&mut temp_bytes,
273+
)
274+
.expect("merge_records_get_temp_bytes failed");
275+
}
276+
let d_temp_storage = if temp_bytes == 0 {
277+
DeviceBuffer::<u8>::new()
278+
} else {
279+
DeviceBuffer::<u8>::with_capacity(temp_bytes)
280+
};
281+
unsafe {
282+
inventory::merge_records(
283+
&d_in_records,
284+
in_num_records,
285+
&d_initial_mem,
286+
&d_addr_space_offsets,
287+
&d_tmp_records,
288+
&d_out_records,
289+
&d_flags,
290+
&d_positions,
291+
&d_temp_storage,
292+
temp_bytes,
293+
&d_out_num_records,
294+
)
295+
.expect("merge_records failed");
296+
}
297+
let out_num_records = d_out_num_records.to_host().unwrap()[0] as usize;
298+
self.boundary
299+
.finalize_records_persistent_device::<DIGEST_WIDTH>(
300+
d_out_records,
337301
out_num_records,
338-
false,
339302
);
340-
Some(merkle_tree_ctx)
303+
let d_merkle_records = DeviceBuffer::<u32>::with_capacity(out_words);
304+
unsafe {
305+
cuda_memcpy::<true, true>(
306+
d_merkle_records.as_mut_raw_ptr(),
307+
self.boundary.persistent_records().as_raw_ptr(),
308+
out_words * std::mem::size_of::<u32>(),
309+
)
310+
.expect("failed to copy boundary records for merkle update");
341311
}
312+
persistent.merkle_records = Some(d_merkle_records);
313+
314+
let unpadded_merkle_height =
315+
persistent.merkle_tree.calculate_unpadded_height(&partition);
316+
#[cfg(feature = "metrics")]
317+
{
318+
self.unpadded_merkle_height = unpadded_merkle_height;
319+
}
320+
321+
persistent.merkle_tree.finalize();
322+
let merkle_tree_ctx = persistent.merkle_tree.update_with_touched_blocks(
323+
unpadded_merkle_height,
324+
persistent
325+
.merkle_records
326+
.as_ref()
327+
.expect("missing merkle records"),
328+
out_num_records,
329+
false,
330+
);
331+
Some(merkle_tree_ctx)
342332
}
343333
}
344334
TouchedMemory::Volatile(partition) => {
@@ -347,11 +337,9 @@ impl MemoryInventoryGPU {
347337
None
348338
}
349339
};
350-
mem.tracing_info("boundary tracegen");
351340
let mut ret = vec![self.boundary.generate_proving_ctx(())];
352341
if let Some(merkle_proof_ctx) = merkle_proof_ctx {
353342
ret.push(merkle_proof_ctx);
354-
mem.tracing_info("dropping merkle tree");
355343
let persistent = self.persistent.as_mut().unwrap();
356344
persistent.merkle_tree.drop_subtrees();
357345
persistent.initial_memory = Vec::new();

0 commit comments

Comments
 (0)