Embed non-linked builtins in OCaml strings for easier availability at runtime

lukstafi · lukstafi · commit 59194177aeac · 2025-08-01T22:24:24.000+02:00
Note: Metal backend still broken
diff --git a/arrayjit/lib/builtins_cuda_large.ml b/arrayjit/lib/builtins_cuda_large.ml
@@ -1,3 +1,4 @@
+let source = {|
 #include <cuda_runtime.h>
 #include <stdint.h>
 
@@ -107,4 +108,6 @@ __device__ uint4x32_t arrayjit_threefry4x32_impl(uint4x32_t key, uint4x32_t coun
     return result;
 }
 
-__device__ uint4x32_t ( *arrayjit_threefry4x32)(uint4x32_t key, uint4x32_t counter) = arrayjit_threefry4x32_impl;
+__device__ uint4x32_t ( *arrayjit_threefry4x32)(uint4x32_t key, uint4x32_t counter) = arrayjit_threefry4x32_impl;
+
+|}
diff --git a/arrayjit/lib/builtins_cuda_small.ml b/arrayjit/lib/builtins_cuda_small.ml
@@ -1,4 +1,4 @@
-
+let source = {|
 typedef struct {
     uint32_t v[4];
 } uint4x32_t;
@@ -181,4 +181,5 @@ __device__ uint8x16_t uint4x32_to_u8_uniform_vec(uint4x32_t x) {
     result.v[i*4 + 3] = (uint8_t)((x.v[i] >> 24) & 0xFF);
   }
   return result;
-}
+}
+|}
diff --git a/arrayjit/lib/builtins_metal.ml b/arrayjit/lib/builtins_metal.ml
@@ -1,3 +1,4 @@
+let source = {|
 #include <metal_stdlib>
 using namespace metal;
 
@@ -367,4 +368,5 @@ uint4 half_to_uint4x32(uint16_t x) {
 
 uint4 fp8_to_uint4x32(uint8_t x) {
     return uint4(uint32_t(x), 0, 0, 0);
-}
+}
+|}
diff --git a/arrayjit/lib/cuda_backend.ml b/arrayjit/lib/cuda_backend.ml
@@ -709,11 +709,7 @@ end) : Ir.Backend_impl.Lowered_backend = struct
     if Utils.debug_log_from_routines () then
       Buffer.add_string b "__device__ int printf (const char * format, ... );\n";
     Buffer.add_string b "\n\n";
-    let builtins_path =
-      Stdlib.Filename.concat (Stdlib.Filename.dirname Stdlib.__FILE__) "builtins_small.cu"
-    in
-    let builtins_content = Stdio.In_channel.read_all builtins_path in
-    Buffer.add_string b builtins_content;
+    Buffer.add_string b Builtins_cuda_small.source;
     (* Needs to be after the small builtins, because uses uint4x32_t. *)
     Buffer.add_string b builtins_large_header;
     Buffer.add_string b "\n\n"
diff --git a/arrayjit/lib/dune b/arrayjit/lib/dune
@@ -62,7 +62,7 @@
  (name cuda_backend)
  (public_name arrayjit.cuda_backend)
  (optional)
- (modules cuda_backend)
+ (modules cuda_backend builtins_cuda_large builtins_cuda_small)
  (libraries base cudajit.cuda cudajit.nvrtc utils ir)
  (preprocess
   (pps
@@ -78,7 +78,7 @@
  (name metal_backend)
  (public_name arrayjit.metal_backend)
  (optional)
- (modules metal_backend)
+ (modules metal_backend builtins_metal)
  (libraries base metal utils ir)
  (preprocess
   (pps
diff --git a/arrayjit/lib/metal_backend.ml b/arrayjit/lib/metal_backend.ml
@@ -639,21 +639,14 @@ end) : Ir.Backend_impl.Lowered_backend = struct
       Stdio.prerr_endline error_msg;
       failwith error_msg
 
-  let prepend_builtins b =
-    let builtins_path =
-      Stdlib.Filename.concat (Stdlib.Filename.dirname Stdlib.__FILE__) "builtins.metal"
-    in
-    let builtins_content = Stdio.In_channel.read_all builtins_path in
-    Buffer.add_string b builtins_content;
-    Buffer.add_string b "\n\n"
-
   let compile ~name bindings lowered =
     let module Syntax = C_syntax.C_syntax (C_syntax_config (struct
       let procs = [| lowered |]
     end)) in
     let idx_params = Indexing.bound_symbols bindings in
     let b = Buffer.create 4096 in
-    prepend_builtins b;
+    Buffer.add_string b Builtins_metal.source;
+    Buffer.add_string b "\n";
     let declarations_doc = Syntax.print_declarations () in
     (* Add Metal address space qualifiers *)
     let params, proc_doc = Syntax.compile_proc ~name idx_params lowered in
