Merge pull request #965 from ptersilie/more-bytecode2

Handle some more opcodes.

ykrt/src/compile/jitc_yk/aot_ir.rs

@@ -75,6 +75,12 @@ index!(InstrIdx);
 pub(crate) struct ConstIdx(usize);
 index!(ConstIdx);
 
+/// A global variable index.
+#[deku_derive(DekuRead)]
+#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
+pub(crate) struct GlobalIdx(usize);
+index!(GlobalIdx);
+
 /// A function argument index.
 #[deku_derive(DekuRead)]
 #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
@@ -128,13 +134,13 @@ pub(crate) enum Opcode {
     Store,
     Alloca,
     Call,
-    GetElementPtr,
     Br,
     CondBr,
     Icmp,
     BinaryOperator,
     Ret,
     InsertValue,
+    PtrAdd,
     Unimplemented = 255,
 }
 
@@ -259,12 +265,32 @@ impl IRDisplay for ArgOperand {
     }
 }
 
+/// An operand that describes a global variable.
+#[deku_derive(DekuRead)]
+#[derive(Debug)]
+pub(crate) struct GlobalOperand {
+    global_idx: GlobalIdx,
+}
+
+impl GlobalOperand {
+    pub(crate) fn index(&self) -> GlobalIdx {
+        self.global_idx
+    }
+}
+
+impl IRDisplay for GlobalOperand {
+    fn to_str(&self, m: &Module) -> String {
+        m.globals[self.global_idx].to_str(m)
+    }
+}
+
 const OPKIND_CONST: u8 = 0;
 const OPKIND_LOCAL_VARIABLE: u8 = 1;
 const OPKIND_TYPE: u8 = 2;
 const OPKIND_FUNCTION: u8 = 3;
 const OPKIND_BLOCK: u8 = 4;
 const OPKIND_ARG: u8 = 5;
+const OPKIND_GLOBAL: u8 = 6;
 const OPKIND_UNIMPLEMENTED: u8 = 255;
 
 #[deku_derive(DekuRead)]
