Proof of concept for how to do proper branching on target architecture that works for regular and Eurydice compilation runs

Makefile

@@ -50,7 +50,7 @@ clean-and-test:
 .PRECIOUS: %.llbc
 %.llbc: %.rs .charon_version
 	# --mir elaborated --add-drop-bounds 
-	$(CHARON) rustc --preset=eurydice --dest-file "$@" $(CHARON_EXTRA) -- $<
+	$(CHARON) rustc --preset=eurydice --dest-file "$@" $(CHARON_EXTRA) -- --cfg eurydice $<
 
 out/test-%/main.c: test/main.c
 	mkdir -p out/test-$*
@@ -69,6 +69,10 @@ test/println.llbc: CHARON_EXTRA = \
   --include=core::fmt::Arguments --include=core::fmt::rt::*::new_const \
   --include=core::fmt::rt::Argument
 
+test/intrinsics.llbc: CHARON_EXTRA = \
+  --include=core::fmt::Arguments --include=core::fmt::rt::*::new_const \
+  --include=core::fmt::rt::Argument
+
 test-partial_eq: EXTRA_C = ../../test/partial_eq_stubs.c
 test-nested_arrays: EXTRA = -funroll-loops 0
 test-array: EXTRA = -fcomments
@@ -77,6 +81,17 @@ test-more_str: EXTRA_C = ../../test/core_str_lib.c
 test-more_primitive_types: EXTRA = --config test/more_primitive_types.yaml
 test-global_ref: EXTRA_C = ../../test/core_cmp_lib.c
 
+test-intrinsics: EXTRA = --config test/intrinsics.yaml
+test-intrinsics: EXTRA_C = -I../../test
+
+ifeq ($(shell uname -m),arm64)
+test-intrinsics: EXTRA_C = vectorized.c vectorized_arm.c
+endif
+
+ifeq ($(shell uname -m),x86_64)
+test-intrinsics: EXTRA_C = vectorized.c vectorized_intel.c
+endif
+
 
 test-%: test/%.llbc out/test-%/main.c | all
 	$(EURYDICE) $(EXTRA) --output out/test-$* $<
