execution models explorations

main.rs

use std::time::Instant;

use cranelift_jit::{JITBuilder, JITModule};
use cranelift::prelude::*;
use cranelift_module::{Linkage, Module};

// row:
// |... sized columns/unsized col meta | ... unsized column

enum Value {
    String(String),
    UInt64(u64),
}

impl Value {
    fn serialized_size(&self) -> usize {
        match self {
            Value::String(s) => s.len() + size_of::<u64>(),
            Value::UInt64(_) => size_of::<u64>(),
        }
    }

    fn is_unsized(&self) -> bool {
        match self {
            Value::String(_) => true,
            Value::UInt64(_) => false,
        }
    }

    fn len(&self) -> usize {
        match self {
            Value::String(s) => s.len(),
            Value::UInt64(_) => todo!(),
        }
    }
}

struct Row {
    values: Vec<Value>,
}

impl Row {
    fn serialized_size(&self) -> usize {
        self.values.iter().map(Value::serialized_size).sum()
    }

    fn serialize(&self) -> Vec<u8> {
        let serialized_size = self.serialized_size();
        let mut out = vec![0u8; serialized_size];
        let unsized_data: usize = self.values.iter().filter(|v| v.is_unsized()).map(|v| v.len()).sum();
        let mut unsized_start_offset = serialized_size - unsized_data;
        let mut current_offset = 0;
        for v in self.values.iter() {
            match v {
                Value::String(s) => {
                    let len = s.len();
                    let offset = unsized_start_offset;
                    unsized_start_offset += len;
                    let it = len << 32 | offset;
                    out[offset..offset + len].copy_from_slice(s.as_bytes());
                    out[current_offset..current_offset + size_of::<u64>()].copy_from_slice(&it.to_le_bytes());
                    current_offset += size_of::<u64>();
                },
                Value::UInt64(it) => {
                    out[current_offset..current_offset + size_of::<u64>()].copy_from_slice(&it.to_le_bytes());
                },
            }
        }

        out
    }
}

enum DataType {
    Int,
    String,
}

enum Expr {
    Col { idx: usize, ty: DataType },
    Add { lhs: Box<Self>, rhs: Box<Self>, ty: DataType },
    Mult { lhs: Box<Self>, rhs: Box<Self>, ty: DataType },
    ConsInt { i: u64 }
}

struct EmitCtx<'a, 'b> {
    stack: Vec<cranelift::prelude::Value>,
    builder: &'a mut FunctionBuilder<'b>,
    row: cranelift::prelude::Value,
}

struct EmitVmCtx {
    stack: Vec<usize>,
    reg_no: usize,
    out: Vec<Op>,
}

impl Expr {
    fn emit_vm(&self, pgm: &mut EmitVmCtx) {
        match self {
            Expr::Col { idx, ty } => {
                let reg = pgm.reg_no;
                pgm.reg_no += 1;
                pgm.out.push(Op::Col { offset: *idx, reg });
                pgm.stack.push(reg);
            },
            Expr::Add { lhs, rhs, ty } => {
                lhs.emit_vm(pgm);
                rhs.emit_vm(pgm);
                let rhs = pgm.stack.pop().unwrap();
                let lhs = pgm.stack.pop().unwrap();
                pgm.out.push(Op::Add { r1: lhs, r2: rhs });
                pgm.stack.push(lhs);
            },
            Expr::Mult { lhs, rhs, ty } => {
                lhs.emit_vm(pgm);
                rhs.emit_vm(pgm);
                let rhs = pgm.stack.pop().unwrap();
                let lhs = pgm.stack.pop().unwrap();
                pgm.out.push(Op::Mult { r1: lhs, r2: rhs });
                pgm.stack.push(lhs);
            },
            Expr::ConsInt { i } => todo!(),
        }
    }

