kit

kit
git clone https://git.ryansepassi.com/git/kit.git
Log | Files | Refs | README

rv64_jit_test.c (9908B)

      1 /* RV64 JIT smoke test.
      2  *
      3  * Builds a tiny ELF relocatable object in memory for rv64 containing
      4  * one function:
      5  *
      6  *   .text
      7  *   .globl rv64_jit_answer
      8  *   rv64_jit_answer:
      9  *     addi a0, zero, 42      # 0x02a00513
     10  *     jalr zero, ra, 0       # 0x00008067  (ret)
     11  *
     12  * Feeds it through kit_link_session in KIT_LINK_OUTPUT_JIT mode,
     13  * which exercises the rv64 path of:
     14  *   - executable-memory reservation + W^X protect cycle
     15  *   - relocation application (none needed here, but the path runs)
     16  *   - symbol resolution / lookup by C-mangled name
     17  *   - icache flush (fence.i / __riscv_flush_icache on rv64 hosts)
     18  *
     19  * If we are running on a rv64 host, the test then *calls* the JITed
     20  * function and asserts the return is 42 — that's the native-host
     21  * execution leg the parity checklist asked for.  On non-rv64 hosts
     22  * we still build the image (verifying the in-memory machinery is wired
     23  * end-to-end) but SKIP the actual call: the bytes are valid rv64 but
     24  * the host CPU can't decode them.  The test prints "SKIP <reason>" and
     25  * exits 77 (the GNU autotools "skipped" convention) when this happens.
     26  *
     27  * Wired into mk/test.mk via test-rv64-jit.  Always builds; calls only on
     28  * rv64 Linux.  This mirrors the value-proposition outlined in
     29  * doc/RV64_PARITY_CHECKLIST.md: have the code path in place for the day
     30  * someone runs kit on a rv64 dev box. */
     31 
     32 #include <kit/core.h>
     33 #include <kit/jit.h>
     34 #include <kit/link.h>
     35 #include <kit/object.h>
     36 #include <stdarg.h>
     37 #include <stdint.h>
     38 #include <stdio.h>
     39 #include <stdlib.h>
     40 #include <string.h>
     41 #include <sys/mman.h>
     42 #include <unistd.h>
     43 
     44 #include "lib/kit_unit.h"
     45 
     46 /* Native execution requires the host CPU to be rv64 (any OS that gives
     47  * us POSIX mmap + mprotect, which on rv64 means Linux today).  Anywhere
     48  * else we still build the JIT image but skip the call. */
     49 #if defined(__riscv) && (__riscv_xlen == 64)
     50 #define RV64_HOST_NATIVE 1
     51 #else
     52 #define RV64_HOST_NATIVE 0
     53 #endif
     54 
     55 /* ---- host glue: heap + diag come from the shared KitUnit ---- */
     56 static KitUnit g_u;
     57 
     58 /* ---- execmem with W^X dual-mapping (mirrors test/link/harness) ---- */
     59 static int xm_to_posix(int p) {
     60   int q = 0;
     61   if (p & KIT_PROT_READ) q |= PROT_READ;
     62   if (p & KIT_PROT_WRITE) q |= PROT_WRITE;
     63   if (p & KIT_PROT_EXEC) q |= PROT_EXEC;
     64   return q;
     65 }
     66 
     67 #if defined(__linux__)
     68 #include <sys/syscall.h>
     69 #define XM_DUAL_LINUX 1
     70 #else
     71 #define XM_DUAL_LINUX 0
     72 #endif
     73 
     74 typedef struct XmTok {
     75   void* w;
     76   void* r;
     77   size_t n;
     78 } XmTok;
     79 
     80 static KitStatus xm_reserve_single(size_t n, KitExecMemRegion* out) {
     81   void* p =
     82       mmap(NULL, n, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0);
     83   if (p == MAP_FAILED) return KIT_NOMEM;
     84   out->write = out->runtime = p;
     85   out->size = n;
     86   out->token = NULL;
     87   return KIT_OK;
     88 }
     89 
     90 static KitStatus xm_reserve(void* u, size_t n, int p, KitExecMemRegion* out) {
     91   (void)u;
     92   if (!out || !n) return KIT_INVALID;
     93   if (!(p & KIT_PROT_EXEC)) return xm_reserve_single(n, out);
     94 #if XM_DUAL_LINUX
     95   {
     96     int fd = (int)syscall(SYS_memfd_create, "kit-rv64-jit-test", 0u);
     97     void *w, *r;
     98     XmTok* tok;
     99     if (fd < 0) return KIT_NOMEM;
    100     if (ftruncate(fd, (off_t)n) != 0) {
    101       close(fd);
    102       return KIT_NOMEM;
    103     }
    104     w = mmap(NULL, n, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
    105     if (w == MAP_FAILED) {
    106       close(fd);
    107       return KIT_NOMEM;
    108     }
    109     r = mmap(NULL, n, PROT_READ, MAP_SHARED, fd, 0);
    110     close(fd);
    111     if (r == MAP_FAILED) {
    112       munmap(w, n);
    113       return KIT_NOMEM;
    114     }
    115     tok = (XmTok*)malloc(sizeof(*tok));
    116     if (!tok) {
    117       munmap(r, n);
    118       munmap(w, n);
    119       return KIT_NOMEM;
    120     }
    121     tok->w = w;
    122     tok->r = r;
    123     tok->n = n;
    124     out->write = w;
    125     out->runtime = r;
    126     out->size = n;
    127     out->token = tok;
    128     return KIT_OK;
    129   }
    130 #else
    131   return xm_reserve_single(n, out);
    132 #endif
    133 }
    134 
    135 static KitStatus xm_protect(void* u, void* a, size_t n, int p) {
    136   (void)u;
    137   return mprotect(a, n, xm_to_posix(p)) == 0 ? KIT_OK : KIT_IO;
    138 }
    139 
    140 static void xm_release(void* u, KitExecMemRegion* region) {
    141   (void)u;
    142   if (!region || !region->size) return;
    143   if (region->token) {
    144     XmTok* tok = (XmTok*)region->token;
    145     if (tok->r && tok->r != tok->w) munmap(tok->r, tok->n);
    146     if (tok->w) munmap(tok->w, tok->n);
    147     free(tok);
    148   } else if (region->write) {
    149     munmap(region->write, region->size);
    150   }
    151   region->write = region->runtime = NULL;
    152   region->size = 0;
    153   region->token = NULL;
    154 }
    155 
    156 static void xm_flush(void* u, void* a, size_t n) {
    157   (void)u;
    158 #if defined(__aarch64__) || defined(__arm__) || defined(__riscv)
    159 #if defined(__riscv)
    160   /* Local-hart self-modify ordering; __builtin___clear_cache below also
    161    * issues the cross-hart syscall on Linux. */
    162   __asm__ __volatile__("fence.i" ::: "memory");
    163 #endif
    164   __builtin___clear_cache((char*)a, (char*)a + n);
    165 #else
    166   (void)a;
    167   (void)n;
    168 #endif
    169 }
    170 
    171 static KitExecMem g_execmem = {
    172     16 * 1024, xm_reserve, xm_protect, xm_release, xm_flush, NULL,
    173 };
    174 
    175 /* ---- rv64 instruction encodings used by the test ---- */
    176 /* `addi a0, zero, 42` — I-type: imm[11:0]=42, rs1=0, funct3=000 (ADDI),
    177  *  rd=10 (a0), opcode=0010011. */
    178 #define ENC_ADDI_A0_ZERO_42 0x02a00513u
    179 /* `jalr zero, 0(ra)` (= ret) — I-type: imm=0, rs1=1 (ra), funct3=000,
    180  *  rd=0 (zero), opcode=1100111. */
    181 #define ENC_RET 0x00008067u
    182 
    183 /* ---- the test ---- */
    184 typedef int (*answer_fn)(void);
    185 
    186 int main(void) {
    187   /* Page size for the execmem.  Same dance as the other runners. */
    188   {
    189     long ps = sysconf(_SC_PAGESIZE);
    190     if (ps > 0) g_execmem.page_size = (size_t)ps;
    191   }
    192 
    193   KitTargetSpec target =
    194       kit_unit_target(KIT_ARCH_RV64, KIT_OS_LINUX, KIT_OBJ_ELF);
    195 
    196   kit_unit_init(&g_u);
    197   g_u.ctx.now = -1;
    198 
    199   KitCompiler* c = NULL;
    200   if (kit_unit_compiler_new(&g_u, target, &c) != KIT_OK || !c) {
    201     fprintf(stderr, "rv64_jit_test: compiler_new failed\n");
    202     return 2;
    203   }
    204 
    205   /* Build the object. */
    206   KitObjBuilder* ob = NULL;
    207   if (kit_obj_builder_new(c, &ob) != KIT_OK || !ob) {
    208     fprintf(stderr, "rv64_jit_test: obj_builder_new failed\n");
    209     kit_compiler_free(c);
    210     return 2;
    211   }
    212 
    213   KitObjSectionDesc sec_desc;
    214   memset(&sec_desc, 0, sizeof(sec_desc));
    215   sec_desc.name = kit_sym_intern(c, KIT_SLICE_LIT(".text"));
    216   sec_desc.kind = KIT_SEC_TEXT;
    217   sec_desc.flags = KIT_SF_EXEC | KIT_SF_ALLOC;
    218   sec_desc.align = 4;
    219   KitObjSection text = KIT_SECTION_NONE;
    220   if (kit_obj_builder_section(ob, &sec_desc, &text) != KIT_OK) {
    221     fprintf(stderr, "rv64_jit_test: section failed\n");
    222     return 2;
    223   }
    224 
    225   uint32_t code[2] = {ENC_ADDI_A0_ZERO_42, ENC_RET};
    226   if (kit_obj_builder_write(ob, text, code, sizeof(code)) != KIT_OK) {
    227     fprintf(stderr, "rv64_jit_test: write failed\n");
    228     return 2;
    229   }
    230 
    231   KitObjSymbolDesc sym_desc;
    232   memset(&sym_desc, 0, sizeof(sym_desc));
    233   sym_desc.name = kit_sym_intern(c, KIT_SLICE_LIT("rv64_jit_answer"));
    234   sym_desc.bind = KIT_SB_GLOBAL;
    235   sym_desc.kind = KIT_SK_FUNC;
    236   sym_desc.section = text;
    237   sym_desc.value = 0;
    238   sym_desc.size = sizeof(code);
    239   KitObjSymbol sym = KIT_OBJ_SYMBOL_NONE;
    240   if (kit_obj_builder_symbol(ob, &sym_desc, &sym) != KIT_OK) {
    241     fprintf(stderr, "rv64_jit_test: symbol failed\n");
    242     return 2;
    243   }
    244 
    245   if (kit_obj_builder_finalize(ob) != KIT_OK) {
    246     fprintf(stderr, "rv64_jit_test: finalize failed\n");
    247     return 2;
    248   }
    249 
    250   /* JIT the object.  The host's execmem is the W^X dual-map above. */
    251   KitJitHost jhost;
    252   memset(&jhost, 0, sizeof(jhost));
    253   jhost.execmem = &g_execmem;
    254 
    255   KitLinkSessionOptions opts;
    256   memset(&opts, 0, sizeof(opts));
    257   opts.output_kind = KIT_LINK_OUTPUT_JIT;
    258   opts.entry = KIT_SLICE_LIT("rv64_jit_answer");
    259   opts.jit_host = &jhost;
    260 
    261   KitLinkSession* sess = NULL;
    262   if (kit_link_session_new(c, &opts, &sess) != KIT_OK || !sess) {
    263     fprintf(stderr, "rv64_jit_test: link_session_new failed\n");
    264     return 1;
    265   }
    266   if (kit_link_session_add_obj(sess, ob) != KIT_OK) {
    267     fprintf(stderr, "rv64_jit_test: add_obj failed\n");
    268     kit_link_session_free(sess);
    269     return 1;
    270   }
    271 
    272   KitJit* jit = NULL;
    273   if (kit_link_session_jit(sess, &jit) != KIT_OK || !jit) {
    274     fprintf(stderr, "rv64_jit_test: link_session_jit failed\n");
    275     kit_link_session_free(sess);
    276     return 1;
    277   }
    278   kit_link_session_free(sess);
    279 
    280   void* fn = kit_jit_lookup(jit, KIT_SLICE_LIT("rv64_jit_answer"));
    281   if (!fn) {
    282     fprintf(stderr, "rv64_jit_test: lookup failed\n");
    283     kit_jit_free(jit);
    284     kit_compiler_free(c);
    285     return 1;
    286   }
    287 
    288   /* Reading back the first instruction bytes through the runtime alias
    289    * is always safe and verifies the bytes survived the W^X dance plus
    290    * the icache-flush hook fired without crashing.  This is the portable
    291    * check on non-rv64 hosts. */
    292   uint32_t got = 0;
    293   memcpy(&got, fn, sizeof(got));
    294   if (got != ENC_ADDI_A0_ZERO_42) {
    295     fprintf(stderr,
    296             "rv64_jit_test: bytes corrupted at runtime alias: got 0x%08x "
    297             "expected 0x%08x\n",
    298             (unsigned)got, (unsigned)ENC_ADDI_A0_ZERO_42);
    299     kit_jit_free(jit);
    300     kit_compiler_free(c);
    301     return 1;
    302   }
    303 
    304 #if RV64_HOST_NATIVE
    305   /* Real execution on a rv64 host. */
    306   {
    307     answer_fn f = (answer_fn)(uintptr_t)fn;
    308     int r = f();
    309     if (r != 42) {
    310       fprintf(stderr, "rv64_jit_test: jit fn returned %d, expected 42\n", r);
    311       kit_jit_free(jit);
    312       kit_compiler_free(c);
    313       return 1;
    314     }
    315     printf("rv64_jit_test: PASS (native rv64 execution returned 42)\n");
    316   }
    317 #else
    318   /* Non-rv64 host: JIT plumbing worked end-to-end (image built,
    319    * permissions flipped, lookup resolved, bytes intact at the runtime
    320    * alias).  Skip the actual call — calling rv64 bytes on a non-rv64
    321    * CPU would SIGILL.  Exit-code 77 is the GNU autotools convention
    322    * for "skipped" so test wrappers can distinguish from pass/fail. */
    323   printf(
    324       "rv64_jit_test: SKIP — non-rv64 host (image built, "
    325       "lookup OK, bytes intact)\n");
    326   kit_jit_free(jit);
    327   kit_compiler_free(c);
    328   return 77;
    329 #endif
    330 
    331   kit_jit_free(jit);
    332   kit_compiler_free(c);
    333   return 0;
    334 }