@@ -132,3 +147,6 @@ format-apply:
 .PHONY: clean-llbc
 clean-llbc:
 	rm test/*.llbc || true
+
+%.pp.rs: %.rs
+	rustup run $(shell basename $$($(CHARON) toolchain-path)) rustc -Zunpretty=expanded --cfg eurydice $<

lib/Builtin.ml

@@ -522,6 +522,16 @@ let static_assert, static_assert_ref =
   ( K.DExternal (None, [ Krml.Common.Private; Macro ], 0, 0, name, typ, [ "test"; "msg" ]),
     K.(with_type typ (EQualified name)) )
 
+let targets =
+  let mk name =
+    K.DExternal (None, [ Krml.Common.Private; IfDef ], 0, 0, (["Eurydice"], "target_is_" ^ name), K.TBool, [])
+  in
+  [
+    mk "aarch64";
+    mk "x86";
+    mk "x86_64";
+  ]
+
 (* Replacements, now applied on-the-fly in AstOfLlbc.
 
  IMPORTANT: such replacements are written in abstract syntax that *already* has cleanups applied,
@@ -777,7 +787,7 @@ let files =
            in
            K.DExternal (None, flags, List.length cg_args, n_type_args, name, typ, arg_names))
          builtin_funcs
-       @ [ nonzero_def; static_assert; dst_def; str_t_def ]
+       @ [ nonzero_def; static_assert; dst_def; str_t_def ] @ targets
      in
      "Eurydice", externals);
   ]

lib/Cleanup1.ml

@@ -465,6 +465,23 @@ let remove_slice_eq = object
     | _ -> super#visit_expr ((), e.typ) e
 end
 
+let match_target_ifdefs = object
+  inherit [_] map as super
+  method! visit_EApp (((), _) as env) e es =
+    match e.node, es with
+    | EQualified ([_; "eurylib"], "target_arch"), [ arg ] ->
+        begin match arg with
+        | { node = EFlat [Some "data", {node=EString str; _}; _ ]; _ } ->
+            begin match str with
+            | "x86" | "x86_64" | "aarch64" -> EQualified (["Eurydice"], "target_is_" ^ str)
+            | _ -> failwith ("Unknown target arch: " ^ str)
+            end
+        | _ -> failwith ("Argument to `target_arch` is not a constant string")
+        end
+    | _ ->
+        super#visit_EApp env e es
+end
+
 let cleanup files =
   let files = remove_units#visit_files () files in
   let files = remove_assignments#visit_files AtomMap.empty files in
@@ -474,4 +491,5 @@ let cleanup files =
   let files = remove_terminal_continues#visit_files false files in
   let files = Krml.Simplify.let_to_sequence#visit_files () files in
   let files = remove_slice_eq#visit_files () files in
+  let files = match_target_ifdefs#visit_files () files in
   files

test/eurydice_intrinsics.h

@@ -0,0 +1,8 @@
+#if defined(__aarch64__)
+#include <arm_neon.h>
+typedef uint16x8_t Vec128;
+#elif defined(__i686__)
+typedef __m128i Vec128;
+#else
+#error "Incorrect header inclusion"
+#endif

test/intrinsics.rs

@@ -0,0 +1,215 @@
+#![feature(register_tool)]
+#![register_tool(charon)]
+
+// STEP 1: Implementation multiplexing. We provide two macros, one for runtime, and one for
+// extracting, which generates a semantically equivalent cascade of if-then-elses, knowing that the
+// call to `target_arch` will receive special treatment and be turned into an ifdef.
+//
+// Remember that the easiest way to see the effect of macro expansion is to call
+// `rustc --cfg eurydice -Zunpretty=expanded src/main.rs`
+
+#[cfg(not(eurydice))]
+#[macro_export]
+macro_rules! arch_match {
+    ( $( $(|)? $($x:literal)|* => $e:block )|* | _ => $e_default:block ) => {
+        {
+            $(
+                $(
+                    #[cfg(target_arch=$x)]
+                    if (runtime_support($x)) {
+                        $e
+                    } else {
+                        $e_default
+                    }
+                )*
+            )*
+            #[cfg(all($($(not(target_arch=$x)),*),*))]
+            $e_default
+        }
+    };
+}
+
+/*
+// First attempt: closer to what krml expects to generate the correct #ifdef-based code, but no way
+// to inline `target_arch` prior to running Charon (and reduce the match).
+mod eurylib {
+    #[charon::opaque]
+    pub const TARGET_IS_AARCH64: bool = false;
+    #[charon::opaque]
+    pub const TARGET_IS_X86_64: bool = false;
+}
+
+fn target_arch(x: &str) -> bool {
+    match x {
+      "aarch64" => eurylib::TARGET_IS_AARCH64,
+      "x86_64" => eurylib::TARGET_IS_X86_64,
+      _ => panic!("oh noes")
+    }
+}
+*/
+
+// Second attempt: with dedicated support in Eurydice
+mod eurylib {
+    #[charon::opaque]
+    pub fn target_arch(x: &str) -> bool {
+        true
+    }
+}
+
+use eurylib::*;
+
+fn runtime_support(x: &str) -> bool {
+    // CPUID, etc.
+    true
+}
+
+#[cfg(eurydice)]
+#[macro_export]
+macro_rules! arch_match {
+    ( $( $(|)? $($x:literal)|* => $e:block )|* | _ => $e_default:block ) => {
+        {
+            $(
+                $(
+                    if (target_arch($x) && runtime_support($x)) {
+                        $e
+                    } else 
+                )*
+            )*
+            $e_default
+        }
+    };
+}
+
+// STEP 2: stub out intrinsics so that the Vec128 and both Intel and ARM intrinsics can live in the
+// same scope, for the purposes of extraction.
+
+#[cfg(not(eurydice))]
+mod intrinsics {
+    // Intel
+    #[cfg(target_arch = "x86")]
+    pub use std::arch::x86::*;
+    #[cfg(target_arch = "x86_64")]
+    pub use std::arch::x86_64::*;
+    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+    pub type Vec128 = __m128i;
+
+    // ARM
+    #[cfg(target_arch = "aarch64")]
+    pub use core::arch::aarch64::*;
+    #[cfg(target_arch = "aarch64")]
+    pub type Vec128 = uint16x8_t;
+}
+
+#[cfg(eurydice)]
+mod intrinsics {
+    mod private {
+        #[charon::opaque]
+        pub struct Vec128 { dummy: () }
+    }
+    // FIXME: why go through that indirection?
+    pub use self::private::Vec128;
+    unsafe extern "C" {
+        pub fn vdupq_n_u16(x: u16) -> Vec128;
+        pub fn vst1q_u16(x: *mut u16, y: Vec128);
+
+        pub fn _mm_set1_epi16(x: i16) -> Vec128;
+        pub fn _mm_storeu_si128(x: *mut Vec128, y: Vec128);
+    }
+}
+
+// STEP 3: client implements their library of optimized implementations for both targets
+
+#[cfg(any(target_arch = "x86", target_arch = "x86_64", eurydice))]
+mod vectorized_intel {
+    use super::intrinsics::*;
+
+    pub fn vec128_set_u16(val: u16) -> Vec128 {
+        unsafe { _mm_set1_epi16(val as i16) }
+    }
+
+    pub fn vec128_store(elem: &mut [u16; 8], val: Vec128) {
+        unsafe {
+            let addr = elem.as_mut_ptr();
+            _mm_storeu_si128(addr as *mut Vec128, val);
+        }
+    }
+}
+
+#[cfg(any(target_arch = "aarch64", eurydice))]
+mod vectorized_arm {
+    use super::intrinsics::*;
+
+    pub fn vec128_set_u16(val: u16) -> Vec128 {
+        unsafe { vdupq_n_u16(val) }
+    }
+
+    pub fn vec128_store(elem: &mut [u16; 8], val: Vec128) {
+        unsafe {
+            let addr = elem.as_mut_ptr();
+            vst1q_u16(addr, val);
+        }
+    }
+
+}
+
+// TODO: autogenerate
+#[cfg(eurydice)]
+mod vectorized_stubs {
+    use super::*;
+    use super::intrinsics::Vec128;
+
+    pub fn vec128_set_u16(val: u16) -> Vec128 {
+        arch_match!{
+            | "x86" | "x86_64" => { vectorized_intel::vec128_set_u16(val) }
+            | "aarch64" => { vectorized_arm::vec128_set_u16(val) }
+            | _ => { panic!("no default impl") }
+        }
+    }
+
+    pub fn vec128_store(elem: &mut [u16; 8], val: Vec128) {
+        arch_match!{
+            | "x86" | "x86_64" => { vectorized_intel::vec128_store(elem, val) }
+            | "aarch64" => { vectorized_arm::vec128_store(elem, val) }
+            | _ => { panic!("no default impl") }
+        }
+    }
+}
+
+// STEP 4: regular build picks either one of these, C extraction multiplexes
+
+#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), not(eurydice)))]
+use vectorized_intel::*;
+
+#[cfg(all(target_arch = "aarch64", not(eurydice)))]
+use vectorized_arm::*;
+
+#[cfg(eurydice)]
+use vectorized_stubs::*;
+
+// STEP 5: client multiplexes via the macro
+
+fn default_impl() { println!("ok") }
+
+mod vectorized {
+    use super::*;
+    pub fn vector_impl() {
+        let expected = [1u16; 8];
+        let v = vec128_set_u16(1);
+        let mut actual = [0u16; 8];
+        vec128_store(&mut actual, v);
+        assert_eq!(actual, expected);
+    }
+}
+
+use vectorized::*;
+
+// TODO: intrinsics.c unconditionally includes vectorized.h -- this should only happen when
+// vectorized implementations are available
+fn main() {
+    arch_match!{
+        | "x86" | "x86_64" => { vector_impl() }
+        | "aarch64" => { vector_impl() }
+        | _ => { default_impl() }
+    }
+    println!("Hello, world!");
+}

