support finding dependencies of a variable

[Roger-luo]

I think we discuss this a long time ago in slack @MikeInnes , by dispatching pullbacks or other contextual calls to only variables dependent on the result, it will at least eliminate some annoying errors for functions that does not involve global references etc.

but I'm not sure this is the right way to implement it. I guess you might have some comments. especially, I'm not sure this is right for branches.

[MikeInnes]

This seems like a reasonable approach; it won't yet work for branches because ir[v] doesn't exist for block arguments. You instead have to enumerate each branch to the block and include all passed argument variables as dependencies.

It may be better to rework this to use a dataflow-based approach, which is basically like doing type inference. You walk forward over all variables keeping track of the set of dependencies (which is easy to calculate as the union of dependencies of variables in the Expr). When you get to the end of the block, you check each branch to see if its dependencies have changed, and if so add the target block to a work queue. That will correctly handle things like loops where dependencies can be complex.

It might be useful to have a little tutorial on doing data flow analyses in the docs, since this is a good approach to almost every non-trivial task on IR.

[Roger-luo]

Thanks! @MikeInnes by dataflow-based approach do you mean something like Mjolnir? Or should I PR this feature into Mjolnir instead of IRTools?

[MikeInnes]

No, "data flow analysis" is just a general approach to writing algorithms on IR like this; see for example https://en.wikipedia.org/wiki/Data-flow_analysis. The clearest example in IRTools is probably this function which turns all implicit inter-block variable usage into explicit block arguments. So I'm suggesting that you rework the structure of your algorithm to do a dataflow analysis in the way that I described above.

[Roger-luo]

@MikeInnes I just rework this, should be ready for a review.

[Roger-luo]

err... Not sure why this fails on 1.0

[Roger-luo]

I think this is because the IR on 1.0 is different from later Julia versions due to different implementation of @info. So I mark that as broken.

@MikeInnes any chance to review this?

[MikeInnes]

This is looking nice, thanks! However, I think it'd be useful to consider a test case with a loop, like

function pow(x, n)
  r = 1
  while n > 0
    n -= 1
    r *= x
  end
  return r
end

In particular dependencies currently fails on this because it always follows branches, but instead it needs to determine whether a branch actually needs to be followed (ie the dependencies of an input actually changed). I suspect that find_dependency_path would also fail if there's a circular dependency.

The current setup is a little weird because what dependencies calculates – immediate dependencies – is fairly trivial. It doesn't require a dataflow analysis and can be done efficiently for one variable at a time. So instead I think it'd make sense for dependencies to calculate all dependencies (like find_dependency_path), for all variables at once. You can do this by defining the set of dependencies as the set of immediate uses + their dependencies. When you update basic block arguments you can check if they have any new dependencies, and redo the block if there are.

[MikeInnes]

For Expr you could just iterate over the arguments, rather than the heavier prewalk.

[MikeInnes]

I would probably just avoid @info since its implementation is liable to change at any time; perhaps just write down similar code by hand?

[Roger-luo]

@MikeInnes bump