Implement handling of globals.

There appear two be only two instructions involving globals: store and load. So
instead of hiding globals inside constants, like LLVM does, store them in a
separate vector and add two special opcodes StoreGlobal and LoadGlobal.

tests/ir_lowering/empty.ll

@@ -3,6 +3,7 @@
 ;     # IR format version: 0
 ;     # Num funcs: 1
 ;     # Num consts: 0
+;     # Num globals: 0
 ;     # Num types: 2
 ;
 ;     func main() {

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)]
@@ -258,12 +264,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)]
@@ -282,6 +308,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),
 }
@@ -323,6 +351,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),
         }
@@ -744,6 +773,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({}, tls={})", self.name, self.is_threadlocal)
+    }
+}
+
 /// A bytecode module.
 ///
 /// This is the top-level container for the bytecode.
@@ -763,6 +807,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>,
@@ -861,6 +909,7 @@ impl Module {
         ret.push_str(&format!("# IR format version: {}\n", self.version));
         ret.push_str(&format!("# Num funcs: {}\n", self.funcs.len()));
         ret.push_str(&format!("# Num consts: {}\n", self.consts.len()));
+        ret.push_str(&format!("# Num globals: {}\n", self.globals.len()));
         ret.push_str(&format!("# Num types: {}\n", self.types.len()));
 
         for func in &self.funcs {
@@ -1129,6 +1178,16 @@ mod tests {
         // bytes:
         data.write_u32::<NativeEndian>(50).unwrap();
 
+        // GLOBALS
+        // num_globals:
+        write_native_usize(&mut data, 1);
+
+        // GLOBAL 1
+        // is_threadlocal:
+        let _ = data.write_u8(0);
+        // name:
+        write_str(&mut data, "aaa");
+
         // TYPES
         // num_types:
         write_native_usize(&mut data, 7);
@@ -1198,6 +1257,7 @@ mod tests {
 # IR format version: 0
 # Num funcs: 2
 # Num consts: 3
+# Num globals: 1
 # Num types: 7
 
 func foo($arg0: ptr, $arg1: i32) -> i32 {

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,21 +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,
         }
     }
 }
@@ -233,10 +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),
         }
     }
 }
@@ -252,7 +263,9 @@ macro_rules! instr {
 }
 
 instr!(Load, LoadInstruction);
+instr!(LoadGlobal, LoadGlobalInstruction);
 instr!(Store, StoreInstruction);
+instr!(StoreGlobal, StoreGlobalInstruction);
 instr!(LoadArg, LoadArgInstruction);
 instr!(Call, CallInstruction);
 instr!(PtrAdd, PtrAddInstruction);
@@ -317,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
@@ -418,6 +463,43 @@ impl fmt::Display for StoreInstruction {
     }
 }
 
+/// 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
@@ -591,6 +673,11 @@ mod tests {
     #[test]
     fn instr_size() {
         assert_eq!(mem::size_of::<CallInstruction>(), 7);
+        assert_eq!(mem::size_of::<StoreInstruction>(), 4);
+        assert_eq!(mem::size_of::<LoadInstruction>(), 6);
+        assert_eq!(mem::size_of::<LoadGlobalInstruction>(), 6);
+        assert_eq!(mem::size_of::<StoreGlobalInstruction>(), 6);
+        assert_eq!(mem::size_of::<PtrAddInstruction>(), 6);
         assert!(mem::size_of::<Instruction>() <= mem::size_of::<u64>());
     }
 

ykrt/src/compile/jitc_yk/trace_builder.rs

@@ -146,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(
@@ -169,8 +172,13 @@ impl<'a> TraceBuilder<'a> {
         inst: &aot_ir::Instruction,
     ) -> Result<jit_ir::Instruction, CompilationError> {
         let val = self.handle_operand(inst.operand(0))?;
-        let ptr = self.handle_operand(inst.operand(1))?;
-        Ok(jit_ir::StoreInstruction::new(val, ptr).into())
+        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.
@@ -197,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 {