How to do heterogeneous execution on cpu and gpu?

wrongtest · November 26, 2021, 6:20am

Hi~ Can this unittest case help you?

apache/tvm/blob/be03d62e5b0afd607964365bc73e94f72fdfaaef/tests/python/relay/test_vm.py#L1071


        input(),
        tvm.target.Target("cuda"),
    )


    # The newshape annotation should have been turned into a constant on the CPU.
    assert "VirtualDevice[0]: device type 1" in exe.virtual_devices
    assert "Constant[0]: has shape int64[3] on device index 0" in exe.constants




@tvm.testing.requires_cuda
def test_multi_targets():
    # Build an IRModule.
    n = 10
    x = relay.var("x", shape=(n,))
    y = relay.var("y", shape=(n,))
    z = relay.var("z", shape=(n,))
    f = relay.Function([x, y, z], x + relay.op.annotation.on_device(y + z, tvm.cpu()))
    mod = IRModule.from_expr(f)


    # Compile to VMExecutable.
    with tvm.transform.PassContext(