-
-
Notifications
You must be signed in to change notification settings - Fork 48
Use Eigen Maps for data and transformed data #865
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 13 commits
e2871c8
4527def
b709f6b
7d922a0
4535343
7ad77ba
b944c0c
e16a7ee
1f9ea31
d05fd8f
d59779e
303682d
e8eb9d9
f5d8cee
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -389,8 +389,8 @@ let pp_ctor ppf p = | |||||||||||||||
| match is_input_data with | ||||||||||||||||
| | true -> | ||||||||||||||||
| pp_validate_data ppf (decl_id, st) ; | ||||||||||||||||
| pp_set_size ppf (decl_id, st, DataOnly, false) | ||||||||||||||||
| | false -> pp_set_size ppf (decl_id, st, DataOnly, true) ) | ||||||||||||||||
| pp_assign_data ppf (decl_id, st, false) | ||||||||||||||||
| | false -> pp_assign_data ppf (decl_id, st, true) ) | ||||||||||||||||
| | Unsized _ -> () ) | ||||||||||||||||
| | _ -> pp_statement ppf s | ||||||||||||||||
| in | ||||||||||||||||
|
|
@@ -437,14 +437,26 @@ let pp_ctor ppf p = | |||||||||||||||
| let rec top_level_decls Stmt.Fixed.({pattern; _}) = | ||||||||||||||||
| match pattern with | ||||||||||||||||
| | Decl d when d.decl_id <> "pos__" -> | ||||||||||||||||
| [(d.decl_id, Type.to_unsized d.decl_type, UnsizedType.DataOnly)] | ||||||||||||||||
| [(d.decl_id, Type.to_unsized d.decl_type)] | ||||||||||||||||
| | SList stmts -> List.concat_map ~f:top_level_decls stmts | ||||||||||||||||
| | _ -> [] | ||||||||||||||||
|
|
||||||||||||||||
| (** Print the private data members of the model class *) | ||||||||||||||||
| let pp_model_private ppf {Program.prepare_data; _} = | ||||||||||||||||
| let data_decls = List.concat_map ~f:top_level_decls prepare_data in | ||||||||||||||||
| pf ppf "%a" (list ~sep:cut pp_decl) data_decls | ||||||||||||||||
| (*Filter out Any data that is not an Eigen matrix*) | ||||||||||||||||
| let get_eigen_map (name, ut) = | ||||||||||||||||
| match (UnsizedType.is_eigen_type ut, Transform_Mir.is_opencl_var name) with | ||||||||||||||||
| | true, false -> Some (name, ut) | ||||||||||||||||
| | false, _ -> None | ||||||||||||||||
| | true, _ -> None | ||||||||||||||||
|
||||||||||||||||
| match (UnsizedType.is_eigen_type ut, Transform_Mir.is_opencl_var name) with | |
| | true, false -> Some (name, ut) | |
| | false, _ -> None | |
| | true, _ -> None | |
| if UnsizedType.is_eigen_type ut && not (Transform_Mir.is_opencl_var name) | |
| then Some (name, ut) | |
| else None |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Totally cool with this change, though one thing I don't get yet. When is it better to use an if vs. a match? Is it just a style choice? I feel like when I see ifs I usually see us do a ; at the end I get kind of confused about execution stuff so I'm like "Oh idt I like that" but are ifs totally fine?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I dunno, I think it's just a style choice. Personally I prefer if/when for booleans and match for anything more complex. Also if can omit else so it's a bit more compact if you only need one branch.
If you want to keep match maybe merge | true, _ -> None | false, _ -> None into a single | _, _ -> None?
Actually, thinking more about this, it's more natural to use List.filter here instead of List.filter_map.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Makes sense! Switched this to use filter and if
| Original file line number | Diff line number | Diff line change | ||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -21,42 +21,129 @@ let rec contains_eigen (ut : UnsizedType.t) : bool = | |||||||||||||||||
| | UMatrix | URowVector | UVector -> true | ||||||||||||||||||
| | UInt | UReal | UMathLibraryFunction | UFun _ -> false | ||||||||||||||||||
|
|
||||||||||||||||||
| let pp_set_size ppf (decl_id, st, adtype, (needs_filled : bool)) = | ||||||||||||||||||
| (* TODO: generate optimal adtypes for expressions and declarations *) | ||||||||||||||||||
| let real_nan = | ||||||||||||||||||
| match adtype with | ||||||||||||||||||
| | UnsizedType.AutoDiffable -> "DUMMY_VAR__" | ||||||||||||||||||
| | DataOnly -> "std::numeric_limits<double>::quiet_NaN()" | ||||||||||||||||||
| (*Fill only needs to happen for containers | ||||||||||||||||||
| * Note: This should probably be moved into its own function as data | ||||||||||||||||||
| * does not need to be filled as we are promised user input data has the correct | ||||||||||||||||||
| * dimensions. Transformed data must be filled as incorrect slices could lead | ||||||||||||||||||
| * to elements of objects in transform data not being set by the user. | ||||||||||||||||||
| *) | ||||||||||||||||||
| let pp_filler ppf (decl_id, st, nan_type, needs_filled) = | ||||||||||||||||||
| match (needs_filled, contains_eigen (SizedType.to_unsized st)) with | ||||||||||||||||||
| | true, true -> | ||||||||||||||||||
| pf ppf "@[<hov 2>stan::math::fill(%s, %s);@]@," decl_id nan_type | ||||||||||||||||||
| | _ -> () | ||||||||||||||||||
|
|
||||||||||||||||||
| (*Pretty print a sized type*) | ||||||||||||||||||
| let pp_st ppf (st, adtype) = | ||||||||||||||||||
| pf ppf "%a" pp_unsizedtype_local (adtype, SizedType.to_unsized st) | ||||||||||||||||||
|
|
||||||||||||||||||
| let pp_ut ppf (ut, adtype) = pf ppf "%a" pp_unsizedtype_local (adtype, ut) | ||||||||||||||||||
|
|
||||||||||||||||||
| (*Get a string representing for the NaN type of the given type *) | ||||||||||||||||||
| let nan_type (st, adtype) = | ||||||||||||||||||
| match (adtype, st) with | ||||||||||||||||||
| | UnsizedType.AutoDiffable, _ -> "DUMMY_VAR__" | ||||||||||||||||||
| | DataOnly, SizedType.SInt -> "std::numeric_limits<int>::quiet_NaN()" | ||||||||||||||||||
|
||||||||||||||||||
| | DataOnly, _ -> "std::numeric_limits<double>::quiet_NaN()" | ||||||||||||||||||
|
|
||||||||||||||||||
| (*Pretty printer for the right hand side of expressions to initialize objects. | ||||||||||||||||||
| * For scalar types this sets the value to NaN and for containers initializes the memory. | ||||||||||||||||||
| *) | ||||||||||||||||||
| let rec pp_initialize ppf (st, adtype) = | ||||||||||||||||||
| let init_nan = nan_type (st, adtype) in | ||||||||||||||||||
| match st with | ||||||||||||||||||
| | SizedType.SInt -> pf ppf "std::numeric_limits<int>::min()" | ||||||||||||||||||
| | SReal -> pf ppf "%s" init_nan | ||||||||||||||||||
| | SVector d | SRowVector d -> pf ppf "%a(%a)" pp_st (st, adtype) pp_expr d | ||||||||||||||||||
| | SMatrix (d1, d2) -> | ||||||||||||||||||
| pf ppf "%a(%a, %a)" pp_st (st, adtype) pp_expr d1 pp_expr d2 | ||||||||||||||||||
| | SArray (t, d) -> | ||||||||||||||||||
| pf ppf "%a(%a, %a)" pp_st (st, adtype) pp_expr d pp_initialize (t, adtype) | ||||||||||||||||||
|
|
||||||||||||||||||
| (*Initialize an object of a given size.*) | ||||||||||||||||||
| let pp_assign_sized ppf (decl_id, st, adtype) = | ||||||||||||||||||
| let init_nan = nan_type (st, adtype) in | ||||||||||||||||||
| let pp_assign ppf (decl_id, st, adtype) = | ||||||||||||||||||
| pf ppf "@[<hov 2>%s = %a;@]@," decl_id pp_initialize (st, adtype) | ||||||||||||||||||
| in | ||||||||||||||||||
| let rec pp_size_ctor ppf st = | ||||||||||||||||||
| let pp_st ppf st = | ||||||||||||||||||
| pf ppf "%a" pp_unsizedtype_local (adtype, SizedType.to_unsized st) | ||||||||||||||||||
| in | ||||||||||||||||||
| pf ppf "@[%a%a@]@," pp_assign (decl_id, st, adtype) pp_filler | ||||||||||||||||||
| (decl_id, st, init_nan, true) | ||||||||||||||||||
|
|
||||||||||||||||||
| let%expect_test "set size mat array" = | ||||||||||||||||||
| let int = Expr.Helpers.int in | ||||||||||||||||||
| strf "@[<v>%a@]" pp_assign_sized | ||||||||||||||||||
| ("d", SArray (SArray (SMatrix (int 2, int 3), int 4), int 5), DataOnly) | ||||||||||||||||||
| |> print_endline ; | ||||||||||||||||||
| [%expect | ||||||||||||||||||
| {| | ||||||||||||||||||
| d = std::vector<std::vector<Eigen::Matrix<double, -1, -1>>>(5, std::vector<Eigen::Matrix<double, -1, -1>>(4, Eigen::Matrix<double, -1, -1>(2, 3))); | ||||||||||||||||||
| stan::math::fill(d, std::numeric_limits<double>::quiet_NaN()); |}] | ||||||||||||||||||
|
|
||||||||||||||||||
| (* Initialize Data and Transformed Data | ||||||||||||||||||
| * This function is used in the model's constructor to | ||||||||||||||||||
| * 1. Initialize memory for the data and transformed data | ||||||||||||||||||
| * 2. If an Eigen type, place that memory into the class's Map | ||||||||||||||||||
| * 3. Set the initial values of that data to NaN. | ||||||||||||||||||
| * @param ppf A pretty printer | ||||||||||||||||||
| * @param decl_id The name of the model class member | ||||||||||||||||||
| * @param st The type of the class member | ||||||||||||||||||
| *) | ||||||||||||||||||
| let pp_assign_data ppf | ||||||||||||||||||
| ((decl_id, st, needs_filled) : string * Expr.Typed.t SizedType.t * bool) = | ||||||||||||||||||
| let init_nan = nan_type (st, DataOnly) in | ||||||||||||||||||
| let pp_assign ppf (decl_id, st) = | ||||||||||||||||||
| match st with | ||||||||||||||||||
| | SizedType.SInt -> pf ppf "std::numeric_limits<int>::min()" | ||||||||||||||||||
| | SReal -> pf ppf "%s" real_nan | ||||||||||||||||||
| | SVector d | SRowVector d -> pf ppf "%a(%a)" pp_st st pp_expr d | ||||||||||||||||||
| | SMatrix (d1, d2) -> pf ppf "%a(%a, %a)" pp_st st pp_expr d1 pp_expr d2 | ||||||||||||||||||
| | SArray (t, d) -> pf ppf "%a(%a, %a)" pp_st st pp_expr d pp_size_ctor t | ||||||||||||||||||
| | SizedType.SVector _ | SRowVector _ | SMatrix _ -> | ||||||||||||||||||
| pf ppf "@[<hov 2>%s__ = %a;@]@," decl_id pp_initialize (st, DataOnly) | ||||||||||||||||||
| | SInt | SReal | SArray _ -> | ||||||||||||||||||
| pf ppf "@[<hov 2>%s = %a;@]@," decl_id pp_initialize (st, DataOnly) | ||||||||||||||||||
| in | ||||||||||||||||||
| let print_fill ppf st = | ||||||||||||||||||
| match (contains_eigen (SizedType.to_unsized st), needs_filled) with | ||||||||||||||||||
| | true, true -> pf ppf "stan::math::fill(%s, %s);" decl_id real_nan | ||||||||||||||||||
| | _, _ -> () | ||||||||||||||||||
| let pp_placement_new ppf (decl_id, st) = | ||||||||||||||||||
| match st with | ||||||||||||||||||
| | SizedType.SVector d | SRowVector d -> | ||||||||||||||||||
| pf ppf "@[<hov 2>new (&%s) Eigen::Map<%a>(%s__.data(), %a);@]@," | ||||||||||||||||||
| decl_id pp_st (st, DataOnly) decl_id pp_expr d | ||||||||||||||||||
| | SMatrix (d1, d2) -> | ||||||||||||||||||
| pf ppf "@[<hov 2>new (&%s) Eigen::Map<%a>(%s__.data(), %a, %a);@]@," | ||||||||||||||||||
| decl_id pp_st (st, DataOnly) decl_id pp_expr d1 pp_expr d2 | ||||||||||||||||||
| | _ -> () | ||||||||||||||||||
| in | ||||||||||||||||||
| pf ppf "@[<hov 0>%s = %a;@,%a @]@," decl_id pp_size_ctor st print_fill st | ||||||||||||||||||
| pf ppf "@[%a%a%a@]@," pp_assign (decl_id, st) pp_placement_new (decl_id, st) | ||||||||||||||||||
| pp_filler | ||||||||||||||||||
| (decl_id, st, init_nan, needs_filled) | ||||||||||||||||||
|
|
||||||||||||||||||
| let%expect_test "set size mat array" = | ||||||||||||||||||
| let%expect_test "set size map int array" = | ||||||||||||||||||
| let int = Expr.Helpers.int in | ||||||||||||||||||
| strf "@[<v>%a@]" pp_set_size | ||||||||||||||||||
| ( "d" | ||||||||||||||||||
| , SArray (SArray (SMatrix (int 2, int 3), int 4), int 5) | ||||||||||||||||||
| , DataOnly | ||||||||||||||||||
| , false ) | ||||||||||||||||||
| strf "@[<v>%a@]" pp_assign_data | ||||||||||||||||||
| ("darrmat", SArray (SArray (SInt, int 4), int 5), false) | ||||||||||||||||||
| |> print_endline ; | ||||||||||||||||||
| [%expect | ||||||||||||||||||
| {| | ||||||||||||||||||
| d = std::vector<std::vector<Eigen::Matrix<double, -1, -1>>>(5, std::vector<Eigen::Matrix<double, -1, -1>>(4, Eigen::Matrix<double, -1, -1>(2, 3))); |}] | ||||||||||||||||||
| darrmat = std::vector<std::vector<int>>(5, std::vector<int>(4, std::numeric_limits<int>::min())); |}] | ||||||||||||||||||
|
|
||||||||||||||||||
| let%expect_test "set size map mat array" = | ||||||||||||||||||
| let int = Expr.Helpers.int in | ||||||||||||||||||
| strf "@[<v>%a@]" pp_assign_data | ||||||||||||||||||
| ("darrmat", SArray (SArray (SMatrix (int 2, int 3), int 4), int 5), true) | ||||||||||||||||||
| |> print_endline ; | ||||||||||||||||||
| [%expect | ||||||||||||||||||
| {| | ||||||||||||||||||
| darrmat = std::vector<std::vector<Eigen::Matrix<double, -1, -1>>>(5, std::vector<Eigen::Matrix<double, -1, -1>>(4, Eigen::Matrix<double, -1, -1>(2, 3))); | ||||||||||||||||||
| stan::math::fill(darrmat, std::numeric_limits<double>::quiet_NaN()); |}] | ||||||||||||||||||
|
|
||||||||||||||||||
| let%expect_test "set size map mat" = | ||||||||||||||||||
| let int = Expr.Helpers.int in | ||||||||||||||||||
| strf "@[<v>%a@]" pp_assign_data ("dmat", SMatrix (int 2, int 3), false) | ||||||||||||||||||
| |> print_endline ; | ||||||||||||||||||
| [%expect | ||||||||||||||||||
| {| | ||||||||||||||||||
| dmat__ = Eigen::Matrix<double, -1, -1>(2, 3); | ||||||||||||||||||
| new (&dmat) Eigen::Map<Eigen::Matrix<double, -1, -1>>(dmat__.data(), 2, 3); |}] | ||||||||||||||||||
|
|
||||||||||||||||||
| let%expect_test "set size map int" = | ||||||||||||||||||
| strf "@[<v>%a@]" pp_assign_data ("dint", SInt, true) |> print_endline ; | ||||||||||||||||||
| [%expect {| | ||||||||||||||||||
| dint = std::numeric_limits<int>::min(); |}] | ||||||||||||||||||
|
|
||||||||||||||||||
| (** [pp_for_loop ppf (loopvar, lower, upper, pp_body, body)] tries to | ||||||||||||||||||
| pretty print a for-loop from lower to upper given some loopvar.*) | ||||||||||||||||||
|
|
@@ -70,7 +157,51 @@ let rec integer_el_type = function | |||||||||||||||||
| | SInt -> true | ||||||||||||||||||
| | SArray (st, _) -> integer_el_type st | ||||||||||||||||||
|
|
||||||||||||||||||
| let pp_decl ppf (vident, ut, adtype) = | ||||||||||||||||||
| (* Print the private members of the model class | ||||||||||||||||||
| * Accounting for types that can be moved to OpenCL. | ||||||||||||||||||
| * @param ppf A formatter | ||||||||||||||||||
| * @param vident name of the private member. | ||||||||||||||||||
| * @param ut The unsized type to print. | ||||||||||||||||||
| *) | ||||||||||||||||||
| let pp_data_decl ppf (vident, ut) = | ||||||||||||||||||
| let opencl_check = (Transform_Mir.is_opencl_var vident, ut) in | ||||||||||||||||||
| let pp_type = | ||||||||||||||||||
| match opencl_check with | ||||||||||||||||||
| | _, UnsizedType.(UInt | UReal) | false, _ -> pp_unsizedtype_local | ||||||||||||||||||
| | true, UArray UInt -> fun ppf _ -> pf ppf "matrix_cl<int>" | ||||||||||||||||||
| | true, _ -> fun ppf _ -> pf ppf "matrix_cl<double>" | ||||||||||||||||||
| in | ||||||||||||||||||
| match (opencl_check, ut) with | ||||||||||||||||||
| | (false, _), ut -> ( | ||||||||||||||||||
| match ut with | ||||||||||||||||||
| | UnsizedType.URowVector | UVector | UMatrix -> | ||||||||||||||||||
| pf ppf "%a %s__;" pp_type (DataOnly, ut) vident | ||||||||||||||||||
| | _ -> pf ppf "%a %s;" pp_type (DataOnly, ut) vident ) | ||||||||||||||||||
| | (true, _), _ -> pf ppf "%a %s;" pp_type (DataOnly, ut) vident | ||||||||||||||||||
|
|
||||||||||||||||||
| (*Create strings representing maps of Eigen types*) | ||||||||||||||||||
| let pp_map_decl ppf (vident, ut) = | ||||||||||||||||||
| let scalar = local_scalar ut DataOnly in | ||||||||||||||||||
| match ut with | ||||||||||||||||||
| | UnsizedType.UInt | UReal -> () | ||||||||||||||||||
| | UMatrix -> | ||||||||||||||||||
| pf ppf "Eigen::Map<Eigen::Matrix<%s, -1, -1>> %s{nullptr, 0, 0};" scalar | ||||||||||||||||||
| vident | ||||||||||||||||||
| | URowVector -> | ||||||||||||||||||
| pf ppf "Eigen::Map<Eigen::Matrix<%s, 1, -1>> %s{nullptr, 0};" scalar | ||||||||||||||||||
| vident | ||||||||||||||||||
| | UVector -> | ||||||||||||||||||
| pf ppf "Eigen::Map<Eigen::Matrix<%s, -1, 1>> %s{nullptr, 0};" scalar | ||||||||||||||||||
| vident | ||||||||||||||||||
| | x -> | ||||||||||||||||||
| let bad_type = Fmt.strf "%a" pp_ut (x, DataOnly) in | ||||||||||||||||||
| raise_s | ||||||||||||||||||
| [%message | ||||||||||||||||||
| "Error during Map data construction for " vident " of type " bad_type | ||||||||||||||||||
|
||||||||||||||||||
| let bad_type = Fmt.strf "%a" pp_ut (x, DataOnly) in | |
| raise_s | |
| [%message | |
| "Error during Map data construction for " vident " of type " bad_type | |
| raise_s | |
| [%message | |
| "Error during Map data construction for " vident " of type " | |
| (x : UnsizedType.t) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why does this need adjustment?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is because with the default parallel setting the instances run out of RAM, with -j7 the tests pass but they run for 3 hours instead of 20-30mins! We need to get to the bottom of this.
Does this effect "typical" models or is it just a factor with these huge test models ala this one.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah I think this is just with the huge models. Though @rok-cesnovar I think you and I have both tried these locally without issue. It could be the compiler we use for running these tests? I'm really not sure. My guess as to what's happening is that we are now generating signatures for both
Eigen::Map<>types andEigen::Matrix<>typesThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I havent explored it for this PR yet, but have it on my list for the next few days.