Merge pull request #1082 from stan-dev/stancjs_optimizations

Add optimization flags, debug-* flags to stancjs

Author: WardBrian

src/frontend/Pretty_printing.ml

@@ -565,7 +565,7 @@ let pp_typed_expression ppf e =
 let pretty_print_program ?(bare_functions = false) ?(line_length = 78)
     ?(inline_includes = false) p =
   let result =
-    strf "%a" (pp_program ~bare_functions ~line_length ~inline_includes) p in
+    str "%a" (pp_program ~bare_functions ~line_length ~inline_includes) p in
   check_correctness ~bare_functions p result ;
   result
 

src/stanc/stanc.ml

@@ -301,7 +301,7 @@ let use_file filename =
       if !dump_opt_mir_pretty then Program.Typed.pp Format.std_formatter opt ;
       opt in
     if !output_file = "" then output_file := remove_dotstan !model_file ^ ".hpp" ;
-    let cpp = Fmt.strf "%a" Stan_math_code_gen.pp_prog opt_mir in
+    let cpp = Fmt.str "%a" Stan_math_code_gen.pp_prog opt_mir in
     Out_channel.write_all !output_file ~data:cpp ;
     if !print_model_cpp then print_endline cpp )
 

src/stancjs/stancjs.ml

@@ -66,6 +66,8 @@ let stan2cpp model_name model_string is_flag_set flag_val =
           flag_val "max-line-length"
           |> Option.map ~f:int_of_string
           |> Option.value ~default:78 in
+        let mir = Ast_to_Mir.trans_prog model_name typed_ast in
+        let tx_mir = Transform_Mir.trans_prog mir in
         if is_flag_set "auto-format" || is_flag_set "print-canonical" then
           r.return
             ( Result.Ok
@@ -77,16 +79,47 @@ let stan2cpp model_name model_string is_flag_set flag_val =
                       canonicalizer_settings ) )
             , warnings
             , [] ) ;
+        if is_flag_set "debug-mir" then
+          r.return
+            ( Result.Ok
+                (Sexp.to_string_hum [%sexp (mir : Middle.Program.Typed.t)])
+            , warnings
+            , [] ) ;
+        if is_flag_set "debug-mir-pretty" then
+          r.return (Result.Ok (Fmt.str "%a" Program.Typed.pp mir), warnings, []) ;
         if is_flag_set "debug-generate-data" then
           r.return
             ( Result.Ok (Debug_data_generation.print_data_prog typed_ast)
             , warnings
             , [] ) ;
-        let mir = Ast_to_Mir.trans_prog model_name typed_ast in
-        let tx_mir = Transform_Mir.trans_prog mir in
         let opt_mir =
-          if is_flag_set "O" then Optimize.optimization_suite tx_mir else tx_mir
-        in
+          let opt_lvl =
+            if is_flag_set "O0" then Optimize.O0
+            else if is_flag_set "O1" then Optimize.O1
+            else if is_flag_set "Oexperimental" || is_flag_set "O" then
+              Optimize.Oexperimental
+            else Optimize.O0 in
+          Optimize.optimization_suite
+            ~settings:(Optimize.level_optimizations opt_lvl)
+            tx_mir in
+        if is_flag_set "debug-optimized-mir" then
+          r.return
+            ( Result.Ok
+                (Sexp.to_string_hum [%sexp (opt_mir : Middle.Program.Typed.t)])
+            , warnings
+            , [] ) ;
+        if is_flag_set "debug-optimized-mir-pretty" then
+          r.return
+            (Result.Ok (Fmt.str "%a" Program.Typed.pp opt_mir), warnings, []) ;
+        if is_flag_set "debug-transformed-mir" then
+          r.return
+            ( Result.Ok
+                (Sexp.to_string_hum [%sexp (tx_mir : Middle.Program.Typed.t)])
+            , warnings
+            , [] ) ;
+        if is_flag_set "debug-transformed-mir-pretty" then
+          r.return
+            (Result.Ok (Fmt.str "%a" Program.Typed.pp tx_mir), warnings, []) ;
         let cpp = Fmt.str "%a" Stan_math_code_gen.pp_prog opt_mir in
         let uninit_warnings =
           if is_flag_set "warn-uninitialized" then
@@ -141,4 +174,9 @@ let stan2cpp_wrapped name code (flags : Js.string_array Js.t Js.opt) =
       (warnings @ pedantic_mode_warnings) in
   wrap_result ?printed_filename result ~warnings
 
-let () = Js.export "stanc" stan2cpp_wrapped
+let dump_stan_math_signatures () =
+  Js.string @@ Fmt.str "%a" Stan_math_signatures.pretty_print_all_math_sigs ()
+
+let () =
+  Js.export "dump_stan_math_signatures" dump_stan_math_signatures ;
+  Js.export "stanc" stan2cpp_wrapped

test/stancjs/debug.js

