tvm_tirx.backend.trn.transform.AnnotateKernelArguments

class tvm_tirx.backend.trn.transform.AnnotateKernelArguments(*args, **kwargs)

Bases: ModulePass

Order kernel inputs before outputs and annotate their per-kernel input count.

Run after SplitHostDevice and before MakePackedAPI. The source trn.num_input_args contract remains on each host function until its kernel calls have been matched to the surviving host parameters and captures. Functions without extracted kernels retain their own contract.