Function build_and_run

tests/python/relax/test_codegen_cutlass.py:92–111

(mod, inputs_np, target, legalize=True, cuda_graph=False)

def build_and_run(mod, inputs_np, target, legalize=True, cuda_graph=False):
    with tvm.transform.PassContext(
        config={
            "relax.backend.use_cuda_graph": cuda_graph,
            "relax.transform.apply_legalize_ops": legalize,
        }
    ):
        ex = tvm.compile(mod, target)

    dev = tvm.device(target, 0)
    vm = relax.VirtualMachine(ex, dev)
    f = vm["main"]
    inputs = [tvm.runtime.tensor(inp, dev) for inp in inputs_np]

    # For cuda graph, run the compiled function twice to make sure that we can launch the cached
    # graph on the second run.
    if cuda_graph:
        f(*inputs)

    return f(*inputs).numpy()


def build_cutlass(mod, assert_all_bindings_fused=True, num_final_bindings=1):