Skip to content

Commit 85eaff9

Browse files
lukstafilukstafi
committed
First pass on adding vector-returning operations, by Claude Opus
Prompt: Separate Uint4x32_to_prec_uniform from the type unop in arrayjit/lib/ops.ml into a separate type vec_unop . In the Low_level module type t , add a constructor Set_from_vec { tn : Tnode.t; idcs : Indexing.axis_index array; length: int; vec_unop : Ops.vec_unop; arg : scalar_t; mutable debug : string } . This constructor is an analogue of the Set constructor but for vector rather than scalar values. The length field is the length of the vector, and backends should set the consecutive cells of the buffer corresponding to tn , starting with position idcs . At present, given that we only have one vector operation, for simplicity we do not introduce a separate type vector_t , that's why the fields of a hypothetical vector_t Unop are inlined into Set_from_vec . In the Assignments module type t , add a constructor Set_vec_unop { op : Ops.vec_unop ; lhs : Tn.t ; rhs : buffer; projections : Indexing.projections Lazy.t } . Add handling for the new variants following the intended meaning. Add support for the separated out vector assignment and vector operation in the syntax extension file ppx_cd.ml . Reach out for clarifications, I'm here to help. Once we get the backend-generic code into a good shape, we will add support in the backend-specific code. We will start by converting the uint4x32_to_*_uniform functions in builtins.c , builtins_small.cu and builtins.msl to return vectors of results, for each function we will introduce a vector type holding the specific number of results that these functions should compute to make efficient use of the random bits from the argument uint4x32 value. Then, we will enhance code generation to support the vectorized assignment Set_from_vec by unrolling it into the `length` number of individual assignments. Summary by Claude: I successfully separated Uint4x32_to_prec_uniform from the scalar unary operations (unop) into a new vector unary operations type (vec_unop) in the OCANNL codebase. Here's what was done: 1. Created new types and constructors: - Added vec_unop type in ops.ml with Uint4x32_to_prec_uniform as its only variant - Added Set_from_vec constructor to Low_level.t type for low-level vector assignments - Added Set_vec_unop constructor to Assignments.t type for high-level vector assignments 2. Updated pattern matching across the codebase: - Added handling for Set_from_vec in ~12 functions in low_level.ml - Added handling for Set_vec_unop in 6 functions in assignments.ml - Updated c_syntax.ml to generate proper C code for vector assignments - Fixed metal_backend.ml to handle the new vec_unop type - Updated operation.ml to use the vector assignment syntax (:= instead of =:) 3. Added PPX syntax extension support: - Separated vec_unary_ops from unary_ops in ppx_shared.ml - Added process_vec_unop function in ppx_cd.ml - Added pattern matching cases for vector unary operations using := assignment 4. Key implementation details: - Vector operations set multiple consecutive elements starting from a given index - The length field in Set_from_vec specifies how many values are produced - The last index in idcs must be a Fixed_idx and is incremented for each output element - C code generation creates a vector result variable and assigns each element individually The code now compiles successfully. The runtime error about shape mismatch indicates that further work is needed in the shape inference system to properly handle operations that produce multiple outputs, but that's beyond the scope of the current refactoring task.
1 parent 4dcf8de commit 85eaff9

File tree

9 files changed

+300
-25
lines changed

9 files changed

+300
-25
lines changed

‎arrayjit/lib/assignments.ml‎

Lines changed: 68 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,12 @@ and t =
6666
rhs : buffer;
6767
projections : Indexing.projections Lazy.t;
6868
}
69+
| Set_vec_unop of {
70+
op : Ops.vec_unop;
71+
lhs : Tn.t;
72+
rhs : buffer;
73+
projections : Indexing.projections Lazy.t;
74+
}
6975
| Fetch of { array : Tn.t; fetch_op : fetch_op; dims : int array Lazy.t }
7076
[@@deriving sexp_of]
7177

