Wire binops up to promotion

Author: WardBrian

src/analysis_and_optimization/Partial_evaluator.ml

@@ -6,11 +6,12 @@ open Middle
 
 exception Rejected of Location_span.t * string
 
-let is_int i Expr.Fixed.{pattern; _} =
+let rec is_int i Expr.Fixed.{pattern; _} =
   let nums = List.map ~f:(fun s -> string_of_int i ^ s) [""; "."; ".0"] in
   match pattern with
   | (Lit (Int, i) | Lit (Real, i)) when List.mem nums i ~equal:String.equal ->
       true
+  | Promotion (e, _, _) -> is_int i e
   | _ -> false
 
 let apply_prefix_operator_int (op : string) i =
@@ -108,7 +109,8 @@ let rec eval_expr ?(preserve_stability = false) (e : Expr.Typed.t) =
                 |> Option.value_map
                      ~f:(fun op ->
                        Frontend.Typechecker.operator_stan_math_return_type op
-                         argument_types )
+                         argument_types
+                       |> Option.map ~f:fst )
                      ~default:
                        (Frontend.Typechecker.stan_math_return_type name
                           argument_types ) in

src/frontend/Canonicalize.ml

@@ -146,7 +146,7 @@ let rec no_parens {expr; emeta} =
   | Variable _ | IntNumeral _ | RealNumeral _ | ImagNumeral _ | GetLP
    |GetTarget ->
       {expr; emeta}
-  | TernaryIf _ | BinOp _ | PrefixOp _ | PostfixOp _ ->
+  | TernaryIf _ | BinOp _ | PrefixOp _ | PostfixOp _ | Promotion _ ->
       {expr= map_expression keep_parens ident expr; emeta}
   | Indexed (e, l) ->
       { expr=
@@ -158,7 +158,7 @@ let rec no_parens {expr; emeta} =
                   | i -> map_index keep_parens i )
                 l )
       ; emeta }
-  | ArrayExpr _ | RowVectorExpr _ | FunApp _ | CondDistApp _ | Promotion _ ->
+  | ArrayExpr _ | RowVectorExpr _ | FunApp _ | CondDistApp _ ->
       {expr= map_expression no_parens ident expr; emeta}
 
 and keep_parens {expr; emeta} =

src/frontend/Typechecker.ml

@@ -205,13 +205,15 @@ let stan_math_return_type name arg_tys =
 let operator_stan_math_return_type op arg_tys =
   match (op, arg_tys) with
   | Operator.IntDivide, [(_, UnsizedType.UInt); (_, UInt)] ->
-      Some UnsizedType.(ReturnType UInt)
+      Some (UnsizedType.(ReturnType UInt), [Promotion.NoPromotion; NoPromotion])
   | IntDivide, _ -> None
   | _ ->
       Stan_math_signatures.operator_to_stan_math_fns op
       |> List.filter_map ~f:(fun name ->
              SignatureMismatch.matching_stanlib_function name arg_tys
-             |> match_to_rt_option )
+             |> function
+             | SignatureMismatch.UniqueMatch (rt, _, p) -> Some (rt, p)
+             | _ -> None )
       |> List.hd
 
 let assignmentoperator_stan_math_return_type assop arg_tys =
@@ -220,7 +222,7 @@ let assignmentoperator_stan_math_return_type assop arg_tys =
       SignatureMismatch.matching_stanlib_function "divide" arg_tys
       |> match_to_rt_option
   | Plus | Minus | Times | EltTimes | EltDivide ->
