From cf954b861b49f163dbc096486e9033cefbdefdb4 Mon Sep 17 00:00:00 2001 From: Cedric Jiang Date: Thu, 11 Dec 2025 12:27:39 -0800 Subject: [PATCH 1/2] Simple malloc instruction support --- parasol_cpu/src/proc/assembly.rs | 3 ++ parasol_cpu/src/proc/fhe_processor.rs | 3 ++ parasol_cpu/src/proc/gas_model.rs | 2 +- parasol_cpu/src/proc/ops/malloc.rs | 56 +++++++++++++++++++++++++++ parasol_cpu/src/proc/ops/mod.rs | 1 + 5 files changed, 64 insertions(+), 1 deletion(-) create mode 100644 parasol_cpu/src/proc/ops/malloc.rs diff --git a/parasol_cpu/src/proc/assembly.rs b/parasol_cpu/src/proc/assembly.rs index 295ef98a..24e51477 100644 --- a/parasol_cpu/src/proc/assembly.rs +++ b/parasol_cpu/src/proc/assembly.rs @@ -477,6 +477,9 @@ define_op! { // Raw cmux [0xC1 Cmux (dst dst, 0, Register) (src cond, 0, Register) (src a, 0, Register) (src b, 0, Register)], + // Memory allocation + [0xD1 Malloc (dst dst, 0, Register) (src size, 0, Register)], + // If the a debug handler with the given id is installed, call it passing the `src` register's value. [0xF0 Dbg (src src, 0, Register) (meta handler_id, 32, u32)] } diff --git a/parasol_cpu/src/proc/fhe_processor.rs b/parasol_cpu/src/proc/fhe_processor.rs index d222058f..889b9b51 100644 --- a/parasol_cpu/src/proc/fhe_processor.rs +++ b/parasol_cpu/src/proc/fhe_processor.rs @@ -836,6 +836,9 @@ impl Tomasulo for FheProcessor { Ret() => { Self::retire(&retirement_info, Ok(())); } + Malloc(dst, size) => { + self.malloc(retirement_info, dst, size, &memory, instruction_id, pc); + } Dbg(src, handler_id) => { self.dbg(&retirement_info, src, handler_id, instruction_id, pc); } diff --git a/parasol_cpu/src/proc/gas_model.rs b/parasol_cpu/src/proc/gas_model.rs index d19a6c24..eb6e63fd 100644 --- a/parasol_cpu/src/proc/gas_model.rs +++ b/parasol_cpu/src/proc/gas_model.rs @@ -144,7 +144,7 @@ impl GasModel { // instructions that have trivial cost: either they do not deal with any ciphertext at all // or they don't compute on ciphertext content (just treat them as vector of opaque objects) Load(..) | LoadI(..) | Store(..) | BranchNonZero(..) | BranchZero(..) | Branch(..) - | Move(..) | Dbg(..) | Sext(..) | Zext(..) | Trunc(..) => 1, + | Move(..) | Dbg(..) | Sext(..) | Zext(..) | Trunc(..) | Malloc(..) => 1, // Not has only one input and the cost is non-trivial if that input is ciphertext Not(_, input) => self.figure_out_gas(&[input], input, op), diff --git a/parasol_cpu/src/proc/ops/malloc.rs b/parasol_cpu/src/proc/ops/malloc.rs new file mode 100644 index 00000000..a91be5ef --- /dev/null +++ b/parasol_cpu/src/proc/ops/malloc.rs @@ -0,0 +1,56 @@ +use std::sync::Arc; + +use crate::{ + Error, Memory, Register, Result, + proc::{DispatchIsaOp, fhe_processor::FheProcessor}, + tomasulo::{registers::RobEntryRef, tomasulo_processor::RetirementInfo}, + unwrap_registers, +}; + +impl FheProcessor { + pub fn malloc( + &mut self, + retirement_info: RetirementInfo, + dst: RobEntryRef, + size: RobEntryRef, + memory: &Arc, + instruction_id: usize, + pc: u32, + ) { + let malloc_impl = || -> Result<()> { + unwrap_registers!((mut dst) (size)); + + match size { + Register::Plaintext { val, width } => { + if *width != 32 || *val >= u32::MAX as u128 { + return Err(Error::IllegalOperands { + inst_id: instruction_id, + pc, + }); + } + + let ptr = memory.try_allocate(*val as u32)?; + + *dst = Register::Plaintext { + val: ptr.0 as u128, + width: 32, + }; + + FheProcessor::retire(&retirement_info, Ok(())); + } + _ => { + return Err(Error::IllegalOperands { + inst_id: instruction_id, + pc, + }); + } + }; + + Ok(()) + }; + + if let Err(e) = malloc_impl() { + FheProcessor::retire(&retirement_info, Err(e)); + } + } +} diff --git a/parasol_cpu/src/proc/ops/mod.rs b/parasol_cpu/src/proc/ops/mod.rs index cbc40c29..3b5bcf79 100644 --- a/parasol_cpu/src/proc/ops/mod.rs +++ b/parasol_cpu/src/proc/ops/mod.rs @@ -23,6 +23,7 @@ mod comparisons; mod dbg; mod load; mod loadi; +mod malloc; mod mov; mod mul; mod neg; From d8d082f56631c0e4c0ddb7905bfce7e7f918bf23 Mon Sep 17 00:00:00 2001 From: Cedric Jiang Date: Thu, 11 Dec 2025 23:27:25 -0800 Subject: [PATCH 2/2] Add test; cosmetic --- parasol_cpu/src/proc/fhe_processor.rs | 2 +- parasol_cpu/src/proc/ops/malloc.rs | 2 +- parasol_cpu/src/proc/tests/malloc.rs | 136 ++++++++++++++++++++++++++ parasol_cpu/src/proc/tests/mod.rs | 1 + 4 files changed, 139 insertions(+), 2 deletions(-) create mode 100644 parasol_cpu/src/proc/tests/malloc.rs diff --git a/parasol_cpu/src/proc/fhe_processor.rs b/parasol_cpu/src/proc/fhe_processor.rs index 889b9b51..21c62953 100644 --- a/parasol_cpu/src/proc/fhe_processor.rs +++ b/parasol_cpu/src/proc/fhe_processor.rs @@ -837,7 +837,7 @@ impl Tomasulo for FheProcessor { Self::retire(&retirement_info, Ok(())); } Malloc(dst, size) => { - self.malloc(retirement_info, dst, size, &memory, instruction_id, pc); + self.malloc(retirement_info, &memory, dst, size, instruction_id, pc); } Dbg(src, handler_id) => { self.dbg(&retirement_info, src, handler_id, instruction_id, pc); diff --git a/parasol_cpu/src/proc/ops/malloc.rs b/parasol_cpu/src/proc/ops/malloc.rs index a91be5ef..38588edd 100644 --- a/parasol_cpu/src/proc/ops/malloc.rs +++ b/parasol_cpu/src/proc/ops/malloc.rs @@ -11,9 +11,9 @@ impl FheProcessor { pub fn malloc( &mut self, retirement_info: RetirementInfo, + memory: &Arc, dst: RobEntryRef, size: RobEntryRef, - memory: &Arc, instruction_id: usize, pc: u32, ) { diff --git a/parasol_cpu/src/proc/tests/malloc.rs b/parasol_cpu/src/proc/tests/malloc.rs new file mode 100644 index 00000000..a391f7dd --- /dev/null +++ b/parasol_cpu/src/proc/tests/malloc.rs @@ -0,0 +1,136 @@ +use std::sync::Arc; + +use crate::{ + ArgsBuilder, Error, Memory, proc::IsaOp, register_names::*, test_utils::make_computer_128, +}; + +use parasol_runtime::{fluent::UInt8, test_utils::get_secret_keys_128}; + +#[test] +fn can_malloc_from_plaintext_size_and_pass_plaintext_data() { + let (mut proc, _) = make_computer_128(); + + let memory = Arc::new(Memory::new_default_stack()); + + let program = memory.allocate_program(&[ + IsaOp::Load(T0, SP, 32, 0), + IsaOp::Load(T1, SP, 32, 4), + IsaOp::LoadI(T2, 1, 32), + IsaOp::Malloc(T3, T2), + IsaOp::Load(T3, T0, 8, 0), + IsaOp::Store(T1, T3, 8, 0), + IsaOp::Ret(), + ]); + + let input_ptr = memory.try_allocate(4).unwrap(); + let output_ptr = memory.try_allocate(4).unwrap(); + let args = ArgsBuilder::new() + .arg(input_ptr) + .arg(output_ptr) + .no_return_value(); + + let val: u8 = 42; + + memory.try_write_type(input_ptr, &val).unwrap(); + + proc.run_program(program, &memory, args).unwrap(); + + let output: u8 = memory.try_load_type(output_ptr).unwrap(); + + assert_eq!(output, 42); +} + +#[test] +fn can_malloc_from_plaintext_size_and_pass_ciphertext_data() { + let (mut proc, enc) = make_computer_128(); + let sk = get_secret_keys_128(); + + let memory = Arc::new(Memory::new_default_stack()); + + let program = memory.allocate_program(&[ + IsaOp::Load(T0, SP, 32, 0), + IsaOp::Load(T1, SP, 32, 4), + IsaOp::LoadI(T2, 1, 32), + IsaOp::Malloc(T3, T2), + IsaOp::Load(T3, T0, 8, 0), + IsaOp::Store(T1, T3, 8, 0), + IsaOp::Ret(), + ]); + + let input_ptr = memory.try_allocate(4).unwrap(); + let output_ptr = memory.try_allocate(4).unwrap(); + let args = ArgsBuilder::new() + .arg(input_ptr) + .arg(output_ptr) + .no_return_value(); + + let val = 42; + + memory + .try_write_type(input_ptr, &UInt8::encrypt_secret(val, &enc, &sk)) + .unwrap(); + + proc.run_program(program, &memory, args).unwrap(); + + let output: UInt8 = memory.try_load_type(output_ptr).unwrap(); + + assert_eq!(output.decrypt(&enc, &sk), val); +} + +#[test] +fn cannot_malloc_from_ciphertext_size() { + let (mut proc, enc) = make_computer_128(); + let sk = get_secret_keys_128(); + + let mut case = || { + let memory = Arc::new(Memory::new_default_stack()); + + let program = memory.allocate_program(&[ + IsaOp::Load(T0, SP, 32, 0), + IsaOp::Load(T0, T0, 32, 0), + IsaOp::Malloc(T1, T0), + IsaOp::Ret(), + ]); + + let encrypted: [UInt8; 4] = [0; 4] + .into_iter() + .map(|x| UInt8::encrypt_secret(x, &enc, &sk)) + .collect::>() + .try_into() + .unwrap_or_else(|_| unreachable!()); + let ciphertext_size_ptr = memory.try_allocate_type(&encrypted).unwrap(); + + let args = ArgsBuilder::new() + .arg(ciphertext_size_ptr) + .no_return_value(); + + let result = proc.run_program(program, &memory, args); + + assert!(matches!(result, Err(Error::IllegalOperands { .. }))); + }; + + case(); +} + +#[test] +fn cannot_malloc_from_plaintext_size_with_wrong_width() { + let (mut proc, _) = make_computer_128(); + + let mut case = || { + let memory = Arc::new(Memory::new_default_stack()); + + let program = memory.allocate_program(&[ + IsaOp::LoadI(T0, 4, 16), + IsaOp::Malloc(T1, T0), + IsaOp::Ret(), + ]); + + let args = ArgsBuilder::new().no_return_value(); + + let result = proc.run_program(program, &memory, args); + + assert!(matches!(result, Err(Error::IllegalOperands { .. }))); + }; + + case(); +} diff --git a/parasol_cpu/src/proc/tests/mod.rs b/parasol_cpu/src/proc/tests/mod.rs index 5a160d5a..dc43f2e4 100644 --- a/parasol_cpu/src/proc/tests/mod.rs +++ b/parasol_cpu/src/proc/tests/mod.rs @@ -9,6 +9,7 @@ mod comparisons; mod dbg; mod faults; mod load_store; +mod malloc; mod mov; mod mul; mod neg;