Function test_tensorrt_only

tests/python/relax/test_transform_codegen_pass.py:132–157

(entry_func_name)

@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
@requires_tensorrt_runtime
def test_tensorrt_only(entry_func_name):
    mod, inputs, expected = setup_test()

    if entry_func_name != "main":
        mod[entry_func_name] = mod
        del mod["main"]

    # Define patterns that we want to offload to byoc
    # This test will offload entire model
    # Thus, define patterns for both `multiply` and `add` ops
    patterns = [
        ("tensorrt.multiply", is_op("relax.multiply")(wildcard(), wildcard())),
        ("tensorrt.add", is_op("relax.add")(wildcard(), wildcard())),
    ]

    new_mod = tvm.transform.Sequential(
        [
            relax.transform.FuseOpsByPattern(patterns),
            relax.transform.MergeCompositeFunctions(),
            relax.transform.RunCodegen(),
        ]
    )(mod)

    ex0 = tvm.compile(new_mod, target, params={})
    # Sanity check for the correctness and roundtrip
    check_roundtrip(ex0, dev, inputs, expected, entry_func_name)


@pytest.mark.gpu