tvm_tirx.tirx.transform.DecorateDeviceScope

tvm_tirx.tirx.transform.DecorateDeviceScope()

Wrap each function body in a tirx.device_scope region.

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(x: T.int32):
        T.evaluate(x)

@I.ir_module
class Expected:
    @T.function(private=True)
    def main(x: T.int32):
        with T.device_scope():
            T.evaluate(x)

def run_pass():
    return tr.DecorateDeviceScope()(Before)

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