[Relay][RFC] Pass Manager

After several discussions with @jroesch and some people tagged below, we come up with the following RFC for Relay Pass Manager design.

Motivation

Relay features a series of optimization passes which improve performance metrics of models such as mean inference, memory footprint, power consumption, or for specific devices. There are a suite of standard optimizations as well as machine learning specific optimizations including constant folding, dead code elimination, operator layout alteration, and operator fusion. Each of these passes are structured as Relay-to-Relay transformations on the abstract syntax tree (AST) using the analysis result collected during and/or before traversal. Many existing production compilers, such as GCC and LLVM, employ pass managers to effectively manage the execution of passes. Initially managing passes is straightforward as the number of passes is small, but mature compilers will contain hundreds of individual passes. Often external users will want to have custom passes correctly scheduled without having to modify a single handcrafted pass order.

The responsibilities of pass managers usually involve:

• Organizing the execution order of optimization passes though not necessarily providing an optimal ordering.
• Collect the and maintain the required analysis information during optimization.
• Simplifying the implementation of new passes for compiler developers, and enabling extensibility for users of the stack.

One might argue that the current number of Relay passes do not necessitate the effort to design such a pass manager. However, the management of the passes, e.g. applying them in a convenient manner, and finding a good order becomes intractable as the number of them increases. Relay is only a few months old and we are rapidly moving towards dozens of passes.

We would like to write the code to configure and select these passes once, and deduplicate the definition of multiple optimize functions. Currently similar code appears in relay.build, the interpreter, the virtual machine, the ahead of time compiler and so on.

In addition, we do not have a standard signature for all passes. For example we currently pass the Module object to only some optimization passes but not to all. One example that does not receive the relay.Module is operator fusion. This limitation makes it impossible to execute multiple function programs in some cases. As we build more complex models we will need to use functions and globals. Separating repetitive code such as traversing the module, caching, etc would alleviate the time and effort needed to write passes. Of course this introduces extra engineering effort as well as introduce bugs. In our opinion factoring the code into a single interface will bring greater to stability to all passes in the long run.

This RFC proposes a design for Relay’s pass manager, enabling better programmatic orchestration of optimizations. The design is largely inspired by the hierarchical pass manager used in LLVM. The initial implementation lives in https://github.com/zhiics/tvm/tree/pass_manager and is a work in progress.

Proposed Design

Our main design goal is ease of extension for users. We want it to be possible for users to quickly add new passes without loss of backward compatibility. For example our design allows users to implement Relay passes in Python and pass them to C++ by making use the PackedFunc system TVM provides.

Backend

We provide a Pass object to contain the basic information needed by a pass. name is the pass name and opt_level indicates at which optimization level the pass will be enabled. For example, it could be a function pass, a module pass, or a sequential pass.

class Pass;

class PassNode : RelayNode {
  std::string name;
  int opt_level;
  
  virtual std::vector<std::string> Required() const = 0;
  virtual Module operator()(const Module& mod) = 0;
  virtual Module Run(const Module& mod) const;

  void SetContext(const PassContext& pass_ctx) { pass_ctx_ = pass_ctx; } 
 protected:
   PassContext pass_ctx_;
};

PassContext is designed to contain useful information for an optimization pass. For example, it could contain error reporting system so optimization authors can provide diagnostics about why an optimization fails.

class PassContext;
class PassContextNode : RelayNode {
  ErrorReporter reporter;
  // more to be added
};

We also propose to design subclasses to implement different types of optimization passes, e.g. function-level passes, module-level passes, and sequential passes. Each subclass itself could act as a pass manager. For instance, they could glob the require passes and execute them or build a dependency graph based on the given meta data.

// Module passes are mainly for IPO purpose.
class ModulePassNode : PassNode {
PassFunc pass_func;
// Other stuff...
};

// Function passes are performed on one function at a time. It produces a rewritten Relay function. 
class FunctionPassNode : PassNode {
 PassFunc pass_func;
// Other stuff...
};

class SequentialPassNode : PassNode {
   // Passes need to be executed.
  Array<Pass> passes;
  // Passes are disabled during optimization
  std::vector<std::string> disabled;
};

This has to deal with the node system between Python and C++, but it doesn’t require much specialization of different kinds of passes.

Each module or function optimization pass takes some required info, such as name and opt_level. In addition, it also accepts a pass_func argument, which is a curried function used for flexible design of interfaces. For example, it allows Module -> Module transformations on each Relay Module or Function.

template <typename T>
using PassFunc = TypedPackedFunc<TypedPackedFunc<T(T)>(const Module&)>;

SequentialPass is similar to Pytorch nn.Sequential which contains a host of passes for execution. It also has a disabled list that contains the list of disabled passes. For example, we may have a configuration which performs all passes at opt_level = 3 with some disabled passes disable-pass=xx, disable=yy. With SequentailPass, we could glob all passes at opt_level = 3 and exclude those in disabled.

In addition, some helper functions are provided to create each type of these functions, i.e.:

FunctionPass CreateFunctionPass(std::string name, int opt_level, PassFunc pass_func);
ModulePass CreateModulePass(std::string name, int opt_level, PassFunc pass_func);
SequentialPass CreateSequentialPass(std::string name, int opt_level, Array<Pass> passes, Array<tvm::Expr> disabled);

Frontend

Only some simple APIs will be needed for the frontend side. For example, we can provide users the following API to create a pass. The backend receives the information and decides which function it should use to create a Pass object.

def create_pass(pass_name, opt_level, pass_kind, pass_func, passes, disabled):
    pass

This following is an example that the pass_func would look like from Python.

def pass_fun(pass_state):
    my_class = MyClass(pass_state)
    def _transform(func):
        return my_class.visit(func)
    return _transform(func)

Unsolved problems

The current design only sketches the high-level class and interfaces, we propose to first implement the abstract interface with an eye for future compatibility. We believe there are a few more problems to be addressed in the future, and welcome community contributions in these areas:

• Add constraints to passes
• Build dependencies between passes (it could be a dependence graph built from metadata like pass name/id)
• Preserve analysis
• Validate analysis
• …

Any comments are welcome:)

@icemelon9 @wweic @tqchen @ZihengJiang @yzhliu @yidawang

[tqchen]

@zhiics can you summarize the discussion, update your post with the latest interface and then tag folks to see if they like it?

[tqchen]

Let me elaborate on what I mean by functional style and imperative execution. Currently, most of the current pass such as relay.ir_pass.FoldScaleAxis, is a simple function of Expr->Expr. This is a simple function, without any states.

I feel that most of the passes should be exposed as a function, which allows us to do Module->Module or Function->Function. Ideally, there should be no pass states, which is hard to reason about. The pass construction function could be a functor that takes in the configurations and returns you a Pass, which contains an operator(), or call that runs the transformation. In that way, you can use a Pass in the same way as you use current imperative pass functions, except that you might get more meta-data field about the pass for reasoning.

Why it is Important to Support Imperative Style

This is the same story when we talk about declarative vs imperative style code for deep learning.

The imperative style code is popular because we can easily show examples to your audience in jupyter notebook, it lowers the bar of entry, we can construct an Expr, run a FoldScaleAxis, print it 😃

The drawback of imperative style is that it lacks metadata, so that we cannot do automatic reasoning and addition of passes. In the first generation of DL frameworks, we have graph and execution session(which has a rough analogy to the pass manager style) .

What we really want, is the second generation of deep learning frameworks. e.g. functional Keras, or hybridize in gluon for passes. Which exposes the same imperative API, but uses the graph underneath.