    fn emit(&self, ctx: &mut EmitCtx) {
        match self {
            Expr::Col { idx, ty } => {
                match ty {
                    DataType::Int => {
                        let mut mem_flags = MemFlags::trusted();
                        mem_flags.set_readonly();
                        mem_flags.set_heap();

                        let val = ctx.builder.ins().load(Type::int(64).unwrap(), mem_flags,ctx.row, *idx as i32);
                        ctx.stack.push(val);
                    },
                    DataType::String => todo!(),
                }
            },
            Expr::Add { lhs, rhs, ty } => {
                lhs.emit(ctx);
                rhs.emit(ctx);
                let rhs = ctx.stack.pop().unwrap();
                let lhs = ctx.stack.pop().unwrap();
                match ty {
                    DataType::Int => {
                        ctx.stack.push(ctx.builder.ins().iadd(lhs, rhs));
                    },
                    DataType::String => todo!(),
                }
            },
            Expr::Mult { lhs, rhs, ty } => {
                lhs.emit(ctx);
                rhs.emit(ctx);
                let rhs = ctx.stack.pop().unwrap();
                let lhs = ctx.stack.pop().unwrap();
                match ty {
                    DataType::Int => {
                        ctx.stack.push(ctx.builder.ins().imul(lhs, rhs));
                    },
                    DataType::String => todo!(),
                }
            },
            Expr::ConsInt { i } => {
                ctx.stack.push(ctx.builder.ins().iconst(Type::int(64).unwrap(), *i as i64));
            },
        }

    }
}

type RawRow<'a> = &'a [u8];

struct Jit {
    module: JITModule,
}

impl Jit {
    fn new() -> Self {
        let mut flag_builder = settings::builder();
        flag_builder.set("use_colocated_libcalls", "false").unwrap();
        flag_builder.set("is_pic", "false").unwrap();
        flag_builder.set("opt_level", "none").unwrap();
        let isa_builder = cranelift_native::builder().unwrap_or_else(|msg| {
            panic!("host machine is not supported: {}", msg);
        });
        let isa = isa_builder
            .finish(settings::Flags::new(flag_builder))
            .unwrap();
        let builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
        let mut module = JITModule::new(builder);

        Self {
            module,
        }
    }

    fn compile_fun(&mut self, e: &Expr) -> for<'a> fn(RawRow<'a>) -> u64 {
        // let guard = pprof::ProfilerGuardBuilder::default().frequency(10000).blocklist(&["libc", "libgcc", "pthread", "vdso"]).build().unwrap(); 
        let before = Instant::now();
        let mut ctx = self.module.make_context();
        dbg!(before.elapsed());
        let ptr = self.module.target_config().pointer_type();
        dbg!(before.elapsed());
        ctx.func.signature.params.push(AbiParam::new(ptr));
        ctx.func.signature.returns.push(AbiParam::new(Type::int(64).unwrap()));
        let mut builder_ctx = FunctionBuilderContext::new();
        dbg!(before.elapsed());

        let mut builder = FunctionBuilder::new(&mut ctx.func, &mut builder_ctx);

        let entry_block = builder.create_block();

        // Since this is the entry block, add block parameters corresponding to
        // the function's parameters.
        //
        // TODO: Streamline the API here.
        builder.append_block_params_for_function_params(entry_block);
        builder.switch_to_block(entry_block);
        builder.seal_block(entry_block);

        dbg!(before.elapsed());
        let val = builder.block_params(entry_block)[0];
        let var = Variable::new(0);
        builder.declare_var(var, ptr);
        builder.def_var(var, val);

        dbg!(before.elapsed());
        let ret = {
            let mut ctx = EmitCtx {
                stack: Vec::with_capacity(32),
                builder: &mut builder,
                row: val,
            };

            e.emit(&mut ctx);

            ctx.stack.pop().unwrap()
        };

        dbg!(before.elapsed());
        builder.ins().return_(&[ret]);

        builder.finalize();
        dbg!(before.elapsed());
        let id = self.module
            .declare_function(&uuid::Uuid::new_v4().to_string(), Linkage::Export, &ctx.func.signature)
            .map_err(|e| e.to_string()).unwrap();

        dbg!(before.elapsed());
        self.module
            .define_function(id, &mut ctx)
            .unwrap();

        // Now that compilation is finished, we can clear out the context state.
        self.module.clear_context(&mut ctx);

        // Finalize the functions which we just defined, which resolves any
        // outstanding relocations (patching in addresses, now that they're
        // available).
        self.module.finalize_definitions().unwrap();

        // We can now retrieve a pointer to the machine code.
        let f = self.module.get_finalized_function(id);

        // if let Ok(report) = guard.report().build() {
        //     let file = std::fs::File::create(format!("flamegraph-{}.svg", uuid::Uuid::new_v4())).unwrap();
        //     report.flamegraph(file).unwrap();
        // };

        unsafe { std::mem::transmute::<*const u8, for<'a> fn(RawRow<'a>) -> u64>(f) }
    }
}