@@ -283,6 +309,8 @@ pub(crate) enum Operand {
     Block(BlockOperand),
     #[deku(id = "OPKIND_ARG")]
     Arg(ArgOperand),
+    #[deku(id = "OPKIND_GLOBAL")]
+    Global(GlobalOperand),
     #[deku(id = "OPKIND_UNIMPLEMENTED")]
     Unimplemented(#[deku(until = "|v: &u8| *v == 0", map = "deserialise_string")] String),
 }
@@ -324,6 +352,7 @@ impl IRDisplay for Operand {
             Self::Type(t) => m.types[t.type_idx].to_str(m),
             Self::Function(f) => m.funcs[f.func_idx].name.to_owned(),
             Self::Block(bb) => bb.to_str(m),
+            Self::Global(g) => g.to_str(m),
             Self::Arg(a) => a.to_str(m),
             Self::Unimplemented(s) => format!("?op<{}>", s),
         }
@@ -382,8 +411,8 @@ impl Instruction {
         self.opcode == Opcode::Store
     }
 
-    pub(crate) fn is_gep(&self) -> bool {
-        self.opcode == Opcode::GetElementPtr
+    pub(crate) fn is_ptr_add(&self) -> bool {
+        self.opcode == Opcode::PtrAdd
     }
 
     pub(crate) fn is_control_point(&self, aot_mod: &Module) -> bool {
@@ -745,6 +774,21 @@ impl IRDisplay for Constant {
     }
 }
 
+/// A constant.
+#[deku_derive(DekuRead)]
+#[derive(Debug)]
+pub(crate) struct Global {
+    is_threadlocal: bool,
+    #[deku(until = "|v: &u8| *v == 0", map = "deserialise_string")]
+    name: String,
+}
+
+impl IRDisplay for Global {
+    fn to_str(&self, m: &Module) -> String {
+        format!("Global({})", self.name)
+    }
+}
+
 /// A bytecode module.
 ///
 /// This is the top-level container for the bytecode.
@@ -764,6 +808,10 @@ pub(crate) struct Module {
     #[deku(count = "num_consts", map = "deserialise_into_ti_vec")]
     consts: TiVec<ConstIdx, Constant>,
     #[deku(temp)]
+    num_globals: usize,
+    #[deku(count = "num_globals", map = "deserialise_into_ti_vec")]
+    globals: TiVec<GlobalIdx, Global>,
+    #[deku(temp)]
     num_types: usize,
     #[deku(count = "num_types", map = "deserialise_into_ti_vec")]
     types: TiVec<TypeIdx, Type>,
@@ -1030,7 +1078,7 @@ mod tests {
         // type_idx:
         write_native_usize(&mut data, 2);
         // opcode:
-        data.write_u8(Opcode::GetElementPtr as u8).unwrap();
+        data.write_u8(Opcode::PtrAdd as u8).unwrap();
         // num_operands:
         data.write_u32::<NativeEndian>(1).unwrap();
         // OPERAND 0

ykrt/src/compile/jitc_yk/jit_ir.rs

@@ -130,6 +130,11 @@ index_16bit!(ExtraArgsIdx);
 pub(crate) struct ConstIdx(u16);
 index_16bit!(ConstIdx);
 
+/// A global index.
+#[derive(Copy, Clone, Debug, PartialEq)]
+pub(crate) struct GlobalIdx(U24);
+index_24bit!(GlobalIdx);
+
 /// An instruction index.
 ///
 /// One of these is an index into the [Module::instrs].
@@ -210,17 +215,25 @@ pub(crate) enum Constant {
 #[derive(Debug)]
 pub enum Instruction {
     Load(LoadInstruction),
+    LoadGlobal(LoadGlobalInstruction),
     LoadArg(LoadArgInstruction),
     Call(CallInstruction),
+    PtrAdd(PtrAddInstruction),
+    Store(StoreInstruction),
+    StoreGlobal(StoreGlobalInstruction),
 }
 
 impl Instruction {
     /// Returns `true` if the instruction defines a local variable.
     pub(crate) fn is_def(&self) -> bool {
         match self {
             Self::Load(..) => true,
+            Self::LoadGlobal(..) => true,
             Self::LoadArg(..) => true,
             Self::Call(..) => true, // FIXME: May or may not define. Ask func sig.
+            Self::PtrAdd(..) => true,
+            Self::Store(..) => false,
+            Self::StoreGlobal(..) => false,
         }
     }
 }
@@ -229,8 +242,12 @@ impl fmt::Display for Instruction {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Self::Load(i) => write!(f, "{}", i),
+            Self::LoadGlobal(i) => write!(f, "{}", i),
             Self::LoadArg(i) => write!(f, "{}", i),
             Self::Call(i) => write!(f, "{}", i),
+            Self::PtrAdd(i) => write!(f, "{}", i),
+            Self::Store(i) => write!(f, "{}", i),
+            Self::StoreGlobal(i) => write!(f, "{}", i),
         }
     }
 }
@@ -246,8 +263,12 @@ macro_rules! instr {
 }
 
 instr!(Load, LoadInstruction);
+instr!(LoadGlobal, LoadGlobalInstruction);
+instr!(Store, StoreInstruction);
+instr!(StoreGlobal, StoreGlobalInstruction);
 instr!(LoadArg, LoadArgInstruction);
 instr!(Call, CallInstruction);
+instr!(PtrAdd, PtrAddInstruction);
 
 /// The operands for a [Instruction::Load]
 ///
@@ -309,6 +330,38 @@ impl LoadArgInstruction {
     }
 }
 
+/// The operands for a [Instruction::LoadGlobal]
+///
+/// # Semantics
+///
+/// Loads a value from a given global variable.
+///
+#[derive(Debug)]
+pub struct LoadGlobalInstruction {
+    /// The pointer to load from.
+    global_idx: GlobalIdx,
+    /// The type of the pointee.
+    ty_idx: TypeIdx,
+}
+
+impl LoadGlobalInstruction {
+    pub(crate) fn new(
+        global_idx: aot_ir::GlobalIdx,
+        ty_idx: TypeIdx,
+    ) -> Result<Self, CompilationError> {
+        Ok(Self {
+            global_idx: GlobalIdx::from_aot(global_idx)?,
+            ty_idx,
+        })
+    }
+}
+
+impl fmt::Display for LoadGlobalInstruction {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "LoadGlobal {}", self.global_idx.0.to_usize())
+    }
+}
+
 /// The operands for a [Instruction::Call]
 ///
 /// # Semantics
@@ -380,6 +433,93 @@ impl CallInstruction {
     }
 }
 
