support finding dependencies of a variable

src/IRTools.jl

@@ -50,7 +50,7 @@ let exports = :[
       definitions, usages, dominators, domtree, domorder, domorder!, renumber,
       merge_returns!, expand!, prune!, ssa!, inlineable!, log!, pis!, func, evalir,
       Simple, Loop, Multiple, reloop, stackify, functional, cond, WorkQueue,
-      Graph, liveness, interference, colouring,
+      Graph, liveness, interference, colouring, dependencies, find_dependency_path,
       # Reflection, Dynamo
       Meta, Lambda, meta, dynamo, transform, refresh, recurse!, self,
       varargs!, slots!,

src/passes/passes.jl

@@ -41,6 +41,73 @@ function usages(b::Block)
   return uses
 end
 
+usages(st::Statement) = usages(st.expr)
+usages(ex) = Set{Variable}()
+
+function usages(ex::Expr)
+  uses = Set{Variable}()
+  prewalk(ex) do x
+    x isa Variable && push!(uses, x)
+    return x
+  end
+  return uses
+end
+
+"""
+  dependencies(ir::IR)
+
+Return the list of direct dependencies for each variable.
+"""
+function dependencies(ir::IR)
+  worklist = [block(ir, 1)]
+  deps = Dict()
+  while !isempty(worklist)
+    b = pop!(worklist)
+    for (v, st) in b
+      set = get!(deps, v, Set{Variable}())
+      union!(set, usages(st))
+    end
+
+    brs = BasicBlock(b).branches
+    jump_next_block = true
+    for br in brs
+      jump_next_block = jump_next_block && (br.condition !== nothing)
+      if br.block > 0
+        next_block = block(ir, br.block)
+        push!(worklist, next_block)
+        if !isempty(br.args)
+          for (x, y) in zip(arguments(next_block), br.args)
+            set = get!(deps, x, Set{Variable}())
+            push!(set, y)
+          end
+        end
+      end
+    end
+    jump_next_block && push!(worklist, block(ir, b.block+1))
+  end
+  return deps
+end
+
+function _find_dependency_path(deps::Dict, x::Variable)
+  path = Set{Variable}()
+  stack = Variable[x]
+  while !isempty(stack)
+    curr = pop!(stack)
+    if haskey(deps, curr)
+      for v in deps[curr]
+        push!(path, v)
+        push!(stack, v)
+      end
+    end
+  end
+  return path
+end
+
+function find_dependency_path(ir::IR, x::Variable)
+  deps = dependencies(ir)
+  return _find_dependency_path(deps, x)
+end
+
 function usecounts(ir::IR)
   counts = Dict{Variable,Int}()
   prewalk(ir) do x

test/analysis.jl

@@ -1,5 +1,5 @@
 using IRTools, Test
-using IRTools: CFG, dominators, domtree
+using IRTools: CFG, dominators, domtree, dependencies, find_dependency_path, var
 
 relu(x) = (y = x > 0 ? x : 0)
 ir = @code_ir relu(1)
@@ -10,3 +10,48 @@ ir = @code_ir relu(1)
 @test domtree(CFG(ir)) == (1 => [2 => [], 3 => [], 4 => []])
 
 @test domtree(CFG(ir)', entry = 4) == (4 => [1 => [], 2 => [], 3 => []])
+
+function f(x)
+  x = sin(x)
+  y = cos(x)
+
+  if x > 1
+    x = cos(x) + 1
+  else
+    x = y + 1
+  end
+  return x
+end
+
+ir = @code_ir f(1.0)
+
+@test dependencies(ir) == Dict(
+  var(9) => Set([var(8), var(7)]),
+  var(8) => Set([var(4)]),
+  var(7) => Set([var(6)]),
+  var(3) => Set([var(2)]),
+  var(5) => Set([var(3)]),
+  var(6) => Set([var(3)]),
+  var(4) => Set([var(3)]),
+)
+
+@test find_dependency_path(ir, var(9)) == Set(var.([2, 8, 7, 3, 4, 6]))
+
+# function f(x)
+#   x = sin(x)
+#   y = cos(x)
+
+#   @info "warning" x
+
+#   if x > 1
+#     x = cos(x) + 1
+#   else
+#     x = y + 1
+#   end
+#   return x
+# end
+
+# ir = @code_ir f(1.0)
+
+# # NOTE: this test is broken due to different IR generated by 1.0 and 1.2+
+# @test_broken find_dependency_path(ir, var(48)) == Set(var.([47, 46, 42, 45, 41, 3, 4, 2]))