Merge pull request #382 from probcomp/20200216_ztangent_irprettyprint

Pretty print for static IR

Author: ztangent

src/static_ir/print_ir.jl

@@ -0,0 +1,36 @@
+Base.show(io::IO, ::MIME"text/plain", ir::StaticIR) =
+    print_ir(io, ir)
+
+function print_ir(io::IO, ir::StaticIR)
+    println(io, "== Static IR ==")
+    args = join((string(arg.name) for arg in ir.arg_nodes), ", ")
+    println(io, "Arguments: ($args)")
+    for node in ir.nodes
+        node in ir.arg_nodes && continue
+        print(io, "  "); print_ir(io, node); println(io)
+    end
+    print(io, "  return $(ir.return_node.name)")
+end
+
+function print_ir(io::IO, node::TrainableParameterNode)
+    print(io, "@param $(node.name)::$(node.typ)")
+end
+
+function print_ir(io::IO, node::JuliaNode)
+    inputs = join((string(i.name) for i in node.inputs), ", ")
+    print(io, "$(node.name) = $(node.fn)($inputs)")
+end
+
+function print_ir(io::IO, node::GenerativeFunctionCallNode)
+    inputs = join((string(i.name) for i in node.inputs), ", ")
+    gen_fn_name = ir_name(node.generative_function)
+    print(io, "$(node.name) = @trace($(gen_fn_name)($inputs), :$(node.addr))")
+end
+
+function print_ir(io::IO, node::RandomChoiceNode)
+    inputs = join((string(i.name) for i in node.inputs), ", ")
+    print(io, "$(node.name) = @trace($(node.dist)($inputs), :$(node.addr))")
+end
+
+ir_name(fn::GenerativeFunction) = nameof(typeof(fn))
+ir_name(fn::DynamicDSLFunction) = nameof(fn.julia_function)

src/static_ir/static_ir.jl

@@ -52,6 +52,7 @@ function generate_generative_function(ir::StaticIR, name::Symbol, options::Stati
             params::Dict{Symbol,Any}
         end
         (gen_fn::$gen_fn_type_name)(args...) = $(GlobalRef(Gen, :propose))(gen_fn, args)[3]
+        $(GlobalRef(Gen, :get_ir))(::$gen_fn_type_name) = $(QuoteNode(ir))
         $(GlobalRef(Gen, :get_ir))(::Type{$gen_fn_type_name}) = $(QuoteNode(ir))
         $(GlobalRef(Gen, :get_trace_type))(::Type{$gen_fn_type_name}) = $trace_struct_name
         $(GlobalRef(Gen, :has_argument_grads))(::$gen_fn_type_name) = $(QuoteNode(has_argument_grads))
@@ -63,6 +64,7 @@ function generate_generative_function(ir::StaticIR, name::Symbol, options::Stati
     Expr(:block, trace_defns, gen_fn_defn, Expr(:call, gen_fn_type_name, :(Dict{Symbol,Any}()), :(Dict{Symbol,Any}())))
 end
 
+include("print_ir.jl")
 include("render_ir.jl")
 
 ###########################

test/static_ir/static_ir.jl

@@ -22,6 +22,7 @@ z = add_julia_node!(builder, (u, v) -> u + v, inputs=[u, v], name=:z)
 set_return_node!(builder, z)
 ir = build_ir(builder)
 bar = eval(generate_generative_function(ir, :bar, track_diffs=false, cache_julia_nodes=false))
+@test occursin("== Static IR ==", repr("text/plain", ir))
 
 #@gen (static, nojuliacache) function foo(a, b)
     #@param theta::Float64
@@ -45,6 +46,7 @@ w = add_julia_node!(builder, (z, a, theta) -> z + 1 + a + theta, inputs=[z, a, t
 set_return_node!(builder, w)
 ir = build_ir(builder)
 foo = eval(generate_generative_function(ir, :foo, track_diffs=false, cache_julia_nodes=false))
+@test occursin("== Static IR ==", repr("text/plain", ir))
 
 theta_val = rand()
 set_param!(foo, :theta, theta_val)
@@ -58,6 +60,7 @@ one = add_constant_node!(builder, 2)
 set_return_node!(builder, one)
 ir = build_ir(builder)
 const_fn = eval(generate_generative_function(ir, :const_fn, track_diffs=false, cache_julia_nodes=false))
+@test occursin("== Static IR ==", repr("text/plain", ir))
 
 Gen.load_generated_functions()