-      operator_stan_math_return_type assop arg_tys
+      operator_stan_math_return_type assop arg_tys |> Option.map ~f:fst
   | _ -> None )
   |> Option.bind ~f:(function
        | ReturnType rtype
@@ -234,9 +236,9 @@ let assignmentoperator_stan_math_return_type assop arg_tys =
 let check_binop loc op le re =
   let rt = [le; re] |> get_arg_types |> operator_stan_math_return_type op in
   match rt with
-  | Some (ReturnType type_) ->
+  | Some (ReturnType type_, [p1; p2]) ->
       mk_typed_expression
-        ~expr:(BinOp (le, op, re))
+        ~expr:(BinOp (Promotion.promote le p1, op, Promotion.promote re p2))
         ~ad_level:(expr_ad_lub [le; re])
         ~type_ ~loc
   | _ ->
@@ -246,7 +248,7 @@ let check_binop loc op le re =
 let check_prefixop loc op te =
   let rt = operator_stan_math_return_type op [arg_type te] in
   match rt with
-  | Some (ReturnType type_) ->
+  | Some (ReturnType type_, _) ->
       mk_typed_expression
         ~expr:(PrefixOp (op, te))
         ~ad_level:(expr_ad_lub [te])
@@ -256,7 +258,7 @@ let check_prefixop loc op te =
 let check_postfixop loc op te =
   let rt = operator_stan_math_return_type op [arg_type te] in
   match rt with
-  | Some (ReturnType type_) ->
+  | Some (ReturnType type_, _) ->
       mk_typed_expression
         ~expr:(PostfixOp (te, op))
         ~ad_level:(expr_ad_lub [te])

src/frontend/Typechecker.mli

@@ -32,7 +32,7 @@ val check_program :
 val operator_stan_math_return_type :
      Middle.Operator.t
   -> (Middle.UnsizedType.autodifftype * Middle.UnsizedType.t) list
-  -> Middle.UnsizedType.returntype option
+  -> (Middle.UnsizedType.returntype * Promotion.t list) option
 
 val stan_math_return_type :
      string

src/stan_math_backend/Expression_gen.ml

@@ -229,6 +229,9 @@ and pp_scalar_binary ppf op fn es =
   else pp_binary_f ppf fn es
 
 and gen_operator_app op ppf es =
+  let remove_basic_promotion (e : 'a Expr.Fixed.t) =
+    match e.pattern with Promotion (e, UReal, _) -> e | _ -> e in
+  let es = List.map ~f:remove_basic_promotion es in
   match op with
   | Operator.Plus -> pp_scalar_binary ppf "+" "stan::math::add" es
   | PMinus ->

src/stan_math_backend/Statement_gen.ml

@@ -305,7 +305,6 @@ let pp_bool_expr ppf expr =
   | _ -> pp_expr ppf expr
 
 let rec pp_statement (ppf : Format.formatter) Stmt.Fixed.{pattern; meta} =
-  (* ({stmt; smeta} : (mtype_loc_ad, 'a) stmt_with) = *)
   let remove_promotions (e : 'a Expr.Fixed.t) =
     (* assignment handles one level of promotion internally, don't do it twice *)
     match e.pattern with Promotion (e, _, _) -> e | _ -> e in

test/integration/good/code-gen/cpp.expected

@@ -2570,123 +2570,124 @@ class complex_scalar_model final : public model_base_crtp<complex_scalar_model>
       current_statement__ = 125;
       gq_complex = (z + y);
       current_statement__ = 126;
-      gq_complex = (z + gq_r);
+      gq_complex = (z + stan::math::to_complex(gq_r, 0));
       current_statement__ = 127;
-      gq_complex = (gq_r + z);
+      gq_complex = (stan::math::to_complex(gq_r, 0) + z);
       current_statement__ = 128;
-      gq_complex = (z + gq_i);
+      gq_complex = (z + stan::math::to_complex(gq_i, 0));
       current_statement__ = 129;
-      gq_complex = (gq_i + z);
+      gq_complex = (stan::math::to_complex(gq_i, 0) + z);
       current_statement__ = 130;
       gq_complex = (d_complex + p_complex);
       current_statement__ = 131;
-      gq_complex = (d_complex + d_r);
+      gq_complex = (d_complex + stan::math::to_complex(d_r, 0));
       current_statement__ = 132;
-      gq_complex = (d_complex + p_r);
+      gq_complex = (d_complex + stan::math::to_complex(p_r, 0));
       current_statement__ = 133;
-      gq_complex = (d_r + p_complex);
+      gq_complex = (stan::math::to_complex(d_r, 0) + p_complex);
       current_statement__ = 134;
-      gq_complex = (p_complex + p_r);
+      gq_complex = (p_complex + stan::math::to_complex(p_r, 0));
       current_statement__ = 135;
-      gq_complex = (d_complex + gq_i);
+      gq_complex = (d_complex + stan::math::to_complex(gq_i, 0));
       current_statement__ = 136;
-      gq_complex = (gq_i + p_complex);
+      gq_complex = (stan::math::to_complex(gq_i, 0) + p_complex);
       current_statement__ = 137;
       gq_complex = (z - y);
       current_statement__ = 138;
-      gq_complex = (z - gq_r);
+      gq_complex = (z - stan::math::to_complex(gq_r, 0));
       current_statement__ = 139;
-      gq_complex = (gq_r - z);
+      gq_complex = (stan::math::to_complex(gq_r, 0) - z);
       current_statement__ = 140;
-      gq_complex = (z - gq_i);
+      gq_complex = (z - stan::math::to_complex(gq_i, 0));
       current_statement__ = 141;
-      gq_complex = (gq_i - z);
+      gq_complex = (stan::math::to_complex(gq_i, 0) - z);
       current_statement__ = 142;
       gq_complex = (d_complex - p_complex);
       current_statement__ = 143;
-      gq_complex = (d_complex - d_r);
+      gq_complex = (d_complex - stan::math::to_complex(d_r, 0));
       current_statement__ = 144;
-      gq_complex = (d_complex - p_r);
+      gq_complex = (d_complex - stan::math::to_complex(p_r, 0));
       current_statement__ = 145;
-      gq_complex = (d_r - p_complex);
+      gq_complex = (stan::math::to_complex(d_r, 0) - p_complex);
       current_statement__ = 146;
-      gq_complex = (p_complex - p_r);
+      gq_complex = (p_complex - stan::math::to_complex(p_r, 0));
       current_statement__ = 147;
-      gq_complex = (d_complex - gq_i);
+      gq_complex = (d_complex - stan::math::to_complex(gq_i, 0));
       current_statement__ = 148;
-      gq_complex = (gq_i - p_complex);
+      gq_complex = (stan::math::to_complex(gq_i, 0) - p_complex);
       current_statement__ = 149;
       gq_complex = (z * y);
       current_statement__ = 150;
-      gq_complex = (z * gq_r);
+      gq_complex = (z * stan::math::to_complex(gq_r, 0));
       current_statement__ = 151;
-      gq_complex = (gq_r * z);
+      gq_complex = (stan::math::to_complex(gq_r, 0) * z);
       current_statement__ = 152;
-      gq_complex = (z * gq_i);
+      gq_complex = (z * stan::math::to_complex(gq_i, 0));
       current_statement__ = 153;
-      gq_complex = (gq_i * z);
+      gq_complex = (stan::math::to_complex(gq_i, 0) * z);
       current_statement__ = 154;
       gq_complex = (d_complex * p_complex);
       current_statement__ = 155;
-      gq_complex = (d_complex * d_r);
+      gq_complex = (d_complex * stan::math::to_complex(d_r, 0));
       current_statement__ = 156;
-      gq_complex = (d_complex * p_r);
+      gq_complex = (d_complex * stan::math::to_complex(p_r, 0));
       current_statement__ = 157;
-      gq_complex = (d_r * p_complex);
+      gq_complex = (stan::math::to_complex(d_r, 0) * p_complex);
       current_statement__ = 158;
-      gq_complex = (p_complex * p_r);
+      gq_complex = (p_complex * stan::math::to_complex(p_r, 0));
       current_statement__ = 159;
-      gq_complex = (d_complex * gq_i);
+      gq_complex = (d_complex * stan::math::to_complex(gq_i, 0));
       current_statement__ = 160;
-      gq_complex = (gq_i * p_complex);
+      gq_complex = (stan::math::to_complex(gq_i, 0) * p_complex);
       current_statement__ = 161;
       gq_complex = (z / y);
       current_statement__ = 162;
-      gq_complex = (z / gq_r);
+      gq_complex = (z / stan::math::to_complex(gq_r, 0));
       current_statement__ = 163;
-      gq_complex = (gq_r / z);
+      gq_complex = (stan::math::to_complex(gq_r, 0) / z);
       current_statement__ = 164;
-      gq_complex = (z / gq_i);
+      gq_complex = (z / stan::math::to_complex(gq_i, 0));
       current_statement__ = 165;
-      gq_complex = (gq_i / z);
+      gq_complex = (stan::math::to_complex(gq_i, 0) / z);
       current_statement__ = 166;
       gq_complex = (d_complex / p_complex);
       current_statement__ = 167;
-      gq_complex = (d_complex / d_r);
+      gq_complex = (d_complex / stan::math::to_complex(d_r, 0));
       current_statement__ = 168;
-      gq_complex = (d_complex / p_r);
+      gq_complex = (d_complex / stan::math::to_complex(p_r, 0));
       current_statement__ = 169;
-      gq_complex = (d_r / p_complex);
+      gq_complex = (stan::math::to_complex(d_r, 0) / p_complex);
       current_statement__ = 170;
-      gq_complex = (p_complex / p_r);
+      gq_complex = (p_complex / stan::math::to_complex(p_r, 0));
       current_statement__ = 171;
-      gq_complex = (d_complex / gq_i);
+      gq_complex = (d_complex / stan::math::to_complex(gq_i, 0));
       current_statement__ = 172;
-      gq_complex = (gq_i / p_complex);
+      gq_complex = (stan::math::to_complex(gq_i, 0) / p_complex);
       current_statement__ = 173;
       gq_complex = stan::math::pow(z, y);
       current_statement__ = 174;
       gq_complex = stan::math::pow(z, gq_r);
       current_statement__ = 175;
-      gq_complex = stan::math::pow(gq_r, z);
+      gq_complex = stan::math::pow(stan::math::to_complex(gq_r, 0), z);
       current_statement__ = 176;
       gq_complex = stan::math::pow(z, gq_i);
       current_statement__ = 177;
-      gq_complex = stan::math::pow(gq_i, z);
+      gq_complex = stan::math::pow(stan::math::to_complex(gq_i, 0), z);
       current_statement__ = 178;
       gq_complex = stan::math::pow(d_complex, p_complex);
       current_statement__ = 179;
       gq_complex = stan::math::pow(d_complex, d_r);
       current_statement__ = 180;
       gq_complex = stan::math::pow(d_complex, p_r);
       current_statement__ = 181;
-      gq_complex = stan::math::pow(d_r, p_complex);
+      gq_complex = stan::math::pow(stan::math::to_complex(d_r, 0), p_complex);
       current_statement__ = 182;
       gq_complex = stan::math::pow(p_complex, p_r);
       current_statement__ = 183;
       gq_complex = stan::math::pow(d_complex, gq_i);
       current_statement__ = 184;
-      gq_complex = stan::math::pow(gq_i, p_complex);
+      gq_complex = stan::math::pow(stan::math::to_complex(gq_i, 0),
+                     p_complex);
       current_statement__ = 185;
       gq_complex = -z;
       current_statement__ = 186;
@@ -2724,7 +2725,8 @@ class complex_scalar_model final : public model_base_crtp<complex_scalar_model>
       current_statement__ = 202;
       gq_i = stan::math::logical_eq(d_complex, d_r);
       current_statement__ = 203;
-      gq_i = stan::math::logical_eq(p_r, d_complex);
+      gq_i = stan::math::logical_eq(stan::math::to_complex(p_r, 0),
+               d_complex);
       current_statement__ = 204;
       gq_i = stan::math::logical_eq(p_complex, d_r);
       current_statement__ = 205;
@@ -2744,7 +2746,8 @@ class complex_scalar_model final : public model_base_crtp<complex_scalar_model>
       current_statement__ = 212;
       gq_i = stan::math::logical_neq(d_complex, d_r);
       current_statement__ = 213;
-      gq_i = stan::math::logical_neq(p_r, d_complex);
+      gq_i = stan::math::logical_neq(stan::math::to_complex(p_r, 0),
+               d_complex);
       current_statement__ = 214;
       gq_i = stan::math::logical_neq(p_complex, d_r);
       current_statement__ = 215;
@@ -3017,7 +3020,7 @@ class complex_scalar_model final : public model_base_crtp<complex_scalar_model>
          std::complex<double>(std::numeric_limits<double>::quiet_NaN(),
            std::numeric_limits<double>::quiet_NaN());
       current_statement__ = 336;
-      zi = (1 + stan::math::to_complex(0, 3.14));
+      zi = (stan::math::to_complex(1, 0) + stan::math::to_complex(0, 3.14));
       current_statement__ = 337;
       zi = (zi * stan::math::to_complex(0, 0));
       std::complex<double> yi =