tvm_tirx.tirx.transform.Apply

Contents

tvm_tirx.tirx.transform.Apply#

tvm_tirx.tirx.transform.Apply(ftransform)#

Apply ftransform to each function in the Module.

This function is a thin wrapper around function_pass()

Parameters:

ftransform (tvm_tirx.tirx.Function -> tvm_tirx.tirx.Function) – The transformation pass.

Returns:

fpass (tvm_tirx.ir.transform.Pass) – The result pass

Examples

import tvm_tirx as tvm
from tvm_tirx.script import ir as I, tirx as T
from tvm_tirx.tirx import transform as tr

@I.ir_module
class Before:
    @T.function(private=True)
    def main():
        pass

@I.ir_module
class Expected:
    @T.function(private=True)
    def main():
        T.func_attr({"example": 1})

def run_pass():
    return tr.Apply(lambda f: f.with_attr("example", 1))(Before)

tvm.ir.assert_structural_equal(run_pass(), Expected)
Contents