Improve performance of `ir_traverse._call_with_optional_args`

[EricRahm]

Overview

As a follow-up to the analysis of #118 I found that the inspect.getfullargspec in _call_with_optional_args is taking about 10% of runtime:

A proof-of-concept fully removed this overhead.

Background

This call is invoked as part of the fast_traverse_ir_top_down implementation. The action and incidental_action operations are all run through this helper function which is invoked about 100,000 times in my test. Currently inspect.getfullargspec is called every time. Additionally the argument extraction does a fair amount of list creation and iteration that could also be avoided if memoized:

def _call_with_optional_args(function, positional_arg, keyword_args):
  """Calls function with whatever keyword_args it will accept."""
  argspec = inspect.getfullargspec(function)
  if argspec.varkw:
    # If the function accepts a kwargs parameter, then it will accept all
    # arguments.
    # Note: this isn't technically true if one of the keyword arguments has the
    # same name as one of the positional arguments.
    return function(positional_arg, **keyword_args)
  else:
    ok_arguments = {}
    for name in keyword_args:
      if name in argspec.args[1:] or name in argspec.kwonlyargs:
        ok_arguments[name] = keyword_args[name]
    for name in argspec.args[1:len(argspec.args) - len(argspec.defaults or [])]:
      assert name in ok_arguments, (
          "Attempting to call '{}'; missing '{}' (have '{!r}')".format(
              function.__name__, name, list(keyword_args.keys())))
    return function(positional_arg, **ok_arguments)

Possible Fix

Memoize the extracted lists of valid and required args from inspect.getfullargspecfor a given function. Something like:

@simple_memoizer.memoize
def get_function_args(function):
  argspec = inspect.getfullargspec(function)
  return {
     "valid": argspec.args[1:] or name in argspec.kwonlyargs
     "required": argspec.args[1:len(argspec.args) - len(argspec.defaults or [])]
  }

Drawbacks

We currently use some lambda functions in addition to static methods for actions. We'd need to convert those to static methods, otherwise we'd end up with a memoization entry for each dynamic lambda instance.

[EricRahm]

120 fixed this.