Simple malloc instruction support

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)]
 }

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, &memory, dst, size, instruction_id, pc);
+            }
             Dbg(src, handler_id) => {
                 self.dbg(&retirement_info, src, handler_id, instruction_id, pc);
             }

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),

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<DispatchIsaOp>,
+        memory: &Arc<Memory>,
+        dst: RobEntryRef<Register>,
+        size: RobEntryRef<Register>,
+        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));
+        }
+    }
+}

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;

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::<Vec<_>>()
+            .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();
+}

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;