Thanks for the clarifications. This discussion thread relates quite a few key design principles which I try to summarize as follows:
P0: Every (composite) IRModule transformation should be a Pass
As an analogy in deep learning frameworks. We tend to design every tensor transformations as nn.Module or a composite Layer. The build_extern function certainly is a IRModule transformation. However, it is not presented as a composite Pass.
The main advantage of presenting everything as a Pass is the uniformity, and further composition of the pass itself. For example, we could further push the BuildExtern(ignoring the naming issue here, we need a better name than âexternâ) into another Sequential, which combines it with further optimizations.
The key insight is that having everything as a pass encourages uniformity and brings changes for further improvements. For example, we could develop a meta pass that tries a list of possible transformations, or search over them.
# Try each pass in order
opt_pass = transform.TrySequential(
[BuildExtern("dnnl"), NormalPipeline(), BuildExtern("xyzlib")])
mod = opt_pass(mod)
P1: Always decouple target capability from dispatch strategy
Different target backends are likely going to be implemented by different people.
- Target capability means âtarget X can do Yâ. For example, âdnnl can dispatch conv2dâ.
- Strategy means âwhat to do given capabilitiesâ. For example, in the case of external compilation, try to dispatch to dnnl first.
The current registration API mixes the two. This means that for each type of the backend, developers have to customize the same set of functions and can not mix and match target compilation capabilities. If we have a strategy that mixes xyzlib and dnnl in the future, we have to implement another xyz_plus_xnnl compiler somewhere.
The following straw man, on the other hand, allows different target capability registration sits in a specific backend folder. Note that I am not pushing for this specific convention, but would like to point out that this decouple principle is important and should be taken into consideration seriously.
# register a dnnl specific target
reg.register("nn.conv2d", "target.dnnl", batch_norm_func)
# register xyzlib specific logic
reg.register("nn.conv2d", "target.xyzlib", batch_norm_func)
In terms of the ["dnnl", "xyzlib]" example, I did not mean that we need the best strategy now(There could be multiple in the future), but instead we should design the target capability registration to be simple enough so that we can easily add new strategies in the future. This principle also relates to @haichenâs AutoTVM strategy design.
Meta-P2: Customization as a Natural Consequence of Existing Design
This point could be a bit controversial. In particular, I think we should try our best to eliminate terms like âcustomizationâ, âextensionâ or âexternâ. Let me clarify what do I mean.
In the case of bring your own codegen. What we are really doing here is to design a Target that have a specific compilation flow to the runtime.Module. Target specific compilation flow is a term that could be designed internal to the project. Then bring your own codegen could simply become a way to define new Target and its compilation flow. We could also introduce a separate concept besides Target, if we think we need something else.
The name extern usually results in a single point of extension(here is a plugin API that you can put your rules). A good design in the infrastructure usually can results in multiple mix and match. And make the normal compilation flow and the custom one under the unified scheme.
In the case of build_extern, this is a specific function that provides a new feature. However, forcing it into an existing known term(Pass) will provide more uniformity and advantages for the project.
To summarize, these customization points should be natural consequence of the existing infra. If existing infra does not handle something well, we should think if we can improve it to encapsulate the things we need eventually.
Note that this is a meta point, and we do not have to do it now, but it would be great if we can collectively think how we can improve âexternâ, âcustomâ as much as possible as part of the development process.