@@ -0,0 +1,40 @@
+var stanc = require('../../src/stancjs/stancjs.bc.js');
+var utils = require("./utils/utils.js");
+
+let basic_model = `
+parameters {
+real y;
+}
+model {
+y~std_normal();
+}
+`
+
+let debug_mir_test = stanc.stanc("basic", basic_model, ["debug-mir"]);
+var ind = debug_mir_test.result.search("#include \\<stan/model/");
+console.assert(ind < 0, "ERROR: MIR printing is not valid.")
+
+
+let debug_mir_pretty_test = stanc.stanc("basic", basic_model, ["debug-mir-pretty"]);
+var ind = debug_mir_pretty_test.result.search("#include \\<stan/model/");
+console.assert(ind < 0, "ERROR: MIR pretty printing is not valid.")
+
+
+let debug_opt_mir_test = stanc.stanc("basic", basic_model, ["01", "debug-optimized-mir"]);
+var ind = debug_opt_mir_test.result.search("#include \\<stan/model/");
+console.assert(ind < 0, "ERROR: Optimized MIR printing is not valid.")
+
+
+let debug_opt_mir_pretty_test = stanc.stanc("basic", basic_model, ["01", "debug-optimized-mir-pretty"]);
+var ind = debug_opt_mir_pretty_test.result.search("#include \\<stan/model/");
+console.assert(ind < 0, "ERROR: Optimized MIR pretty printing is not valid.")
+
+let debug_tx_mir_test = stanc.stanc("basic", basic_model, ["debug-transformed-mir"]);
+var ind = debug_tx_mir_test.result.search("#include \\<stan/model/");
+console.assert(ind < 0, "ERROR: Transformed MIR printing is not valid.")
+
+
+let debug_tx_mir_pretty_test = stanc.stanc("basic", basic_model, ["debug-transformed-mir-pretty"]);
+var ind = debug_tx_mir_pretty_test.result.search("#include \\<stan/model/");
+console.assert(ind < 0, "ERROR: Transformed MIR pretty printing is not valid.")
+

test/stancjs/math_sigs.js

@@ -0,0 +1,6 @@
+var stanc = require('../../src/stancjs/stancjs.bc.js');
+var utils = require("./utils/utils.js");
+
+let stan_math_sigs = stanc.dump_stan_math_signatures();
+console.assert(stan_math_sigs.includes("bernoulli_cdf(int, real) => real"), "Failed to find bernoulli signature!")
+console.assert(stan_math_sigs.includes("zeros_array(int) => array[] real"), "Failed to find zeros_array signature!")

test/stancjs/optimization.js

@@ -11,15 +11,32 @@ transformed data {
 `
 var opt_test = stanc.stanc("optimization", opt_model, []);
 var ind = opt_test.result.search("int t = 1; t <= 5; \\+\\+t");
-if (ind == -1) {
-    console.log("ERROR: Optimization without the O flag!")
-}
+console.assert(ind > -1, "ERROR: Optimization without the O flag!")
 
 var opt_test = stanc.stanc("optimization", opt_model, ["O"]);
 var ind = opt_test.result.search("int t = 1; t <= 5; \\+\\+t");
-if (ind > -1) {
-    console.log("ERROR: No optimization without the O flag!")
+console.assert(ind < 0, "ERROR: No optimization without the O flag!")
+
+var ad_model = `
+data {
+    matrix[10, 10] X_d;
+}
+parameters {
+    matrix[10, 10] X_p;
+}
+
+transformed parameters {
+    matrix[10, 10] X_tp1 = X_d;
 }
+`
+
+var opt_test = stanc.stanc("optimization", ad_model, ["O1"]);
+var ind = opt_test.result.search("\\<double, -1, -1\\> X_tp1");
+console.assert(ind > -1, "ERROR: No AD optimization with the O1 flag!")
+
+var opt_test = stanc.stanc("optimization", ad_model, ["O0"]);
+var ind = opt_test.result.search("\\<local_scalar_t__, -1, -1\\> X_tp1");
+console.assert(ind > -1, "ERROR: AD optimization without the O1 flag!")
 
 var glm_model = `
 data {
@@ -42,16 +59,14 @@ data {
 var no_opencl_test = stanc.stanc("no_opencl", glm_model);
 utils.print_error(no_opencl_test)
 var ind = no_opencl_test.result.search("matrix_cl<int> y_opencl__");
-if (ind > -1) {
-    console.log("ERROR: OpenCL code found without the use-opencl flag!")
-}
+console.assert(ind < 0, "ERROR: OpenCL code found without the use-opencl flag!")
+
 
 var opencl_test = stanc.stanc("opencl", glm_model, ["use-opencl"]);
 utils.print_error(opencl_test)
 var ind = opencl_test.result.search("matrix_cl<int> y_opencl__");
-if (ind == -1) {
-    console.log("ERROR: No OpenCL code found with the use-opencl flag!")
-}
+console.assert(ind > -1, "ERROR: No OpenCL code found with the use-opencl flag!")
+
 
 // multiple flags
 
@@ -86,6 +101,5 @@ utils.print_error(opencl_test)
 var opencl_test = stanc.stanc("opencl", glm_model2, ["use-opencl", "allow-undefined"]);
 utils.print_error(opencl_test)
 var ind = opencl_test.result.search("matrix_cl<int> y_opencl__");
-if (ind == -1) {
-    console.log("ERROR: No OpenCL code found with the use-opencl flag!")
-}
+console.assert(ind > -1, "ERROR: No OpenCL code found with the use-opencl flag!")
+

test/stancjs/stancjs.expected

@@ -67,6 +67,8 @@ dim(z) = (3, 4)
 dim(w) = (3)
 dim(p) = (4, 3)
 
+$ node debug.js
+
 $ node filename.js
 Semantic error in 'good_filename', line 6, column 4 to column 5: 
 Identifier 'z' not in scope.
@@ -179,6 +181,8 @@ $ node info.js
   "distributions": [  ],
   "included_files": [  ] }
 
+$ node math_sigs.js
+
 $ node optimization.js
 Semantic error in 'string', line 3, column 8 to column 11: 
 Function is declared without specifying a definition.