Module Analysis_and_optimization.Memory_patterns

type demotion = int * Middle.Mem_pattern.t * string
val compare_demotion : demotion -> demotion -> int
val demotion_reasons : demotion list Core.ref
val get_warnings : unit -> string list
val user_warning_op : Middle.Mem_pattern.t -> int -> string -> string -> unit
val concat_set_str : string Core.Set.Poly.t -> string

Return a Var expression of the name for each type containing an eigen matrix

val query_var_eigen_names : Middle.Expr.Typed.t -> string Core.Set.Poly.t

Return a set of all types containing autodiffable Eigen matrices in an expression.

val is_nonzero_subset : set:('a, 'b) Core.Set.t -> subset:('a, 'b) Core.Set.t -> bool

Check whether one set is a nonzero subset of another set.

val count_single_idx : int -> Middle.Expr.Typed.t Middle.Index.t -> int

Check an Index to count how many times we see a single index.

  • parameter acc

    An accumulator from previous folds of multiple expressions.

  • parameter idx

    An Index to match. For Single types this adds 1 to the acc. For Upfrom and MultiIndex types we check the inner expression for a Single index. All and Between cannot be Single cell access and so pass acc along.

val is_uni_eigen_loop_indexing : bool -> Middle.UnsizedType.t -> Middle.Expr.Typed.t Middle.Index.t list -> bool

Find indices on Matrix and Vector types that perform single cell access. Returns true if it finds a vector, row vector, matrix, or matrix with single cell access as well as an array of any of the above that is accessing the inner matrix types cell.

  • parameter ut

    An UnsizedType to match against.

  • parameter index

    This list is checked for Single cell access either at the top level or within the Index types of the list.

val query_stan_math_mem_pattern_support : string -> Middle.UnsizedType.argumentlist -> bool
val is_fun_soa_supported : string -> Middle.Expr.Typed.t list -> bool

Validate whether a function can support SoA matrices

val query_initial_demotable_expr : bool -> int -> acc:string Core.Set.Poly.t -> Middle.Expr.Typed.t -> string Core.Set.Poly.t

Query to find the initial set of objects that cannot be SoA. This is mostly recursing over expressions, with the exceptions being functions and indexing expressions. For the logic on functions see the docs for query_initial_demotable_funs.

  • parameter in_loop

    a boolean to signify if the expression exists inside of a loop. If so, the names of matrix and vector like objects will be returned if the matrix or vector is accessed by single cell indexing.

val query_initial_demotable_funs : bool -> int -> string Core.Set.Poly.t -> Middle.Expr.Typed.Meta.t Middle.Expr.t Middle.Fun_kind.t -> Middle.Expr.Typed.t list -> string Core.Set.Poly.t

Query a function to detect if it or any of its used expression's objects or expressions should be demoted to AoS.

The logic here demotes the expressions in a function to AoS if the function's inner expression returns has a meta type containing a matrix and either of :

  1. The function is user defined and the UDFs inputs are matrices.
  2. The Stan math function cannot support AoS
  • parameter in_loop

    A boolean to specify the logic of indexing expressions. See query_initial_demotable_expr for an explanation of the logic.

  • parameter kind

    The function type, for StanLib functions we check if the function supports SoA and for UserDefined functions we always fail and return back all of the names of the objects passed in expressions to the UDF. exprs The expression list passed to the functions.

val extract_nonderived_admatrix_types : Middle.Expr.Typed.t -> (Middle.UnsizedType.autodifftype * Middle.UnsizedType.t) list

Recurse through subexpressions and return a list of Unsized types. Recursion continues until

  1. A non-autodiffable type is found
  2. An autodiffable scalar is found
  3. A `Var` type is found that is an autodiffable matrix

Recurse through functions to find nonderived ad matrix types. Special cases for StanLib functions are for

  • `check_matching_dims`: compiler function that has no effect on optimization
  • `rep_*vector` These are templated in the C++ to cast up to `Var<Matrix>` types
  • `rep_matrix`. When it's only a scalar being propogated an math library overload can upcast to `Var<Matrix>`
val contains_at_least_one_ad_matrix_or_all_data : Middle.UnsizedType.argumentlist -> bool

Checks if a list of types contains at least on ad matrix or if everything is derived from data

val query_initial_demotable_stmt : bool -> string Core.Set.Poly.t -> (Middle.Expr.Typed.Meta.t, Middle.Stmt.Located.Meta.t) Middle.Stmt.t -> string Core.Set.Poly.t

