Edit on GitHub

Modeling Language Implementation

Parsing @gen functions

Gen's built-in modeling languages are designed to preserve Julia's syntax as far as possible, apart from the Tilde syntax for calling generative functions, and the restrictions imposed on the Static Modeling Language. In order to preserve that syntax, including the use of non-Gen macros within @gen functions, we relegate as much of the parsing of @gen functions as possible to Julia's macro-expander and parser.

In particular, we adopt an implementation strategy that enforces a separation between the surface syntax associated with Gen-specific macros (i.e., @trace and @param) and their corresponding implementations, which differ across the Dynamic Modeling Language (DML) and the Static Modeling Language (SML). We do this by introducing the custom expressions Expr(:gentrace, call, addr) and Expr(:genparam, name, type), which serve as intermediate representations in the macro-expanded abstract syntax tree.

Each modeling language can then handle these custom expressions in their own manner, either by parsing them to nodes in the Static Computation Graph (for the SML), or by substituting them with their implementations (for the DML). This effectively allows the SML and DML to have separate implementations of @trace and @param.

For clarity, below is a procedural description of how the @gen macro processes Julia function syntax:

  1. macroexpand the entire function body with respect to the calling module. This expands any (properly-scoped) @trace calls to Expr(:gentrace, ...) expressions, and any (properly-scoped) @param calls to Expr(:genparam, ...) expressions, while also expanding non-Gen macros.
  2. Desugar any tilde expressions x ~ gen_fn(), including those that may have been generated by macros, to Expr(:gentrace, ...) expressions.
  3. Pass the macro-expanded and de-sugared function body on to make_static_gen_function or make_dynamic_gen_function accordingly.
  4. For static @gen functions, match :gentrace expressions when adding address nodes to the static computation graph, and match :genparam expressions when adding parameter nodes to the static computation graph. A StaticIRGenerativeFunction is then compiled from the static computation graph.
  5. For dynamic @gen functions, rewrite any :gentrace expression with its implementation dynamic_trace_impl, and rewrite any :genparam expression with its implementation dynamic_param_impl. The rewritten syntax tree is then evaluated as a standard Julia function, which serves as the implementation of the constructed DynamicDSLFunction.