inline inferred blocks

Author: MikeInnes

Manifest.toml

@@ -9,9 +9,11 @@ uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 
 [[IRTools]]
 deps = ["InteractiveUtils", "MacroTools", "Test"]
-git-tree-sha1 = "90ee39f9beaaa186e4968417ea2b8ed5673c91c0"
+git-tree-sha1 = "5b2ef472bd7001923ea0d703583865ade78c80a3"
+repo-rev = "master"
+repo-url = "https://github.com/MikeInnes/IRTools.jl.git"
 uuid = "7869d1d1-7146-5819-86e3-90919afe41df"
-version = "0.3.3"
+version = "0.4.0"
 
 [[InteractiveUtils]]
 deps = ["Markdown"]

src/infer.jl

@@ -86,6 +86,7 @@ mutable struct Frame
   edges::Vector{Any}
   stmts::Vector{Vector{Variable}}
   rettype::AType
+  outer::Bool
 end
 
 getblock(fr::Frame, b, follow) =
@@ -95,7 +96,7 @@ getblock(fr::Frame, b, follow) =
 # TODO clear up inlined edges
 uninline!(fr::Frame, b, follow) = deleteat!(fr.inlined[b], follow+1:length(fr.inlined[b]))
 
-Frame(ir::IR) = Frame(ir, [IR[] for _ = 1:length(blocks(ir))], [], keys.(blocks(ir)), Union{})
+Frame(ir::IR; outer = false) = Frame(ir, [IR[] for _ = 1:length(blocks(ir))], [], keys.(blocks(ir)), Union{}, outer)
 
 function frame(ir::IR, args...)
   prepare_ir!(ir)
@@ -200,7 +201,7 @@ function step!(inf::Inference)
       error("Unrecognised expression $(st.expr)")
     end
   elseif (brs = openbranches(block); length(brs) == 1 && !isreturn(brs[1])
-          && !(brs[1].block == length(frame.ir.blocks)))
+          && !(frame.outer && brs[1].block == length(frame.ir.blocks)))
     inferbranch!(inf, frame, b, f, brs[1])
   else
     for br in brs
@@ -249,3 +250,19 @@ function return_type(ir::IR, args...)
   infer!(inf)
   return fr.rettype
 end
+
+function inlineable(ir::IR, inf::Inference)
+  for (v, st) in ir
+    Ts = exprtype.((ir,), st.expr.args)
+    haskey(inf.frames, Ts) && return (v, Ts)
+  end
+end
+
+function inlineall(ir::IR, inf::Inference)
+  while (next = inlineable(ir, inf)) != nothing
+    v, Ts = next
+    subir = inf.frames[Ts].ir
+    ir = inline(ir, v, subir)
+  end
+  return ir
+end

src/trace.jl

@@ -133,7 +133,11 @@ function abstract!(tr, bl, env)
   for i = 1:length(args)
     argtypes(ir)[i] = exprtype(tr.ir, rename(env, args[i]))
   end
-  infer!(Inference(Frame(ir), tr.primitives))
+  inf = Inference(Frame(ir, outer = true), tr.primitives)
+  infer!(inf)
+  # TODO weird that we need to expand here; perhaps inlining is impicitly pruning
+  # somehow.
+  ir = inlineall(ir, inf) |> expand!
   inline!(tr.ir, ir, rename.((env,), args))
   for (k, v) in zip(arguments(block(bl.ir, after)), arguments(blocks(tr.ir)[end]))
     env[k] = v