Constant nodes in Static IR

[alex-lew]

Consider the following SML function:

@gen (static) function f()
  x ~ normal(0, 1)
  y ~ normal(0, 1)
  z ~ normal(0, 1)
end

Each of the six constants (the three zeros and the three ones) are treated as "arbitrary Julia code": Gen creates functions () -> 0, () -> 1, and in the trace type, creates six fields to cache the three zeros and the three ones. Furthermore, these fields are not annotated with types.

As such, the following equivalent function yields more performant GFI operations, by up to 5x in my experiments:

@gen (static) function f()
  zero_var::Float64 = 0.0
  one_var::Float64 = 1.0
  x ~ normal(zero_var, one_var)
  y ~ normal(zero_var, one_var)
  z ~ normal(zero_var, one_var)
end

It seems to me that the SML parser could automatically recognize when a Julia expression is a constant (has no dependence on args or upstream vars), and include a special 'constant node' in the DAG. Then the code generators could be updated to directly insert the constant values into generated code, rather than looking them up in the trace. (Ideally the constants could be evaluated at compile-time, and their values inserted, rather than stored as quoted code to insert.)