fn main() {
    let e = Expr::Mult { 
        lhs: Expr::Add { 
            lhs: Expr::Col { idx: 0, ty: DataType::Int }.into(),
            rhs: Expr::Col { idx: 8, ty: DataType::Int }.into(),
            ty: DataType::Int
        }.into(),
        rhs: Expr::Col { idx: 16, ty: DataType::Int }.into(),
        ty: DataType::Int
    };

    let e2 = Expr::Add { 
        lhs: Expr::Add { 
            lhs: Expr::Col { idx: 0, ty: DataType::Int }.into(),
            rhs: Expr::Col { idx: 8, ty: DataType::Int }.into(),
            ty: DataType::Int
        }.into(),
        rhs: Expr::Col { idx: 16, ty: DataType::Int }.into(),
        ty: DataType::Int
    };

    let mut jit = Jit::new();

    let before= Instant::now();
    let f = jit.compile_fun(&e);
    dbg!(before.elapsed());

    let before= Instant::now();
    let f = jit.compile_fun(&e2);
    dbg!(before.elapsed());

    let mut out= Vec::new();
    for _ in 0..65536 {
        out.extend(1u64.to_le_bytes());
        out.extend(2u64.to_le_bytes());
        out.extend(3u64.to_le_bytes());
    }

    
    let before= Instant::now();
    out.chunks(3*size_of::<u64>()).for_each(|e|  { f(e); });
    dbg!(before.elapsed());

    let before= Instant::now();
    out.chunks(3*size_of::<u64>()).for_each(|c| {
        let vals: &[u64] = unsafe {
            std::mem::transmute(c)
        };

        let _ = (vals[0] + vals[1]) * vals[2];
    });
    dbg!(before.elapsed());

    let before= Instant::now();
    let mut pgm = EmitVmCtx {
        stack: Vec::new(),
        reg_no: 0,
        out: Vec::new(),
    };

    pgm.out.push(Op::Next);
    e.emit_vm(&mut pgm);
    pgm.out.push(Op::Goto { i: 0 });

    let mut vm = Vm {
        regs: vec![0; 16],
    };
    dbg!(before.elapsed());

    let before= Instant::now();
    vm.run(&pgm.out, &out);
    dbg!(before.elapsed());
}

struct Vm {
    regs: Vec<u64>,
}

impl Vm {
    fn run(&mut self, pgm: &Vec<Op>, raw: &[u8]) {
        let mut chunks = raw.chunks(64);
        let mut current: Option<&[u8]> = None;
        let mut ip = 0;
        loop {
            let op = &pgm[ip];
            match op {
                Op::Col { offset, reg } => {
                    self.regs[*reg] = u64::from_le_bytes(current.as_ref().unwrap()[*offset..*offset + 8].try_into().unwrap());
                },
                Op::Add { r1, r2 } => {
                    self.regs[*r1] += self.regs[*r2];
                },
                Op::Mult { r1, r2 } => {
                    self.regs[*r1] *= self.regs[*r2];
                },
                Op::Next => {
                    current = chunks.next();
                    if current.is_none() {
                        break
                    }
                }
                Op::Goto { i } => {
                    ip = *i;
                    continue;
                },
            }

            ip += 1;
        }
    }
}

#[derive(Debug)]
enum Op {
    Col { offset: usize, reg: usize },
    Add { r1: usize, r2: usize },
    Mult { r1: usize, r2: usize },
    Next,
    Goto { i: usize }
}