Skip to content

Commit

Permalink
Merge pull request #185 from near/sha_full
Browse files Browse the repository at this point in the history 
Fix Loads and Stores for SHA256 benchmark
  • Loading branch information
@aborg-dev
aborg-dev authored Feb 6, 2024
2 parents be4c275 + 2f5bc55 commit 3d0d0a7
Show file tree
Hide file tree
Showing 11 changed files with 910 additions and 187 deletions.
23 changes: 23 additions & 0 deletions cranelift/codegen/src/isa/zkasm/inst/args.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -802,6 +802,19 @@ impl LoadOP {
Self::Fld => 0b011,
}
}
pub(crate) fn width(self) -> u32 {
match self {
Self::I8 => 1,
Self::I16 => 2,
Self::I32 => 4,
Self::U32 => 4,
Self::U8 => 1,
Self::U16 => 2,
Self::U64 => 8,
Self::Flw => unimplemented!(),
Self::Fld => unimplemented!(),
}
}
}

impl StoreOP {
Expand Down Expand Up @@ -843,6 +856,16 @@ impl StoreOP {
Self::Fsd => 0b011,
}
}
pub(crate) fn width(self) -> u32 {
match self {
Self::I8 => 1,
Self::I16 => 2,
Self::I32 => 4,
Self::I64 => 8,
Self::Fsw => unimplemented!(),
Self::Fsd => unimplemented!(),
}
}
}

impl IntSelectOP {
Expand Down
120 changes: 117 additions & 3 deletions cranelift/codegen/src/isa/zkasm/inst/emit.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -591,14 +591,73 @@ impl MachInstEmit for Inst {
match from {
AMode::RegOffset(r, ..) => {
debug_assert_eq!(r, e0());
// TODO(#43): Implement the conversion using verifiable computations.
put_string(
&format!("${{ ({}) % 8 }} => A\n", access_reg_with_offset(r, offset),),
sink,
);
put_string(
&format!(
"${{ ({}) / 8 }} => {}\n",
access_reg_with_offset(r, offset),
reg_name(r)
),
sink,
);

// Handle the case when read spans two slots.
// We do this read first in case `rd == r == E` which will clobber the
// address register after second `MLOAD`.
put_string(
&format!(
"$ => {} :MLOAD(MEM:{})\n",
reg_name(d0()),
access_reg_with_offset(r, 1)
),
sink,
);
put_string(
&format!(
"$ => {} :MLOAD(MEM:{})\n",
reg_name(rd.to_reg()),
access_reg_with_offset(r, offset)
reg_name(r)
),
sink,
);
// TODO(#34): Implement unaligned and narrow reads using verifiable
// computations.
put_string(
&format!(
"${{ {} >> (8 * A) }} => {}\n",
reg_name(rd.to_reg()),
reg_name(rd.to_reg())
),
sink,
);

// Merge in the read from the second slot.
put_string(
&format!(
"${{ (D << (128 - 8 * (A + {}))) | {} }} => {}\n",
op.width(),
reg_name(rd.to_reg()),
reg_name(rd.to_reg()),
),
sink,
);

// Mask the value to the width of the resulting type.
if op.width() < 8 {
put_string(
&format!(
"${{ {} & ((1 << {}) - 1) }} => {}\n",
reg_name(rd.to_reg()),
8 * op.width(),
reg_name(rd.to_reg()),
),
sink,
);
}
}
AMode::SPOffset(..) | AMode::NominalSPOffset(..) | AMode::FPOffset(..) => {
assert_eq!(offset % 8, 0);
Expand Down Expand Up @@ -629,11 +688,66 @@ impl MachInstEmit for Inst {
match to {
AMode::RegOffset(r, ..) => {
debug_assert_eq!(r, e0());
// TODO(#43): Implement the conversion using verifiable computations.
put_string(
&format!("${{ ({}) % 8 }} => A\n", access_reg_with_offset(r, offset)),
sink,
);
put_string(
&format!(
"${{ ({}) / 8 }} => {}\n",
access_reg_with_offset(r, offset),
reg_name(r)
),
sink,
);

// TODO(#34): Implement unaligned and narrow writes using verifiable
// computations.
let width = op.width() as i64;
if width < 8 {
put_string(
&format!(
"${{ {} & ((1 << {}) - 1) }} => {}\n",
reg_name(src),
8 * width,
reg_name(src),
),
sink,
);
}
put_string(
&format!("$ => {} :MLOAD(MEM:{})\n", reg_name(d0()), reg_name(r)),
sink,
);
put_string(
&format!(
"${{ (D & ~(((1 << {}) - 1) << (8 * A))) | ({} << (8 * A)) }} :MSTORE(MEM:{})\n",
8 * width,
reg_name(src),
reg_name(r),
),
sink,
);

// Handle the case when write spans two slots.
put_string(
&format!("${{ ((A + {width}) > 8) ? (A + {width} - 8) : 0 }} => A\n"),
sink,
);
put_string(
&format!(
"$ => {} :MLOAD(MEM:{})\n",
reg_name(d0()),
access_reg_with_offset(r, 1)
),
sink,
);
put_string(
&format!(
"{} :MSTORE(MEM:{})\n",
"${{ (D & ~((1 << (8 * A)) - 1)) | ({} & ((1 << (8 * A)) - 1)) }} :MSTORE(MEM:{})\n",
reg_name(src),
access_reg_with_offset(r, offset)
access_reg_with_offset(r, 1),
),
sink,
);
Expand Down
10 changes: 9 additions & 1 deletion cranelift/codegen/src/isa/zkasm/inst/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -336,13 +336,21 @@ fn zkasm_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut OperandC
if let Some(r) = from.get_allocatable_register() {
collector.reg_fixed_use(r, e0());
}
let mut clobbered = PRegSet::empty();
clobbered.add(a0().to_real_reg().unwrap().into());
clobbered.add(d0().to_real_reg().unwrap().into());
collector.reg_clobbers(clobbered);
collector.reg_def(rd);
}
&Inst::Store { to, src, .. } => {
if let Some(r) = to.get_allocatable_register() {
collector.reg_fixed_use(r, e0());
}
collector.reg_use(src);
let mut clobbered = PRegSet::empty();
clobbered.add(a0().to_real_reg().unwrap().into());
clobbered.add(d0().to_real_reg().unwrap().into());
collector.reg_clobbers(clobbered);
collector.reg_late_use(src);
}
&Inst::Args { ref args } => {
for arg in args {
Expand Down
2 changes: 1 addition & 1 deletion cranelift/filetests/src/test_zkasm.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ mod tests {

// Generate const data segments definitions.
for (offset, data) in data_segments {
program.push(format!(" {offset} => E"));
program.push(format!(" {} => E", offset / 8));
// Each slot stores 8 consecutive u8 numbers, with earlier addresses stored in lower
// bits.
for (i, chunk) in data.chunks(8).enumerate() {
Expand Down
Loading

0 comments on commit 3d0d0a7

Please sign in to comment.