-
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathcuda_codegen.ml
More file actions
135 lines (121 loc) · 5.37 KB
/
Copy pathcuda_codegen.ml
File metadata and controls
135 lines (121 loc) · 5.37 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
(* SPDX-License-Identifier: MPL-2.0 *)
(* SPDX-FileCopyrightText: 2026 Jonathan D.A. Jewell *)
(** CUDA C++ kernel sublanguage emitter (MVP).
Same kernel shape as the WGSL backend: first param is the global index,
remaining params are buffers. Lowers to a [__global__ void] function
plus a host wrapper that the user can call from C++. *)
open Ast
open Kernel_sublang
let mangle s = s
let scalar_of_type_name = function
| "Int" -> "int"
| "Float" -> "float"
| "Bool" -> "bool"
| n -> unsupported ("type not allowed in CUDA kernel: " ^ n)
let scalar_of (te : type_expr) : string =
match strip_ownership te with
| TyCon id -> scalar_of_type_name id.name
| _ -> unsupported "complex type not allowed in CUDA kernel"
let array_element (te : type_expr) : string =
scalar_of_type_name (require_array_element "Array[Int|Float]" te)
let const_qual = function
| Some Mut -> ""
| _ -> "const "
let rec gen_expr (e : expr) : string =
match e with
| ExprLit lit -> gen_lit lit
| ExprVar id -> mangle id.name
| ExprBinary (a, op, b) ->
let s = match op with
| OpAdd -> "+" | OpSub -> "-" | OpMul -> "*" | OpDiv -> "/" | OpMod -> "%"
| OpEq -> "==" | OpNe -> "!="
| OpLt -> "<" | OpLe -> "<=" | OpGt -> ">" | OpGe -> ">="
| OpAnd -> "&&" | OpOr -> "||"
| OpBitAnd -> "&" | OpBitOr -> "|" | OpBitXor -> "^"
| OpShl -> "<<" | OpShr -> ">>"
| OpConcat -> unsupported "concat not supported in CUDA"
in
"(" ^ gen_expr a ^ " " ^ s ^ " " ^ gen_expr b ^ ")"
| ExprUnary (OpNeg, x) -> "(-" ^ gen_expr x ^ ")"
| ExprUnary (OpNot, x) -> "(!" ^ gen_expr x ^ ")"
| ExprUnary (OpBitNot, x) -> "(~" ^ gen_expr x ^ ")"
| ExprUnary _ -> unsupported "unary op not supported in CUDA kernel"
| ExprIf { ei_cond; ei_then; ei_else } ->
let f = match ei_else with Some e -> gen_expr e | None -> "0" in
Printf.sprintf "(%s ? %s : %s)" (gen_expr ei_cond) (gen_expr ei_then) f
| ExprIndex (a, i) -> Printf.sprintf "%s[%s]" (gen_expr a) (gen_expr i)
| ExprApp (callee, args) ->
let name = match callee with
| ExprVar id -> id.name
| _ -> unsupported "indirect call"
in
if not (is_math_builtin name || name = "fabs") then
unsupported ("call to non-builtin in CUDA kernel: " ^ name);
Printf.sprintf "%s(%s)" name
(String.concat ", " (List.map gen_expr args))
| ExprSpan (inner, _) -> gen_expr inner
| _ -> unsupported "expression form not supported in CUDA kernel"
and gen_lit = function
| LitInt (n, _) -> string_of_int n
| LitFloat (f, _) ->
let s = string_of_float f in
let s = if String.length s > 0 && s.[String.length s - 1] = '.' then s ^ "0" else s in
s ^ "f"
| LitBool (true, _) -> "true"
| LitBool (false, _) -> "false"
| _ -> unsupported "literal form not supported in CUDA kernel"
let rec gen_stmt (s : stmt) : string =
match s with
| StmtLet { sl_pat = PatVar id; sl_value; sl_ty; _ } ->
let ty = match sl_ty with Some t -> scalar_of t | None -> "int" in
Printf.sprintf "%s %s = %s;" ty (mangle id.name) (gen_expr sl_value)
| StmtLet _ -> unsupported "destructuring let not supported in CUDA"
| StmtAssign (lhs, op, rhs) ->
let s = match op with
| AssignEq -> "=" | AssignAdd -> "+=" | AssignSub -> "-="
| AssignMul -> "*=" | AssignDiv -> "/=" in
Printf.sprintf "%s %s %s;" (gen_expr lhs) s (gen_expr rhs)
| StmtExpr e -> gen_expr e ^ ";"
| StmtWhile (c, b) ->
Printf.sprintf "while (%s) { %s }" (gen_expr c)
(String.concat " " (List.map gen_stmt b.blk_stmts))
| StmtFor _ -> unsupported "for-in not supported in CUDA kernel"
let pick_kernel = pick_entry
let validate_kernel = validate_compute_kernel_shape
let generate (program : program) (_symbols : Symbol.t) : string =
let buf = Buffer.create 1024 in
Buffer.add_string buf "// Generated by AffineScript compiler (CUDA C++)\n";
Buffer.add_string buf "// SPDX-License-Identifier: MPL-2.0\n\n";
Buffer.add_string buf "#include <cuda_runtime.h>\n\n";
let fd = pick_kernel program in
validate_kernel fd;
let idx = match fd.fd_params with first :: _ -> first.p_name.name | _ -> "i" in
let bufs = match fd.fd_params with _ :: rest -> rest | [] -> [] in
let buf_decls = List.map (fun (p : param) ->
Printf.sprintf "%s%s *%s"
(const_qual p.p_ownership) (array_element p.p_ty) p.p_name.name
) bufs in
Buffer.add_string buf "__global__\n";
Buffer.add_string buf
(Printf.sprintf "void %s(%s) {\n" (mangle fd.fd_name.name)
(String.concat ", " buf_decls));
Buffer.add_string buf
(Printf.sprintf " int %s = blockIdx.x * blockDim.x + threadIdx.x;\n" idx);
(match fd.fd_body with
| FnExpr e ->
Buffer.add_string buf (Printf.sprintf " (void)(%s);\n" (gen_expr e))
| FnBlock b ->
List.iter (fun s ->
Buffer.add_string buf (" " ^ gen_stmt s ^ "\n")
) b.blk_stmts;
(match b.blk_expr with
| Some e -> Buffer.add_string buf (Printf.sprintf " (void)(%s);\n" (gen_expr e))
| None -> ()));
Buffer.add_string buf "}\n";
Buffer.contents buf
let codegen_cuda (program : program) (symbols : Symbol.t) : (string, string) result =
try Ok (generate program symbols)
with
| Unsupported m -> Error ("CUDA backend: " ^ m)
| Failure m -> Error ("CUDA codegen error: " ^ m)
| e -> Error ("CUDA codegen error: " ^ Printexc.to_string e)