[RFC] Simplifying Module System

Development
#1

Based on some conversation with @jroesch. relay.Module has become the central object of interest for IR transformations after we introduced the new pass infra. It is imperative to make sure that the Module API is simple, concise and easy to use.

This RFC summarizes a few thoughts on that direction.

Have a Separation between Local and Extern Functions

Right now every function is named, and need to have a unique name. This creates certain problems in cases where we might want to easily allocate a function name that is local. Like most low-level compiler infra, we might want to have a distinction between local and extern functions, and only allow string based access to extern functions.

Fix entry function to main

At the moment there is an entry_func field of global var that indicates the global function. Given most recent discussions there seems to be an interest to harden the name to “main”. We might consider remove or keep the entry_func field, depending on whether we also store a mapping in
extern_function_map_

Store analysis results as function attributes

Many passes rely on analysis results(e.g. MAC count or layout map). These analysis result can be stored as a function attributes inside the function.

Some possible API Simplification

Because the module is at the center of evolution, it is important to make creation and manipulation of module easy. We might want to introduce clear sugars that reduce the amount of code a user might want to write.

Here are a few possible ideas about API simplifications

# short hand for relay.Module.from_expr(relay.Function([x], x))
mod = relay.Module.main([x], x)
# add a new extern function to module
# need to debate whether allow override, I think we should
mod["myfunc"] = relay.Function([x], relay.log(x))
# explicit update API, like dictionary's update
mod.update({"myfunc": relay.Function([x], relay.log(x))})
# get an external function
main_func = mod["main"]

def less_common_code():
    # get a global var out before setting a function
    gv2 = mod.get_global_var("name_hint")
    # similar to function declaration. 
    # allocate global variable, avoid conflict
    # private function might be pruned if not reference by extern
    gv1 = mod.declare_global_var("name_hint", linkage="private")
    mod[gv1] = make_function(gv1)
3 Likes
#2

For some expressions, we probably cannot build a Module on top of them as FromExpr needs to perform InferType for updating the Module.

For example,

x = relay.Var("x")
f = relay.Function([x], x + x)

We may want to perform certain optimizations on f (e.g. dce), but it looks we cannot directly create a module as there is no typing info available. Should we only allow Expr -> Expr in this case?

#3

I think this is a design trade-off.

i.e. Does module allow functions that do not have type information. The current design favour the fact that the functions always need to contain full type information. I think is good overall to make that assumption.

There are cases however, when types comes later (e.g. when importing a model and parameters are still unknown). Or certain optimisations that do not depends on type info. I think the current design will dis-favour such optimisation. After any type is introduced, the optimization can be performed as the program still type check with Any.

cc @jroesch

#4

I agree we probably want to keep Module strongly typed. But I am not quite sure if Any solves this problem, as Any will still be a given type. In the particular case above, the type is null. Or we could make this case default to Any.