Something went wrong. Try again.
forked from raven-ml/raven
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159(*--------------------------------------------------------------------------- Copyright (c) 2026 The Raven authors. All rights reserved. SPDX-License-Identifier: ISC ---------------------------------------------------------------------------*)
(* A description on each GPU driver the machine has that runs launches: it leaves the bytes the same submission leaves when made by hand. Kernels come from the driver's conformance fixtures (Rig_gpu_support.Conformance.launch_binary). *)
open Windtrapmodule B = Rig.Buffermodule Sub = Rig.Submissionmodule G = Rig_program
module type Gpu = Rig_gpu_support.Conformance
let words = 256
let le64 v = let b = Bytes.create 8 in Bytes.set_int64_le b 0 (Int64.of_int v); Bytes.to_string b
let le32 v = let b = Bytes.create 4 in Bytes.set_int32_le b 0 (Int32.of_int v); Bytes.to_string b
let get b = let n = B.length b in let h = B.create Rig.host n in B.copy ~src:b ~dst:h; let s = Bytes.create n in B.blit_to_bytes h 0 s 0 n; Bytes.to_string s
(* [ids] into slot 0 over 4 groups of 64 threads, [a = 5], [b = 3], [f = 2.5]: word [k] is [5 + 3k + 2]. Then [twice] from slot 0 into slot 1, [c = 1], its parameters padded to 24 bytes, the size of Metal's struct of them. *)let ids_params = le64 0 ^ le64 5 ^ le32 3 ^ le32 (Int32.to_int (Int32.bits_of_float 2.5))
let twice_params = le64 0 ^ le64 0 ^ le32 1 ^ le32 0
let by_hand d image x y = let launch kernel params refs = let fn = require_some (Rig.Image.resolve image kernel) in Sub.launch ~refs fn ~params:(String.length params) in let s = Sub.make d [| launch "ids" ids_params [| Sub.write 0 ~at:0 |]; launch "twice" twice_params [| Sub.write 1 ~at:0; Sub.read 0 ~at:8 |]; |] in let run = Sub.Run.make s in let store i params ~groups ~threads = let b = Sub.block s i in Sub.Run.groups run b groups 1 1; Sub.Run.threads run b threads 1 1; Sub.Run.shared run b 0; for w = 0 to (String.length params / 4) - 1 do Sub.Run.int32 run b (4 * w) (Int32.to_int (String.get_int32_le params (4 * w))) done in store 0 ids_params ~groups:4 ~threads:64; store 1 twice_params ~groups:words ~threads:1; ignore (Rig.submit run ~buffers:[| x; y |])
let described d binary = let launch name params refs ~groups ~threads = (G.Launch { image = 0; name; params = { G.bytes = params; holes = [||] }; refs; groups = (G.Fixed groups, G.Fixed 1, G.Fixed 1); threads = (G.Fixed threads, G.Fixed 1, G.Fixed 1); shared = G.Fixed 0; }) in { G.devices = [| Rig.arch d |]; ints = 0; memory = [||]; images = [| { G.device = 0; binary = { G.bytes = binary; holes = [||] } } |]; inputs = Iarray.init 2 (fun _ -> { G.device = 0; bytes = 4 * words }); steps = [| G.Submit { device = 0; slots = [| G.Input 0; G.Input 1 |]; fixed = [||]; parts = [| launch "ids" ids_params [| { G.slot = 0; at = 0; use = Write } |] ~groups:4 ~threads:64; launch "twice" twice_params [| { G.slot = 1; at = 0; use = Write }; { G.slot = 0; at = 8; use = Read }; |] ~groups:words ~threads:1; |]; }; |]; }
let same_bytes (module Gpu : Gpu) () = Gpu.with_ @@ fun t -> let d = t.d in let binary = match Gpu.launch_binary () with | Some b -> b | None -> skip ~reason:"the driver runs no launch" () in let image = Result.get_ok (Rig.Image.load d binary) in let x = B.create d (4 * words) and y = B.create d (4 * words) in by_hand d image x y; let p = require_ok ~pp:Format.pp_print_string (G.load (described d binary) [| d |]) in let x' = B.create d (4 * words) and y' = B.create d (4 * words) in ignore (G.run p { inputs = [| x'; y' |]; ints = [||]; step_points = G.step_points p }); let le32s f = String.concat "" (List.init words (fun k -> le32 (f k))) in equal string ~msg:"ids, by hand" (le32s (fun k -> 7 + (3 * k))) (get x); equal string ~msg:"twice, by hand" (le32s (fun k -> 15 + (6 * k))) (get y); equal string ~msg:"ids, described" (get x) (get x'); equal string ~msg:"twice, described" (get y) (get y')
let gpus : (module Gpu) list = [ (module Rig_metal_support); (module Rig_cuda_support); (module Rig_nv_support); (module Rig_amd_support); ]
let () = exit (run "rig.program.gpu" (List.map (fun (module Gpu : Gpu) -> test ~timeout:120. (Gpu.class_ ^ ": a description leaves the bytes of its submission made by hand" ) (same_bytes (module Gpu))) gpus))