[RFC] Unified IR Infra

[tqchen]

Introduction

The TVM stack has been evolving for more than two years now. The current compiler stack contains several components that were designed at different points in time. This RFC proposes a unified IR infrastructure for the TVM stack by combining past lessons.

From a high-level the proposed infrastructure will consist of:

• A unified module, pass and type system for all IR function variants.
• Two major variants of IR expressions and functions: the high-level functional IR(relay)and the tensor-level IR for loop optimizations.
• First-class Python and hybrid script support, and a cross-language in-memory IR structure.
• A unified runtime::Module to enable extensive combination of traditional devices, microcontrollers and NPUs.

Unified IR Module and Pass Infra

The main component of our proposal is to introduce a unified IRModule structure that can contain different variants of functions(relay::Function/te::Function). The namespace te(tensor expression) is a tentative namespace for low-level functions, and comments and suggestions about name choices are more than welcome.

This change will simplify the intermediate representation data structures and terminology used in our previous compilation flow, which can be shown as follows:

 importer: model -> 
 high-level optimizations: relay::Module -> optimizations -> relay::Module  ->
 low-level optimizations: compute/schedule declaration -> Stmt-> ir passes ->Stmt ->
 device-specific codegen: LoweredFunc -> runtime::Module

As shown above, our current flow has different intermediate data structures at different stages (relay::Module, Stmt, LoweredFunc) with different terminologies. Under the new design, we will have a single module structure, and transformations between IRs become ir::Module to ir::Module transformations.

More importantly, we can use specific calling conventions, enabling different function variants to call each other. The following code snippet is a mock-up to demonstrate a module containing both a relay.Function and te.Function. The relay_add_one function can call into the te_add_one function using the destination-passing convention, where outputs are passed as inputs to the function.

def @relay_add_one(%x : Tensor((10,), f32)) {
    call_destination_passing @te_add_one(%x,  out=%b) 
} 

def @te_add_one(%a: NDArray, %b: NDArray) {
    var %n
    %A = decl_buffer(shape=[%n], src=%a)
    %B = decl_buffer(shape=[%n], src=%b) 
    // body ir contents need to be evolved
    for %i = 0 to 10 [data_par] {
        %B[%i] = %A[%i] + 1.0 
    }
}

Enabling both high-level functions(relay.Function) and tensor-expression functions to coexist in the same module enables the potential for cross layer optimizations. Of course, most of the transformations will only operate on one type of function and will simply ignore other functions in the same module.

Most importantly, the proposed change will minimize concepts for developers. Developers only need to know about ir::Module and runtime::Module. Every transformation is ir::Module -> ir::Module.

AutoTVM and other schedule transformations can be viewed as an intelligent way to transform an ir::Module. We are exploring ways to embed the current tensor expressions as a special type of Function in ir::Module.

Timeline Upgrade Path

Once we agree on the general design, we can step several refactor steps to bring the current codebase to the new IR infrastructure.

• Introduce ir::Module and ir::Function
• Move relay::Module to make use of ir::Module
• Move the current low-level IR infra ir_pass and LoweredFunc transformations to Module->Module
• Introduce text format and hybrid script specifications of the unified ir::Module
• Continue to evolve the design of the IRs themselves

[tqchen]

Migrated from https://discuss.tvm.ai/t/ir-unified-tvm-ir-infra/4801

[MarisaKirisame]

maybe some name that has connection to relay is good?
Axon seems quite nice: it is literally under relay, it is brain-related, it is short.

[tqchen]

Nice to consider names, @MarisaKirisame can you also post the name proposal to the original discuss thread to get other people's thoughts. To be clear, the proposal is to name the low-level ir module "axon"

[tqchen]

Move to #4647