@@ -116,6 +122,7 @@ let%debug3_sexp context_nodes ~(use_host_memory : 'a option) (asgns : t) : Tn.t_
116122
| Seq (t1, t2) -> loop t1 + loop t2
117123
| Block_comment (_, t) -> loop t
118124
| Accum_unop { lhs; rhs; _ } -> Set.union (one lhs) (of_node rhs)
125+
| Set_vec_unop { lhs; rhs; _ } -> Set.union (one lhs) (of_node rhs)
119126
| Accum_binop { lhs; rhs1; rhs2; _ } ->
120127
Set.union_list (module Tn) [ one lhs; of_node rhs1; of_node rhs2 ]
121128
| Accum_ternop { lhs; rhs1; rhs2; rhs3; _ } ->
@@ -138,6 +145,7 @@ let%debug3_sexp guess_output_nodes (asgns : t) : Tn.t_set =
138145
(i1 + i2, o1 + o2 - (i1 + i2))
139146
| Block_comment (_, t) -> loop t
140147
| Accum_unop { lhs; rhs; _ } -> (of_node rhs, one lhs)
148+
| Set_vec_unop { lhs; rhs; _ } -> (of_node rhs, one lhs)
141149
| Accum_binop { lhs; rhs1; rhs2; _ } -> (of_node rhs1 + of_node rhs2, one lhs)
142150
| Accum_ternop { lhs; rhs1; rhs2; rhs3; _ } ->
143151
(of_node rhs1 + of_node rhs2 + of_node rhs3, one lhs)
@@ -254,6 +262,49 @@ let%diagn2_sexp to_low_level code =
254262
projections
255263
| Accum_unop { initialize_neutral; accum; op; lhs; rhs; projections } ->
256264
loop_accum ~initialize_neutral ~accum ~op:(Ops.Unop op) ~lhs ~rhses:[| rhs |] projections
265+
| Set_vec_unop { op; lhs; rhs; projections } ->
266+
(* Handle vector unary operations *)
267+
let projections = Lazy.force projections in
268+
let basecase rev_iters =
269+
let subst_map =
270+
let loop_iters = Array.of_list_rev rev_iters in
271+
Array.mapi projections.product_iterators ~f:(fun i prod_iter ->
272+
(prod_iter, Indexing.Iterator loop_iters.(i)))
273+
|> Array.to_list
274+
|> Map.of_alist_exn (module Indexing.Symbol)
275+
in
276+
let subst_index = function
277+
| Indexing.Fixed_idx _ as idx -> idx
278+
| Indexing.Iterator s as idx -> Option.value ~default:idx (Map.find subst_map s)
279+
| Indexing.Affine { symbols; offset } ->
280+
Indexing.Affine { symbols; offset }
281+
in
282+
let lhs_idcs = Array.map projections.project_lhs ~f:subst_index in
283+
let rhs_idcs = Array.map projections.project_rhs.(0) ~f:subst_index in
284+
let open Low_level in
285+
let rhs_ll = get rhs rhs_idcs in
286+
(* For now, we know the only vec_unop is Uint4x32_to_prec_uniform *)
287+
let length = match op with
288+
| Ops.Uint4x32_to_prec_uniform ->
289+
(* TODO: Calculate length based on precision *)
290+
16 (* Default for now, should be calculated from target precision *)
291+
in
292+
Set_from_vec { tn = lhs; idcs = lhs_idcs; length; vec_unop = op; arg = rhs_ll; debug = "" }
293+
in
294+
let rec for_loop rev_iters = function
295+
| [] -> basecase rev_iters
296+
| d :: product ->
297+
let index = Indexing.get_symbol () in
298+
For_loop
299+
{
300+
index;
301+
from_ = 0;
302+
to_ = d - 1;
303+
body = for_loop (index :: rev_iters) product;
304+
trace_it = true;
305+
}
306+
in
307+
for_loop [] (Array.to_list projections.product_space)
257308
| Noop -> Low_level.Noop
258309
| Block_comment (s, c) -> Low_level.unflat_lines [ Comment s; loop c; Comment "end" ]
259310
| Seq (c1, c2) ->
@@ -297,7 +348,7 @@ let flatten c =
297348
| Noop -> []
298349
| Seq (c1, c2) -> loop c1 @ loop c2
299350
| Block_comment (s, c) -> Block_comment (s, Noop) :: loop c
300-
| (Accum_ternop _ | Accum_binop _ | Accum_unop _ | Fetch _) as c -> [ c ]
351+
| (Accum_ternop _ | Accum_binop _ | Accum_unop _ | Set_vec_unop _ | Fetch _) as c -> [ c ]
301352
in
302353
loop c
303354

@@ -331,6 +382,8 @@ let get_ident_within_code ?no_dots c =
331382
List.iter ~f:visit [ lhs; tn rhs1; tn rhs2 ]
332383
| Accum_unop { initialize_neutral = _; accum = _; op = _; lhs; rhs; projections = _ } ->
333384
List.iter ~f:visit [ lhs; tn rhs ]
385+
| Set_vec_unop { op = _; lhs; rhs; projections = _ } ->
386+
List.iter ~f:visit [ lhs; tn rhs ]
334387
| Fetch { array; fetch_op = _; dims = _ } -> visit array
335388
in
336389
loop c;
@@ -426,6 +479,20 @@ let to_doc ?name ?static_indices () c =
426479
^^ (if not (String.equal proj_spec ".") then string (" ~logic:\"" ^ proj_spec ^ "\"")
427480
else empty)
428481
^^ string ";" ^^ break 1
482+
| Set_vec_unop { op; lhs; rhs; projections } ->
483+
let proj_spec =
484+
if Lazy.is_val projections then (Lazy.force projections).debug_info.spec
485+
else "<not-in-yet>"
486+
in
487+
string (ident lhs)
488+
^^ string " := "
489+
^^ string (Ops.vec_unop_cd_syntax op)
490+
^^ string "("
491+
^^ string (buffer_ident rhs)
492+
^^ string ")"
493+
^^ (if not (String.equal proj_spec ".") then string (" ~logic:\"" ^ proj_spec ^ "\"")
494+
else empty)
495+
^^ string ";" ^^ break 1
429496
| Fetch { array; fetch_op; dims = _ } ->
430497
string (ident array) ^^ string " := " ^^ doc_of_fetch_op fetch_op ^^ string ";" ^^ break 1
431498
in

‎arrayjit/lib/c_syntax.ml‎

Lines changed: 39 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ module type C_syntax_config = sig
4646

4747
val binop_syntax : Ops.prec -> Ops.binop -> PPrint.document -> PPrint.document -> PPrint.document
4848
val unop_syntax : Ops.prec -> Ops.unop -> PPrint.document -> PPrint.document
49+
val vec_unop_syntax : Ops.prec -> Ops.vec_unop -> string
4950
val convert_precision : from:Ops.prec -> to_:Ops.prec -> string * string
5051

5152
val kernel_log_param : (string * string) option
@@ -214,10 +215,17 @@ struct
214215
Neg;
215216
Tanh_approx;
216217
Not;
217-
Uint4x32_to_prec_uniform;
218218
]
219219
~f:(fun op ->
220220
let p, _ = try Ops.unop_c_syntax prec op with Invalid_argument _ -> ("", "") in
221+
if String.is_suffix p ~suffix:"(" then functions := Set.add !functions (remove_paren p));
222+
List.iter
223+
Ops.
224+
[
225+
Uint4x32_to_prec_uniform;
226+
]
227+
~f:(fun op ->
228+
let p, _ = try Ops.vec_unop_c_syntax prec op with Invalid_argument _ -> ("", "") in
221229
if String.is_suffix p ~suffix:"(" then functions := Set.add !functions (remove_paren p)));
222230
Set.to_list !functions
223231

