tvm_tirx.tirx.transform.BindTarget

Contents

tvm_tirx.tirx.transform.BindTarget#

tvm_tirx.tirx.transform.BindTarget(target)#

Bind targets and classify host/device helper functions.

Preserves existing device targets and updates host information on exported functions. Helpers used on both sides may be cloned for the host.

Device helper calls must be inside explicit launch_thread or device_entry regions. Lower S-TIR thread bindings before classifying host/device calls.

Parameters:

target (tvm_tirx.target.Target) – target

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

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

def run_pass():
    return tr.BindTarget(tvm.target.Target("cuda"))(Before)

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