vm.fail
symbolic mathematics engine in OCaml with differentiation, integration, simplification, and numerical methods
1open Expr
2open Cse
3
4let c_type_of_var _ = "double"
5
6let rec expr_to_c = function
7 | Const f -> string_of_float f
8 | SymConst Pi -> "M_PI"
9 | SymConst E -> "M_E"
10 | Var v -> v
11 | Add (e1, e2) -> Printf.sprintf "(%s + %s)" (expr_to_c e1) (expr_to_c e2)
12 | Sub (e1, e2) -> Printf.sprintf "(%s - %s)" (expr_to_c e1) (expr_to_c e2)
13 | Mul (e1, e2) -> Printf.sprintf "(%s * %s)" (expr_to_c e1) (expr_to_c e2)
14 | Div (e1, e2) -> Printf.sprintf "(%s / %s)" (expr_to_c e1) (expr_to_c e2)
15 | Pow (e1, e2) -> Printf.sprintf "pow(%s, %s)" (expr_to_c e1) (expr_to_c e2)
16 | Neg e -> Printf.sprintf "(-%s)" (expr_to_c e)
17 | Sin e -> Printf.sprintf "sin(%s)" (expr_to_c e)
18 | Cos e -> Printf.sprintf "cos(%s)" (expr_to_c e)
19 | Tan e -> Printf.sprintf "tan(%s)" (expr_to_c e)
20 | Sinh e -> Printf.sprintf "sinh(%s)" (expr_to_c e)
21 | Cosh e -> Printf.sprintf "cosh(%s)" (expr_to_c e)
22 | Tanh e -> Printf.sprintf "tanh(%s)" (expr_to_c e)
23 | Asin e -> Printf.sprintf "asin(%s)" (expr_to_c e)
24 | Acos e -> Printf.sprintf "acos(%s)" (expr_to_c e)
25 | Atan e -> Printf.sprintf "atan(%s)" (expr_to_c e)
26 | Atan2 (e1, e2) -> Printf.sprintf "atan2(%s, %s)" (expr_to_c e1) (expr_to_c e2)
27 | Exp e -> Printf.sprintf "exp(%s)" (expr_to_c e)
28 | Ln e -> Printf.sprintf "log(%s)" (expr_to_c e)
29 | Log (base_e, arg) -> Printf.sprintf "(log(%s) / log(%s))" (expr_to_c arg) (expr_to_c base_e)
30 | Sqrt e -> Printf.sprintf "sqrt(%s)" (expr_to_c e)
31 | Abs e -> Printf.sprintf "fabs(%s)" (expr_to_c e)
32
33let compile_to_c expr vars =
34 let cse_result = common_subexpression_elimination expr in
35 let params = String.concat ", " (List.map (fun v -> Printf.sprintf "double %s" v) vars) in
36 let body = Buffer.create 1024 in
37
38 Buffer.add_string body "#include <math.h>\n\n";
39 Buffer.add_string body (Printf.sprintf "double compute(%s) {\n" params);
40
41 List.iter (fun (name, e) ->
42 Buffer.add_string body (Printf.sprintf " double %s = %s;\n" name (expr_to_c e))
43 ) cse_result.subexpressions;
44
45 Buffer.add_string body (Printf.sprintf " return %s;\n" (expr_to_c cse_result.final_expr));
46 Buffer.add_string body "}\n";
47
48 Buffer.contents body
49
50let compile_to_cuda expr vars =
51 let cse_result = common_subexpression_elimination expr in
52 let params = String.concat ", " (List.map (fun v -> Printf.sprintf "double %s" v) vars) in
53 let body = Buffer.create 1024 in
54
55 Buffer.add_string body "__device__ double compute_device(";
56 Buffer.add_string body params;
57 Buffer.add_string body ") {\n";
58
59 List.iter (fun (name, e) ->
60 Buffer.add_string body (Printf.sprintf " double %s = %s;\n" name (expr_to_c e))
61 ) cse_result.subexpressions;
62
63 Buffer.add_string body (Printf.sprintf " return %s;\n" (expr_to_c cse_result.final_expr));
64 Buffer.add_string body "}\n\n";
65
66 Buffer.add_string body "__global__ void compute_kernel(double* input, double* output, int n) {\n";
67 Buffer.add_string body " int idx = blockIdx.x * blockDim.x + threadIdx.x;\n";
68 Buffer.add_string body " if (idx < n) {\n";
69 Buffer.add_string body (Printf.sprintf " output[idx] = compute_device(%s);\n"
70 (String.concat ", " (List.mapi (fun i _ -> Printf.sprintf "input[idx * %d + %d]" (List.length vars) i) vars)));
71 Buffer.add_string body " }\n";
72 Buffer.add_string body "}\n";
73
74 Buffer.contents body
75
76let compile_to_vectorized expr vars =
77 let cse_result = common_subexpression_elimination expr in
78 let params = String.concat ", " (List.map (fun v -> Printf.sprintf "const double* restrict %s" v) vars) in
79 let body = Buffer.create 1024 in
80
81 Buffer.add_string body "#include <math.h>\n\n";
82 Buffer.add_string body (Printf.sprintf "void compute_vectorized(%s, double* restrict output, int n) {\n" params);
83 Buffer.add_string body " #pragma omp simd\n";
84 Buffer.add_string body " for (int i = 0; i < n; i++) {\n";
85
86 List.iter (fun (name, e) ->
87 let vec_expr = Str.global_replace (Str.regexp_string (List.hd vars)) (List.hd vars ^ "[i]") (expr_to_c e) in
88 Buffer.add_string body (Printf.sprintf " double %s = %s;\n" name vec_expr)
89 ) cse_result.subexpressions;
90
91 Buffer.add_string body (Printf.sprintf " output[i] = %s;\n" (expr_to_c cse_result.final_expr));
92 Buffer.add_string body " }\n";
93 Buffer.add_string body "}\n";
94
95 Buffer.contents body