test/intrinsics.yaml

@@ -0,0 +1,50 @@
+files:
+  # This one does not generate an extracted file, intentionally.
+  - name: stubs
+    private:
+      - [ intrinsics, intrinsics, "*" ]
+    inline_static: true
+    library: true
+
+  - name: core
+    private:
+      monomorphizations_of:
+        - [ core, "*"]
+      monomorphizations_using:
+        - [ Eurydice, "*" ]
+      patterns:
+        - [ core, "*"]
+      exact:
+        - [ intrinsics, runtime_support]
+        - [ intrinsics, eurylib, target_arch]
+    api:
+      - [Eurydice, "*"]
+
+  - name: vectorized_intel
+    include_in_h:
+      - "eurydice_intrinsics.h"
+    private:
+      - [ intrinsics, vectorized_intel, "*" ]
+
+  - name: vectorized_arm
+    include_in_internal_h:
+      - "\"eurydice_intrinsics.h\""
+    private:
+      - [ intrinsics, vectorized_arm, "*" ]
+
+  - name: vectorized
+    include_in_internal_h:
+      - "\"eurydice_intrinsics.h\""
+    private:
+      - [ intrinsics, vectorized, "*" ]
+      - [ intrinsics, vectorized_stubs, "*" ]
+
+  - name: intrinsics
+    api:
+      - [ "*" ]
+
+
+naming:
+  skip_prefix:
+    - [ intrinsics, intrinsics ]
+    - [ intrinsics, intrinsics, private ]