tvm_tirx.tirx.transform.VerifyMemory

Contents

tvm_tirx.tirx.transform.VerifyMemory#

tvm_tirx.tirx.transform.VerifyMemory()#

Reject illegal host-side access to device memory.

Returns the module unchanged on success and raises an error on invalid access. Functions without a target attribute are skipped, as are functions with a non-default calling convention.

Returns:

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

Examples

Valid device memory accesses are checked without changing the IR.

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(A: T.Tensor((4,), T.float32)):
        T.func_attr({"target": T.target({"kind": "cuda", "arch": "sm_75"})})
        with T.launch_thread("threadIdx.x", 4) as tx:
            A[tx] = T.float32(1.0)

@I.ir_module
class Expected:
    @T.function(private=True)
    def main(A: T.Tensor((4,), T.float32)):
        T.func_attr({"target": T.target({"kind": "cuda", "arch": "sm_75"})})
        with T.launch_thread("threadIdx.x", 4) as tx:
            A[tx] = T.float32(1.0)

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

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