Skip to content

Commit fb9007d

Browse files
jonathanpwangjonathanpwang
authored
chore(perf): use constant for address space in tracing_read when possible (#1899)
I went through the uses of `tracing_read` and switched to using constants for better compilation optimization whenever possible, where the requirements on the address space were determined according to ISA.md I don't think there was really a performance diff: https://github.com/axiom-crypto/openvm-reth-benchmark/actions/runs/16513213487
1 parent 198eb39 commit fb9007d

File tree

6 files changed

+45
-59
lines changed

6 files changed

+45
-59
lines changed

‎crates/toolchain/instructions/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ pub struct VmOpcode(usize);
4141

4242
impl VmOpcode {
4343
/// Returns the corresponding `local_opcode_idx`
44+
#[inline(always)]
4445
pub const fn local_opcode_idx(&self, offset: usize) -> usize {
4546
self.as_usize() - offset
4647
}

‎crates/vm/src/system/native_adapter/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -218,7 +218,7 @@ where
218218
let addr_space = if i == 0 { e } else { f };
219219
reads[i][0] = tracing_read_or_imm_native(
220220
memory,
221-
addr_space.as_canonical_u32(),
221+
addr_space,
222222
*ptr_or_imm,
223223
&mut read_aux.prev_timestamp,
224224
);

‎crates/vm/src/system/native_adapter/util.rs‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@ where
121121
}
122122

123123
/// Reads register value at `ptr` from memory and records the previous timestamp.
124+
/// Reads are only done from address space [NATIVE_AS].
124125
#[inline(always)]
125126
pub fn tracing_read_native<F, const BLOCK_SIZE: usize>(
126127
memory: &mut TracingMemory,
@@ -136,6 +137,7 @@ where
136137
}
137138

138139
/// Writes `ptr, vals` into memory and records the previous timestamp and data.
140+
/// Writes are only done to address space [NATIVE_AS].
139141
#[inline(always)]
140142
pub fn tracing_write_native<F, const BLOCK_SIZE: usize>(
141143
memory: &mut TracingMemory,
@@ -171,20 +173,20 @@ pub fn tracing_write_native_inplace<F, const BLOCK_SIZE: usize>(
171173
#[inline(always)]
172174
pub fn tracing_read_or_imm_native<F>(
173175
memory: &mut TracingMemory,
174-
addr_space: u32,
176+
addr_space: F,
175177
ptr_or_imm: F,
176178
prev_timestamp: &mut u32,
177179
) -> F
178180
where
179181
F: PrimeField32,
180182
{
181183
debug_assert!(
182-
addr_space == RV32_IMM_AS || addr_space == NATIVE_AS,
184+
addr_space == F::ZERO || addr_space == F::from_canonical_u32(NATIVE_AS),
183185
"addr_space={} is not valid",
184186
addr_space
185187
);
186188

187-
if addr_space == RV32_IMM_AS {
189+
if addr_space == F::ZERO {
188190
*prev_timestamp = u32::MAX;
189191
memory.increment_timestamp();
190192
ptr_or_imm

‎extensions/native/circuit/src/adapters/alu_native_adapter.rs‎

Lines changed: 2 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -170,19 +170,9 @@ impl<F: PrimeField32> AdapterTraceStep<F> for AluNativeAdapterStep {
170170
let &Instruction { b, c, e, f, .. } = instruction;
171171

172172
record.b = b;
173-
let rs1 = tracing_read_or_imm_native(
174-
memory,
175-
e.as_canonical_u32(),
176-
b,
177-
&mut record.reads_aux[0].prev_timestamp,
178-
);
173+
let rs1 = tracing_read_or_imm_native(memory, e, b, &mut record.reads_aux[0].prev_timestamp);
179174
record.c = c;
180-
let rs2 = tracing_read_or_imm_native(
181-
memory,
182-
f.as_canonical_u32(),
183-
c,
184-
&mut record.reads_aux[1].prev_timestamp,
185-
);
175+
let rs2 = tracing_read_or_imm_native(memory, f, c, &mut record.reads_aux[1].prev_timestamp);
186176
[rs1, rs2]
187177
}
188178

‎extensions/native/circuit/src/adapters/branch_native_adapter.rs‎

Lines changed: 2 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -165,19 +165,9 @@ where
165165
let &Instruction { a, b, d, e, .. } = instruction;
166166

167167
record.ptrs[0] = a;
168-
let rs1 = tracing_read_or_imm_native(
169-
memory,
170-
d.as_canonical_u32(),
171-
a,
172-
&mut record.reads_aux[0].prev_timestamp,
173-
);
168+
let rs1 = tracing_read_or_imm_native(memory, d, a, &mut record.reads_aux[0].prev_timestamp);
174169
record.ptrs[1] = b;
175-
let rs2 = tracing_read_or_imm_native(
176-
memory,
177-
e.as_canonical_u32(),
178-
b,
179-
&mut record.reads_aux[1].prev_timestamp,
180-
);
170+
let rs2 = tracing_read_or_imm_native(memory, e, b, &mut record.reads_aux[1].prev_timestamp);
181171
[rs1, rs2]
182172
}
183173

‎extensions/rv32im/circuit/src/adapters/loadstore.rs‎

Lines changed: 34 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@ use openvm_circuit_primitives_derive::AlignedBorrow;
2929
use openvm_instructions::{
3030
instruction::Instruction,
3131
program::DEFAULT_PC_STEP,
32-
riscv::{RV32_IMM_AS, RV32_REGISTER_AS},
33-
LocalOpcode,
32+
riscv::{RV32_IMM_AS, RV32_MEMORY_AS, RV32_REGISTER_AS},
33+
LocalOpcode, NATIVE_AS,
3434
};
3535
use openvm_rv32im_transpiler::Rv32LoadStoreOpcode::{self, *};
3636
use openvm_stark_backend::{
@@ -366,8 +366,6 @@ where
366366
} = instruction;
367367

368368
debug_assert_eq!(d.as_canonical_u32(), RV32_REGISTER_AS);
369-
debug_assert!(e.as_canonical_u32() != RV32_IMM_AS);
370-
debug_assert!(e.as_canonical_u32() != RV32_REGISTER_AS);
371369

372370
let local_opcode = Rv32LoadStoreOpcode::from_usize(
373371
opcode.local_opcode_idx(Rv32LoadStoreOpcode::CLASS_OFFSET),
@@ -381,7 +379,6 @@ where
381379
&mut record.rs1_aux_record.prev_timestamp,
382380
));
383381

384-
record.mem_as = e.as_canonical_u32() as u8;
385382
record.imm = c.as_canonical_u32() as u16;
386383
record.imm_sign = g.is_one();
387384
let imm_extended = record.imm as u32 + record.imm_sign as u32 * 0xffff0000;
@@ -397,33 +394,39 @@ where
397394
self.pointer_max_bits
398395
);
399396

400-
let read_data = match local_opcode {
401-
LOADW | LOADB | LOADH | LOADBU | LOADHU => tracing_read(
402-
memory,
403-
e.as_canonical_u32(),
404-
ptr_val,
405-
&mut record.read_data_aux.prev_timestamp,
406-
),
407-
STOREW | STOREH | STOREB => tracing_read(
408-
memory,
409-
RV32_REGISTER_AS,
410-
a.as_canonical_u32(),
411-
&mut record.read_data_aux.prev_timestamp,
412-
),
413-
};
414-
415-
// We need to keep values of some cells to keep them unchanged when writing to those cells
416-
let prev_data = match local_opcode {
397+
// prev_data: We need to keep values of some cells to keep them unchanged when writing to
398+
// those cells
399+
let (read_data, prev_data) = match local_opcode {
400+
LOADW | LOADB | LOADH | LOADBU | LOADHU => {
401+
debug_assert_eq!(e, F::from_canonical_u32(RV32_MEMORY_AS));
402+
record.mem_as = RV32_MEMORY_AS as u8;
403+
let read_data = tracing_read(
404+
memory,
405+
RV32_MEMORY_AS,
406+
ptr_val,
407+
&mut record.read_data_aux.prev_timestamp,
408+
);
409+
let prev_data = memory_read(memory.data(), RV32_REGISTER_AS, a.as_canonical_u32())
410+
.map(u32::from);
411+
(read_data, prev_data)
412+
}
417413
STOREW | STOREH | STOREB => {
418-
if e.as_canonical_u32() == 4 {
414+
let e = e.as_canonical_u32();
415+
debug_assert_ne!(e, RV32_IMM_AS);
416+
debug_assert_ne!(e, RV32_REGISTER_AS);
417+
record.mem_as = e as u8;
418+
let read_data = tracing_read(
419+
memory,
420+
RV32_REGISTER_AS,
421+
a.as_canonical_u32(),
422+
&mut record.read_data_aux.prev_timestamp,
423+
);
424+
let prev_data = if e == NATIVE_AS {
419425
memory_read_native(memory.data(), ptr_val).map(|x: F| x.as_canonical_u32())
420426
} else {
421-
memory_read(memory.data(), e.as_canonical_u32(), ptr_val).map(u32::from)
422-
}
423-
}
424-
LOADW | LOADB | LOADH | LOADBU | LOADHU => {
425-
memory_read(memory.data(), d.as_canonical_u32(), a.as_canonical_u32())
426-
.map(u32::from)
427+
memory_read(memory.data(), e, ptr_val).map(u32::from)
428+
};
429+
(read_data, prev_data)
427430
}
428431
};
429432

@@ -448,8 +451,8 @@ where
448451
} = instruction;
449452

450453
debug_assert_eq!(d.as_canonical_u32(), RV32_REGISTER_AS);
451-
debug_assert!(e.as_canonical_u32() != RV32_IMM_AS);
452-
debug_assert!(e.as_canonical_u32() != RV32_REGISTER_AS);
454+
debug_assert_ne!(e.as_canonical_u32(), RV32_IMM_AS);
455+
debug_assert_ne!(e.as_canonical_u32(), RV32_REGISTER_AS);
453456

454457
let local_opcode = Rv32LoadStoreOpcode::from_usize(
455458
opcode.local_opcode_idx(Rv32LoadStoreOpcode::CLASS_OFFSET),

0 commit comments

Comments
 (0)