Optimization fixes

src/analysis_and_optimization/Monotone_framework.ml

@@ -243,6 +243,7 @@ let constant_propagation_transfer
 (** The transfer function for an expression propagation analysis,
     AKA forward substitution (see page 396 of Muchnick) *)
 let expression_propagation_transfer
+    (can_side_effect_expr : Middle.expr_typed_located -> bool)
     (flowgraph_to_mir : (int, Middle.stmt_loc_num) Map.Poly.t) =
   ( module struct
     type labels = int
@@ -258,7 +259,8 @@ let expression_propagation_transfer
             (* TODO: we are currently only propagating constants for scalars.
              We could do the same for matrix and array expressions if we wanted. *)
             | Middle.Assignment ((s, _, []), e) ->
-                Map.set m ~key:s ~data:(subst_expr m e)
+                if can_side_effect_expr e then m
+                else Map.set m ~key:s ~data:(subst_expr m e)
             | Middle.Decl {decl_id= s; _}
              |Middle.Assignment ((s, _, _ :: _), _) ->
                 Map.remove m s

src/analysis_and_optimization/Optimize.ml

@@ -520,16 +520,15 @@ let propagation
 let constant_propagation =
   propagation Monotone_framework.constant_propagation_transfer
 
-let expression_propagation =
-  propagation Monotone_framework.expression_propagation_transfer
-
-let copy_propagation = propagation Monotone_framework.copy_propagation_transfer
-
 let rec can_side_effect_expr (e : expr_typed_located) =
   match e.expr with
   | Var _ | Lit (_, _) -> false
-  | FunApp (_, f, es) ->
-      String.suffix f 3 = "_lp" || List.exists ~f:can_side_effect_expr es
+  | FunApp (t, f, es) ->
+      String.suffix f 3 = "_lp"
+      || List.exists ~f:can_side_effect_expr es
+      || (t = CompilerInternal && f = string_of_internal_fn FnReadParam)
+      || (t = CompilerInternal && f = string_of_internal_fn FnWriteParam)
+      || (t = CompilerInternal && f = string_of_internal_fn FnUnconstrain)
   | TernaryIf (e1, e2, e3) -> List.exists ~f:can_side_effect_expr [e1; e2; e3]
   | Indexed (e, is) ->
       can_side_effect_expr e || List.exists ~f:can_side_effect_idx is
@@ -541,6 +540,12 @@ and can_side_effect_idx (i : expr_typed_located index) =
   | Single e | Upfrom e | MultiIndex e -> can_side_effect_expr e
   | Between (e1, e2) -> can_side_effect_expr e1 || can_side_effect_expr e2
 
+let expression_propagation =
+  propagation
+    (Monotone_framework.expression_propagation_transfer can_side_effect_expr)
+
+let copy_propagation = propagation Monotone_framework.copy_propagation_transfer
+
 let is_skip_break_continue s =
   match s with Skip | Break | Continue -> true | _ -> false
 
@@ -807,8 +812,13 @@ let optimize_ad_levels mir =
         (module Rev_Flowgraph)
         flowgraph_to_mir initial_ad_variables
     in
+    let insert_constraint_variables vars =
+      Set.Poly.union vars (Set.Poly.map ~f:(fun x -> x ^ "_in__") vars)
+    in
     let optimize_ad_levels_stmt_base i stmt =
-      let autodiffable_variables = (Map.find_exn ad_levels i).exit in
+      let autodiffable_variables =
+        insert_constraint_variables (Map.find_exn ad_levels i).exit
+      in
       match
         map_statement
           (update_expr_ad_levels autodiffable_variables)

src/analysis_and_optimization/Partial_evaluator.ml

@@ -602,19 +602,19 @@ let rec eval_expr (e : Middle.expr_typed_located) =
               match op with
               | "PPlus__" | "PMinus__" | "PNot__" ->
                   apply_prefix_operator_int op (Int.of_string i)
-              | _ -> FunApp (StanLib, op, l) )
+              | _ -> FunApp (t, op, l) )
             | op, [{expr= Lit (Real, r); _}] -> (
               match op with
               | "PPlus__" | "PMinus__" ->
                   apply_prefix_operator_real op (Float.of_string r)
-              | _ -> FunApp (StanLib, op, l) )
+              | _ -> FunApp (t, op, l) )
             | op, [{expr= Lit (Int, i1); _}; {expr= Lit (Int, i2); _}] -> (
               match op with
               | "Plus__" | "Minus__" | "Times__" | "Divide__" | "Modulo__"
                |"Or__" | "And__" | "Equals__" | "NEquals__" | "Less__"
                |"Leq__" | "Greater__" | "Geq__" ->
                   apply_operator_int op (Int.of_string i1) (Int.of_string i2)
-              | _ -> FunApp (StanLib, op, l) )
+              | _ -> FunApp (t, op, l) )
             | op, [{expr= Lit (Real, i1); _}; {expr= Lit (Real, i2); _}]
              |op, [{expr= Lit (Int, i1); _}; {expr= Lit (Real, i2); _}]
              |op, [{expr= Lit (Real, i1); _}; {expr= Lit (Int, i2); _}] -> (
@@ -626,7 +626,7 @@ let rec eval_expr (e : Middle.expr_typed_located) =
                |"Leq__" | "Greater__" | "Geq__" ->
                   apply_logical_operator_real op (Float.of_string i1)
                     (Float.of_string i2)
-              | _ -> FunApp (StanLib, op, l) )
+              | _ -> FunApp (t, op, l) )
             | _ -> FunApp (t, f, l) )
       | TernaryIf (e1, e2, e3) -> (
         match (eval_expr e1, eval_expr e2, eval_expr e3) with

src/stan_math_backend/Stan_math_code_gen.ml

@@ -118,7 +118,13 @@ let pp_fun_def ppf {fdrt; fdname; fdargs; fdbody; _} =
     text
       "local_scalar_t__ DUMMY_VAR__(std::numeric_limits<double>::quiet_NaN());" ;
     pp_unused ppf "DUMMY_VAR__" ;
-    pp_located_error ppf (pp_statement, fdbody) ;
+    let blocked_fdbody =
+      match fdbody.stmt with
+      | SList stmts -> {stmt= Block stmts; smeta= fdbody.smeta}
+      | Block _ -> fdbody
+      | _ -> {stmt= Block [fdbody]; smeta= fdbody.smeta}
+    in
+    pp_located_error ppf (pp_statement, blocked_fdbody) ;
     pf ppf "@ "
   in
   let templates =

src/stan_math_backend/Statement_gen.ml

@@ -151,7 +151,7 @@ let rec pp_statement (ppf : Format.formatter)
   | Break -> string ppf "break;"
   | Continue -> string ppf "continue;"
   | Return e -> pf ppf "return %a;" (option pp_expr) e
-  | Skip -> ()
+  | Skip -> string ppf ";"
   | IfElse (cond, ifbranch, elsebranch) ->
       let pp_else ppf x = pf ppf "else %a" pp_statement x in
       pf ppf "if (@[<hov>%a@]) %a %a" pp_expr cond pp_block_s ifbranch

src/stanc/stanc.ml

@@ -20,9 +20,11 @@ let print_model_cpp = ref false
 let dump_mir = ref false
 let dump_mir_pretty = ref false
 let dump_tx_mir = ref false
+let dump_tx_mir_pretty = ref false
+let dump_opt_mir = ref false
+let dump_opt_mir_pretty = ref false
 let dump_stan_math_sigs = ref false
 let optimize = ref false
-let dump_opt_mir = ref false
 let output_file = ref ""
 let generate_data = ref false
 let warn_uninitialized = ref false
@@ -56,12 +58,20 @@ let options =
       , " For debugging purposes: pretty-print the MIR." )
     ; ( "--debug-optimized-mir"
       , Arg.Set dump_opt_mir
-      , " For debugging purposes: print the MIR after it's been \
-         optimized.Only has an effect when optimizations are turned on." )
+      , " For debugging purposes: print the MIR after it's been optimized. \
+         Only has an effect when optimizations are turned on." )
+    ; ( "--debug-optimized-mir-pretty"
+      , Arg.Set dump_opt_mir_pretty
+      , " For debugging purposes: pretty print the MIR after it's been \
+         optimized. Only has an effect when optimizations are turned on." )
     ; ( "--debug-transformed-mir"
       , Arg.Set dump_tx_mir
       , " For debugging purposes: print the MIR after the backend has \
          transformed it." )
