Add support for WASM Global Variables

cranelift/codegen/src/isa/zkasm/inst.isle

@@ -327,6 +327,12 @@
   (Fsd)
 ))
 
+(type ZkasmBase (enum
+  (Heap)
+  (Global)
+  (Table)
+))
+
 (type AluOPRRR (enum
   ;; base set
   (Add)
@@ -1176,6 +1182,10 @@
 (decl gen_fp_offset_amode (i64 Type) AMode)
 (extern constructor gen_fp_offset_amode gen_fp_offset_amode)
 
+;; Generates a AMode that an access to a global variable at given index.
+(decl gen_global_amode (i64 Type) AMode)
+(extern constructor gen_global_amode gen_global_amode)
+
 ;; Generates an AMode that points to a stack slot + offset.
 (decl gen_stack_slot_amode (StackSlot i64 Type) AMode)
 (extern constructor gen_stack_slot_amode gen_stack_slot_amode)
@@ -1216,6 +1226,9 @@
 (rule 1 (amode_inner (get_stack_pointer) offset ty)
   (gen_sp_offset_amode offset ty))
 
+(rule 1 (amode_inner (symbol_value (zkasm_base ZkasmBase.Global)) index ty)
+  (gen_global_amode index ty))
+
 ;; Similarly if the value is a `stack_addr` we can also turn that into an sp offset.
 (rule 1 (amode_inner (stack_addr ss ss_offset) amode_offset ty)
   (if-let combined_offset (s32_add_fallible ss_offset amode_offset))
@@ -1302,6 +1315,9 @@
 (decl store_op (Type) StoreOP)
 (extern constructor store_op store_op)
 
+(decl zkasm_base (ZkasmBase) GlobalValue)
+(extern extractor zkasm_base zkasm_base)
+
 ;;;; load extern name
 (decl load_ext_name (ExternalName i64) Reg)
 (extern constructor load_ext_name load_ext_name)

cranelift/codegen/src/isa/zkasm/inst/args.rs

@@ -107,6 +107,9 @@ pub enum AMode {
 
     /// A reference to a label.
     Label(MachLabel),
+
+    /// Access to a global variable.
+    Global(i64, Type),
 }
 
 impl AMode {
@@ -117,7 +120,8 @@ impl AMode {
             | AMode::FPOffset(..)
             | AMode::NominalSPOffset(..)
             | AMode::Const(..)
-            | AMode::Label(..) => self,
+            | AMode::Label(..)
+            | AMode::Global(..) => self,
         }
     }
 
@@ -130,7 +134,8 @@ impl AMode {
             | AMode::FPOffset(..)
             | AMode::NominalSPOffset(..)
             | AMode::Const(..)
-            | AMode::Label(..) => None,
+            | AMode::Label(..)
+            | AMode::Global(..) => None,
         }
     }
 
@@ -140,7 +145,7 @@ impl AMode {
             &AMode::SPOffset(..) => Some(stack_reg()),
             &AMode::FPOffset(..) => Some(fp_reg()),
             &AMode::NominalSPOffset(..) => Some(stack_reg()),
-            &AMode::Const(..) | AMode::Label(..) => None,
+            &AMode::Const(..) | AMode::Label(..) | AMode::Global(..) => None,
         }
     }
 
@@ -158,6 +163,7 @@ impl AMode {
             &AMode::FPOffset(offset, _) => offset,
             &AMode::NominalSPOffset(offset, _) => offset,
             &AMode::Const(_) | &AMode::Label(_) => 0,
+            &AMode::Global(offset, _) => offset,
         }
     }
 
@@ -187,6 +193,9 @@ impl Display for AMode {
             &AMode::Label(label) => {
                 write!(f, "[label{}]", label.as_u32())
             }
+            &AMode::Global(offset, ..) => {
+                write!(f, "{}(global)", offset)
+            }
         }
     }
 }

cranelift/codegen/src/isa/zkasm/inst/emit.rs

@@ -622,6 +622,12 @@ impl MachInstEmit for Inst {
                     }
                     AMode::Const(_) => unimplemented!("Load, AMode::Const"),
                     AMode::Label(_) => unimplemented!("Load, AMode::Label"),
+                    AMode::Global(index, _) => {
+                        put_string(
+                            &format!("$ => {} :MLOAD(global_{})\n", reg_name(rd.to_reg()), index),
+                            sink,
+                        );
+                    }
                 };
             }
             &Inst::Store { op, src, flags, to } => {
@@ -662,6 +668,12 @@ impl MachInstEmit for Inst {
                     }
                     AMode::Const(_) => unimplemented!("Store, AMode::Const"),
                     AMode::Label(_) => unimplemented!("Store, AMode::Label"),
+                    AMode::Global(index, _) => {
+                        put_string(
+                            &format!("{} :MSTORE(global_{})\n", reg_name(src), index),
+                            sink,
+                        );
+                    }
                 };
             }
             &Inst::Args { .. } => {

cranelift/codegen/src/isa/zkasm/lower.isle

@@ -355,8 +355,8 @@
 ;;;;;  Rules for `symbol_value`;;;;;;;;;
 (rule
    (lower (symbol_value (symbol_value_data name _ offset)))
-   (load_ext_name name offset)
-)
+   (load_ext_name name offset))
+
 ;;;;;  Rules for `bitcast`;;;;;;;;;
 (rule
    (lower (has_type out_ty (bitcast _ v @ (value_type in_ty))))

cranelift/codegen/src/isa/zkasm/lower/isle.rs

@@ -3,6 +3,7 @@
 // Pull in the ISLE generated code.
 #[allow(unused)]
 pub mod generated_code;
+use cranelift_entity::EntityRef;
 use generated_code::{Context, ExtendOp, MInst};
 
 // Types that the generated ISLE code uses via `use super::*`.
@@ -351,6 +352,10 @@ impl generated_code::Context for ZkAsmIsleContext<'_, '_, MInst, ZkAsmBackend> {
         AMode::Const(c)
     }
 
+    fn gen_global_amode(&mut self, offset: i64, ty: Type) -> AMode {
+        AMode::Global(offset, ty)
+    }
+
     fn valid_atomic_transaction(&mut self, ty: Type) -> Option<Type> {
         if ty.is_int() && ty.bits() <= 64 {
             Some(ty)
@@ -364,6 +369,28 @@ impl generated_code::Context for ZkAsmIsleContext<'_, '_, MInst, ZkAsmBackend> {
     fn store_op(&mut self, ty: Type) -> StoreOP {
         StoreOP::from_type(ty)
     }
+
+    fn zkasm_base(&mut self, global_value: GlobalValue) -> Option<generated_code::ZkasmBase> {
+        if let Some((name, _, offset)) = self.symbol_value_data(global_value) {
+            match name {
+                ExternalName::User(user_name_ref) => {
+                    // All ZKASM "memory"-like accesses use this name, but different offsets.
+                    if user_name_ref.index() != 0 {
+                        return None;
+                    }
+                    return match offset {
+                        0 => Some(generated_code::ZkasmBase::Heap),
+                        1 => Some(generated_code::ZkasmBase::Global),
+                        2 => Some(generated_code::ZkasmBase::Table),
+                        _ => None,
+                    };
+                }
+                _ => {}
+            }
+        }
+        None
+    }
+
     fn load_ext_name(&mut self, name: ExternalName, offset: i64) -> Reg {
         let tmp = self.temp_writable_reg(I64);
         self.emit(&MInst::LoadExtName {

cranelift/filetests/src/test_zkasm.rs

@@ -18,9 +18,41 @@ mod tests {
         let _ = env_logger::builder().is_test(true).try_init();
     }
 
-    fn generate_preamble(start_func_index: usize) -> Vec<String> {
+    fn generate_preamble(
+        start_func_index: usize,
+        globals: &[(cranelift_wasm::GlobalIndex, cranelift_wasm::GlobalInit)],
+    ) -> Vec<String> {
         let mut program: Vec<String> = Vec::new();
+
+        // Generate global variable definitions.
+        for (key, _) in globals {
+            program.push(format!("VAR GLOBAL global_{}", key.index()));
+        }
+
         program.push("start:".to_string());
+        for (key, init) in globals {
+            match init {
+                cranelift_wasm::GlobalInit::I32Const(v) => {
+                    // ZKASM stores constants in 2-complement form, so we need a cast to unsigned.
+                    program.push(format!(
+                        "  {} :MSTORE(global_{})  ;; Global32({})",
+                        *v as u32,
+                        key.index(),
+                        v
+                    ));
+                }
+                cranelift_wasm::GlobalInit::I64Const(v) => {
+                    // ZKASM stores constants in 2-complement form, so we need a cast to unsigned.
+                    program.push(format!(
+                        "  {} :MSTORE(global_{})  ;; Global64({})",
+                        *v as u64,
+                        key.index(),
+                        v
+                    ));
+                }
+                _ => unimplemented!("Global type is not supported"),
+            }
+        }
         program.push("  zkPC + 2 => RR".to_string());
         program.push(format!("  :JMP(function_{})", start_func_index));
         program.push("  :JMP(finalizeExecution)".to_string());
@@ -146,12 +178,16 @@ mod tests {
         translate_module(wasm_module, &mut zkasm_environ).unwrap();
 
         let mut program: Vec<String> = Vec::new();
+
         let start_func = zkasm_environ
             .info
             .start_func
             .expect("Must have a start function");
         // TODO: Preamble should be generated by a linker and/or clift itself.
-        program.append(&mut generate_preamble(start_func.index()));
+        program.append(&mut generate_preamble(
+            start_func.index(),
+            &zkasm_environ.info.global_inits,
+        ));
 
         let num_func_imports = zkasm_environ.get_num_func_imports();
         let mut context = cranelift_codegen::Context::new();
@@ -321,6 +357,7 @@ mod tests {
         add,
         locals,
         locals_simple,
+        global,
         memory,
         counter,
         add_func,

cranelift/wasm/src/environ/zkasm.rs

@@ -17,7 +17,7 @@ use crate::{
 use core::convert::TryFrom;
 use cranelift_codegen::cursor::FuncCursor;
 use cranelift_codegen::ir::immediates::{Imm64, Offset32, Uimm64};
-use cranelift_codegen::ir::{self, GlobalValue, InstBuilder};
+use cranelift_codegen::ir::{self, ExternalName, GlobalValue, InstBuilder};
 use cranelift_codegen::ir::{types::*, UserFuncName};
 use cranelift_codegen::isa::{CallConv, TargetFrontendConfig};
 use cranelift_entity::{EntityRef, PrimaryMap, SecondaryMap};
@@ -82,6 +82,9 @@ pub struct ZkasmModuleInfo {
     /// Globals as provided by `declare_global`.
     pub globals: PrimaryMap<GlobalIndex, Exportable<Global>>,
 
+    /// Inits for globals when available.
+    pub global_inits: Vec<(GlobalIndex, GlobalInit)>,
+
     /// The start function.
     pub start_func: Option<FuncIndex>,
 }
@@ -101,6 +104,7 @@ impl ZkasmModuleInfo {
             tables: PrimaryMap::new(),
             memories: PrimaryMap::new(),
             globals: PrimaryMap::new(),
+            global_inits: Vec::new(),
             start_func: None,
         }
     }
@@ -240,13 +244,20 @@ impl<'zkasm_environment> ZkasmFuncEnvironment<'zkasm_environment> {
         }
     }
 
+    // All ZKASM "memory"-like accesses use this name, but different offsets:
+    // - 0 for heap accesses
+    // - 1 for global variable accesses
+    // - 2 for table accesses
+    fn zkasm_base(func: &mut ir::Function) -> ExternalName {
+        ir::ExternalName::User(func.declare_imported_user_function(ir::UserExternalName {
+            namespace: 0,
+            index: 100,
+        }))
+    }
+
     // TODO(akashin): Ideally, we won't need a function here as the heap base is truly global.
     fn heap_base(func: &mut ir::Function) -> GlobalValue {
-        let name =
-            ir::ExternalName::User(func.declare_imported_user_function(ir::UserExternalName {
-                namespace: 0,
-                index: 100,
-            }));
+        let name = Self::zkasm_base(func);
         func.create_global_value(ir::GlobalValueData::Symbol {
             name,
             offset: Imm64::new(0),
@@ -257,29 +268,21 @@ impl<'zkasm_environment> ZkasmFuncEnvironment<'zkasm_environment> {
 
     // TODO(akashin): Ideally, we won't need a function here as the globals base is truly global.
     fn globals_base(func: &mut ir::Function) -> GlobalValue {
-        let name =
-            ir::ExternalName::User(func.declare_imported_user_function(ir::UserExternalName {
-                namespace: 0,
-                index: 101,
-            }));
+        let name = Self::zkasm_base(func);
         func.create_global_value(ir::GlobalValueData::Symbol {
             name,
-            offset: Imm64::new(65536),
+            offset: Imm64::new(1),
             colocated: false,
             tls: false,
         })
     }
 
     // TODO(akashin): Ideally, we won't need a function here as the heap base is truly global.
     fn table_base(func: &mut ir::Function) -> GlobalValue {
-        let name =
-            ir::ExternalName::User(func.declare_imported_user_function(ir::UserExternalName {
-                namespace: 0,
-                index: 102,
-            }));
+        let name = Self::zkasm_base(func);
         func.create_global_value(ir::GlobalValueData::Symbol {
             name,
-            offset: Imm64::new(2 * 65536),
+            offset: Imm64::new(2),
             colocated: false,
             tls: false,
         })
@@ -312,7 +315,7 @@ impl<'zkasm_environment> FuncEnvironment for ZkasmFuncEnvironment<'zkasm_environ
         func: &mut ir::Function,
         index: GlobalIndex,
     ) -> WasmResult<GlobalVariable> {
-        let offset = i32::try_from((index.index() * 8) + 8).unwrap().into();
+        let offset = i32::try_from(index.index()).unwrap().into();
         Ok(GlobalVariable::Memory {
             gv: ZkasmFuncEnvironment::globals_base(func),
             offset,
@@ -770,8 +773,9 @@ impl<'data> ModuleEnvironment<'data> for ZkasmEnvironment {
         Ok(())
     }
 
-    fn declare_global(&mut self, global: Global, _init: GlobalInit) -> WasmResult<()> {
-        self.info.globals.push(Exportable::new(global));
+    fn declare_global(&mut self, global: Global, init: GlobalInit) -> WasmResult<()> {
+        let index = self.info.globals.push(Exportable::new(global));
+        self.info.global_inits.push((index, init));
         Ok(())
     }
 

cranelift/zkasm_data/benchmarks/fibonacci/generated/from_rust.zkasm

@@ -1,4 +1,10 @@
+VAR GLOBAL global_0
+VAR GLOBAL global_1
+VAR GLOBAL global_2
 start:
+  1048576 :MSTORE(global_0)  ;; Global32(1048576)
+  1048576 :MSTORE(global_1)  ;; Global32(1048576)
+  1048576 :MSTORE(global_2)  ;; Global32(1048576)
   zkPC + 2 => RR
   :JMP(function_1)
   :JMP(finalizeExecution)

cranelift/zkasm_data/benchmarks/fibonacci/state.csv

@@ -1,4 +1,4 @@
 Test,Status,Cycles
-from_rust,pass,53023
+from_rust,pass,53026
 handwritten,pass,50008
 handwritten_wat,pass,170023