Query to find the initial set of objects in statements that cannot be SoA. This is mostly recursive over expressions and statements, with the exception of functions and Assignments. For assignments: We demote the LHS variable if any of the following are true:

  1. A single cell of the LHS is being assigned within a loop.
  2. The top level expression on the RHS is a combination of only data matrices and scalar types. Operations on data matrix and scalar values in Stan math will return a AoS matrix. We currently have no way to tell Stan math to return a SoA matrix.
  3. None of the RHS's functions are able to accept SoA matrices and the rhs is not an internal compiler function.

We demote RHS variables if any of the following are true:

  1. The LHS variable has previously or through this iteration been marked AoS.
  2. The LHS is a tuple projection

For functions see the documentation for query_initial_demotable_funs for the logic on demotion rules.

  • parameter in_loop

    A boolean to specify the logic of indexing expressions. See query_initial_demotable_expr for an explanation of the logic.

val query_demotable_stmt : string Core.Set.Poly.t -> Middle.Stmt.Located.Non_recursive.t -> string Core.Set.Poly.t

Look through a statement to see whether the objects used in it need to be modified from SoA to AoS. Returns the set of object names that need demoted in a statement, if any. This function looks at Assignment statements, and returns back the set of top level object names given:

  1. If the name of the lhs assignee is in the aos_exits, all the names of the expressions with a type containing a matrix are returned.
  2. If the names of the rhs objects containing matrix types are in the subset of aos_exits.
  • parameter aos_exits

    A set of variables that can be demoted.

  • parameter pattern

    The Stmt pattern to query.

val modify_kind : ?force_demotion:bool -> string Core.Set.Poly.t -> Middle.Expr.Typed.t Middle.Fun_kind.t -> Middle.Expr.Typed.t list -> Middle.Expr.Typed.t Middle.Fun_kind.t * Middle.Expr.Typed.t list

Modify a function and it's subexpressions from SoA <-> AoS and vice versa. This performs demotion for sub expressions recursively. The top level expression and it's sub expressions are demoted to SoA if

  1. The names of the variables in the subexpressions returning objects holding matrices are all in the modifiable set.
  2. The function does not support SoA 3. The force argument is true
  • parameter force_demotion

    If true, forces an expression and it's sub-expressions to be AoS.

  • parameter modifiable_set

    The set of names that are either demotable to AoS or promotable to SoA.

  • parameter kind

    A Fun_kind.t

  • parameter exprs

    A list of expressions going into the function. *

val modify_expr_pattern : ?force_demotion:bool -> string Core.Set.Poly.t -> Middle.Expr.Typed.t Middle.Expr.Pattern.t -> Middle.Expr.Typed.t Middle.Expr.Pattern.t

Modify an expression and it's subexpressions from SoA <-> AoS and vice versa. The only real paths in the below is on the functions and ternary expressions. The logic for functions is defined in modify_kind. TernaryIf is forcefully demoted to AoS if the type of the expression contains a matrix.

  • parameter force_demotion

    If true, forces an expression and it's sub-expressions to be AoS.

  • parameter modifiable_set

    The name of the variables whose associated expressions we want to modify.

  • parameter pattern

    The expression to modify.

val modify_expr : ?force_demotion:bool -> string Core.Set.Poly.t -> Middle.Expr.Typed.t -> Middle.Expr.Typed.t

Given a Set of strings containing the names of objects that can be modified from AoS <-> SoA and vice versa, modify them within the expression.

  • parameter mem_pattern

    The memory pattern to change expressions to.

  • parameter modifiable_set

    The name of the variables whose associated expressions we want to modify.

  • parameter expr

    the expression to modify.

Modify statement patterns in the MIR from AoS <-> SoA and vice versa For Decl and Assignment's reading in parameters, we demote to AoS if the decl_id (or assign name) is in the modifiable set and otherwise promote the statement to SoA. For general Assignment statements, we check if the assignee is in the demotable set. If so, we force demotion of all of the rhs expressions. All other statements recurse over their statements and expressions. *

  • parameter pattern

    The statement pattern to modify

  • parameter modifiable_set

    The name of the variable we are searching for.

val modify_stmt : Middle.Stmt.Located.t -> string Core.Set.Poly.t -> Middle.Stmt.Located.t

Modify statement patterns in the MIR from AoS <-> SoA and vice versa

  • parameter mem_pattern

    A mem_pattern to modify expressions to. For the given memory pattern, this modifies statement patterns and expressions to it.

  • parameter stmt

    The statement to modify.

  • parameter modifiable_set

    The name of the variable we are searching for.

val collect_mem_pattern_variables : Middle.Stmt.Located.t Core.List.t -> (string * Middle.Expr.Typed.Meta.t Middle.Expr.t Middle.SizedType.t) list
val pp_mem_patterns : Stdlib.Format.formatter -> Middle.Program.Typed.t -> unit