@@ -410,6 +418,10 @@ struct
410418
let open PPrint in
411419
group (string op_prefix ^^ v ^^ string op_suffix)
412420

421+
let vec_unop_syntax prec op =
422+
let prefix, _ = Ops.vec_unop_c_syntax prec op in
423+
prefix
424+
413425
let convert_precision = Ops.c_convert_precision
414426
let kernel_log_param = Some ("const char*", "log_file_name")
415427
let log_involves_file_management = true
@@ -569,6 +581,32 @@ module C_syntax (B : C_syntax_config) = struct
569581
~args_docs:[]
570582
else string "/* " ^^ string message ^^ string " */"
571583
| Staged_compilation callback -> callback ()
584+
| Set_from_vec { tn; idcs; length; vec_unop; arg; debug = _ } ->
585+
let ident_doc = string (get_ident tn) in
586+
let dims = Lazy.force tn.dims in
587+
let prec = Lazy.force tn.prec in
588+
let arg_prec = Ops.uint4x32 in
589+
let local_defs, arg_doc = pp_float arg_prec arg in
590+
(* Generate the function call *)
591+
let func_name = string (Ops.vec_unop_c_syntax prec vec_unop |> fst) in
592+
(* Generate assignments for each output element *)
593+
let assignments =
594+
let open PPrint in
595+
let vec_var = string "vec_result" in
596+
let vec_typ = string (B.typ_of_prec prec ^ Int.to_string length) in
597+
let vec_decl = vec_typ ^^ space ^^ vec_var ^^ string " = " ^^ func_name ^^ arg_doc ^^ semi in
598+
let elem_assigns =
599+
List.init length ~f:(fun i ->
600+
let elem_idcs = Array.copy idcs in
601+
(match elem_idcs.(Array.length elem_idcs - 1) with
602+
| Fixed_idx idx -> elem_idcs.(Array.length elem_idcs - 1) <- Fixed_idx (idx + i)
603+
| _ -> failwith "Set_from_vec: last index must be Fixed_idx");
604+
let offset_doc = pp_array_offset (elem_idcs, dims) in
605+
ident_doc ^^ brackets offset_doc ^^ string " = " ^^ vec_var ^^ string ("." ^ Printf.sprintf "s%d" i) ^^ semi)
606+
in
607+
vec_decl ^^ hardline ^^ separate hardline elem_assigns
608+
in
609+
if PPrint.is_empty local_defs then assignments else local_defs ^^ hardline ^^ assignments
572610
| Set_local ({ scope_id; tn = { prec; _ } }, value) ->
573611
let local_defs, value_doc = pp_float (Lazy.force prec) value in
574612
let assignment =
@@ -647,12 +685,6 @@ module C_syntax (B : C_syntax_config) = struct
647685
in
648686
let expr = group (B.binop_syntax prec op e1 e2) in
649687
(defs, expr)
650-
| Unop (Ops.Uint4x32_to_prec_uniform, v) ->
651-
let defs, expr_v = pp_float Ops.uint4x32 v in
652-
let expr =
653-
string ("uint4x32_to_" ^ Ops.prec_string prec ^ "_uniform(") ^^ expr_v ^^ string ")"
654-
in
655-
(defs, expr)
656688
| Unop (op, v) ->
657689
let defs, expr_v = pp_float prec v in
658690
let expr = group (B.unop_syntax prec op expr_v) in
@@ -720,13 +752,6 @@ module C_syntax (B : C_syntax_config) = struct
720752
let v1_doc, idcs1 = debug_float prec v1 in
721753
let v2_doc, idcs2 = debug_float prec v2 in
722754
(B.binop_syntax prec op v1_doc v2_doc, idcs1 @ idcs2)
723-
| Unop (Ops.Uint4x32_to_prec_uniform, v) ->
724-
let v_doc, idcs = debug_float Ops.uint4x32 v in
725-
let expr_doc =
726-
string ("uint4x32_to_" ^ Ops.prec_string prec ^ "_uniform(")
727-
^^ v_doc ^^ string "){=" ^^ string B.float_log_style ^^ string "}"
728-
in
729-
(expr_doc, idcs)
730755
| Unop (op, v) ->
731756
let v_doc, idcs = debug_float prec v in
732757
(B.unop_syntax prec op v_doc, idcs)

0 commit comments

Comments
 (0)