tvm_tirx.backend.cuda.transform.LowerIket

Contents

tvm_tirx.backend.cuda.transform.LowerIket#

tvm_tirx.backend.cuda.transform.LowerIket()#

Lower frontend-only NVIDIA IKET annotations.

Run after SplitHostDevice and before MakePackedAPI. Strip annotations for regular compilation; emit metadata and NativeDump placeholders for IKET-enabled modules. The external run-iket process owns trace collection and postprocessing.

Examples

import tvm_tirx as tvm
from tvm_tirx.script import tirx as T
from tvm_tirx.backend.cuda.transform import LowerIket

@T.ir_module
class Before:
    @T.function
    def main(A: T.Tensor((1,), "int32")):
        T.cuda.iket_mark("store")
        A[0] = 1

@T.ir_module
class Expected:
    @T.function
    def main(A: T.Tensor((1,), "int32")):
        A[0] = 1

after = LowerIket()(Before)
tvm.ir.assert_structural_equal(after, Expected)
Contents