+/// The operands for a [Instruction::Store]
+///
+/// # Semantics
+///
+/// Stores a value into a pointer.
+///
+#[derive(Debug)]
+pub struct StoreInstruction {
+    /// The value to store.
+    val: PackedOperand,
+    /// The pointer to store into.
+    ptr: PackedOperand,
+}
+
+impl StoreInstruction {
+    pub(crate) fn new(val: Operand, ptr: Operand) -> Self {
+        // FIXME: assert type of pointer
+        Self {
+            val: PackedOperand::new(&val),
+            ptr: PackedOperand::new(&ptr),
+        }
+    }
+}
+
+impl fmt::Display for StoreInstruction {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "Store {}, {}", self.val.get(), self.ptr.get())
+    }
+}
+
+/// The operands for a [Instruction::StoreGlobal]
+///
+/// # Semantics
+///
+/// Stores a value into a global.
+///
+#[derive(Debug)]
+pub struct StoreGlobalInstruction {
+    /// The value to store.
+    val: PackedOperand,
+    /// The pointer to store into.
+    global_idx: GlobalIdx,
+}
+
+impl StoreGlobalInstruction {
+    pub(crate) fn new(
+        val: Operand,
+        global_idx: aot_ir::GlobalIdx,
+    ) -> Result<Self, CompilationError> {
+        Ok(Self {
+            val: PackedOperand::new(&val),
+            global_idx: GlobalIdx::from_aot(global_idx)?,
+        })
+    }
+}
+
+impl fmt::Display for StoreGlobalInstruction {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "StoreGlobal {}, {}",
+            self.val.get(),
+            self.global_idx.0.to_usize()
+        )
+    }
+}
+
+/// A pointer offsetting instruction.
+///
+/// # Semantics
+///
+/// Returns a pointer value that is the result of adding the specified (byte) offset to the input
+/// pointer operand.
+#[derive(Debug)]
+pub struct PtrAddInstruction {
+    /// The pointer to offset
+    ptr: PackedOperand,
+    /// The offset.
+    off: u32,
+}
+
+impl fmt::Display for PtrAddInstruction {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "PtrAdd {}, {}", self.ptr.get(), self.off)
+    }
+}
+
 /// The `Module` is the top-level container for JIT IR.
 ///
 /// The IR is conceptually a list of word-sized instructions containing indices into auxiliary

ykrt/src/compile/jitc_yk/trace_builder.rs

@@ -63,7 +63,7 @@ impl<'a> TraceBuilder<'a> {
             if inst.is_store() {
                 last_store = Some(inst);
             }
-            if inst.is_gep() {
+            if inst.is_ptr_add() {
                 let op = inst.operand(0);
                 // unwrap safe: we know the AOT code was produced by ykllvm.
                 if trace_input
@@ -95,6 +95,7 @@ impl<'a> TraceBuilder<'a> {
             let jit_inst = match inst.opcode() {
                 aot_ir::Opcode::Load => self.handle_load(inst),
                 aot_ir::Opcode::Call => self.handle_call(inst),
+                aot_ir::Opcode::Store => self.handle_store(inst),
                 _ => todo!("{:?}", inst),
             }?;
 
@@ -145,11 +146,14 @@ impl<'a> TraceBuilder<'a> {
         &mut self,
         inst: &aot_ir::Instruction,
     ) -> Result<jit_ir::Instruction, CompilationError> {
-        let jit_op = self.handle_operand(inst.operand(0))?;
-        Ok(
-            jit_ir::LoadInstruction::new(jit_op, jit_ir::TypeIdx::from_aot(inst.type_idx())?)
-                .into(),
-        )
+        let ty_idx = jit_ir::TypeIdx::from_aot(inst.type_idx())?;
+        if let aot_ir::Operand::Global(go) = inst.operand(0) {
+            // Generate a special load instruction for globals.
+            Ok(jit_ir::LoadGlobalInstruction::new(go.index(), ty_idx)?.into())
+        } else {
+            let jit_op = self.handle_operand(inst.operand(0))?;
+            Ok(jit_ir::LoadInstruction::new(jit_op, ty_idx).into())
+        }
     }
 
     fn handle_call(
@@ -163,6 +167,20 @@ impl<'a> TraceBuilder<'a> {
         Ok(jit_ir::CallInstruction::new(&mut self.jit_mod, inst.callee(), &args)?.into())
     }
 
+    fn handle_store(
+        &mut self,
+        inst: &aot_ir::Instruction,
+    ) -> Result<jit_ir::Instruction, CompilationError> {
+        let val = self.handle_operand(inst.operand(0))?;
+        if let aot_ir::Operand::Global(go) = inst.operand(1) {
+            // Generate a special store instruction for globals.
+            Ok(jit_ir::StoreGlobalInstruction::new(val, go.index())?.into())
+        } else {
+            let ptr = self.handle_operand(inst.operand(1))?;
+            Ok(jit_ir::StoreInstruction::new(val, ptr).into())
+        }
+    }
+
     /// Entry point for building an IR trace.
     ///
     /// Consumes the trace builder, returning a JIT module.
@@ -187,9 +205,6 @@ impl<'a> TraceBuilder<'a> {
         };
 
         let firstblk = self.lookup_aot_block(&prev);
-        // FIXME: This unwrap assumes the first block is mappable, but Laurie just merged a change
-        // that strips the initial block (the block we return to from the control point), so I
-        // don't think this assumption necessarily holds any more. Investigate.
         self.create_trace_header(self.aot_mod.block(&firstblk.unwrap()))?;
 
         for tblk in self.mtrace {