+    ; ( "--debug-transformed-mir-pretty"
+      , Arg.Set dump_tx_mir_pretty
+      , " For debugging purposes: pretty print the MIR after the backend has \
+         transformed it." )
     ; ( "--dump-stan-math-signatures"
       , Arg.Set dump_stan_math_sigs
       , "Dump out the list of supported type signatures for Stan Math backend."
@@ -207,13 +217,17 @@ let use_file filename =
       print_warn_uninitialized uninitialized_vars ) ;
     let tx_mir = Transform_Mir.trans_prog mir in
     if !dump_tx_mir then
+      Sexp.pp_hum Format.std_formatter [%sexp (tx_mir : Middle.typed_prog)] ;
+    if !dump_tx_mir_pretty then
       Middle.Pretty.pp_typed_prog Format.std_formatter tx_mir ;
     let opt_mir =
       if !optimize then (
         let opt =
           Optimize.optimization_suite (optimization_settings ()) tx_mir
         in
         if !dump_opt_mir then
+          Sexp.pp_hum Format.std_formatter [%sexp (opt : Middle.typed_prog)] ;
+        if !dump_opt_mir_pretty then
           Middle.Pretty.pp_typed_prog Format.std_formatter opt ;
         opt )
       else tx_mir

test/integration/bad/stanc.expected

@@ -1064,8 +1064,8 @@ Called from file "src/frontend/Ast_to_Mir.ml", line 532, characters 12-144
 Called from file "src/list.ml", line 557, characters 34-40
 Called from file "src/frontend/Ast_to_Mir.ml" (inlined), line 561, characters 22-76
 Called from file "src/frontend/Ast_to_Mir.ml", line 653, characters 21-64
-Called from file "src/stanc/stanc.ml", line 198, characters 14-54
-Called from file "src/stanc/stanc.ml", line 243, characters 9-16
+Called from file "src/stanc/stanc.ml", line 208, characters 14-54
+Called from file "src/stanc/stanc.ml", line 257, characters 9-16
   $ ../../../../install/default/bin/stanc fun-return-type1.stan
 
 Semantic error in 'fun-return-type1.stan', line 2, column 2 to column 72: