Function test_tensor_with_vdevice

tests/python/relax/test_tvmscript_parser.py:774–816

()

def test_tensor_with_vdevice():
    vdevices = [
        VDevice("llvm"),
        VDevice("cuda", 0),
        VDevice("metal", 0, "global"),
        VDevice({"kind": "cuda", "arch": "sm_80"}, 0),
    ]

    @I.ir_module(s_tir=True)
    class TestModule:
        I.module_attrs({"attr": 10})
        I.module_global_infos(
            {
                "vdevice": [
                    I.vdevice("llvm"),
                    I.vdevice("cuda", 0),
                    I.vdevice("metal", 0, "global"),
                    I.vdevice({"kind": "cuda", "arch": "sm_80"}, 0),
                ]
            }
        )

        @R.function
        def foo(
            a: R.Tensor((128, 128), "float32", "cuda:1"),
            b: R.Tensor((128, 128), "float32", "llvm"),
            c: R.Tensor((128, 128), "float32", "vdevice:3"),
        ) -> R.Tensor((128, 128), "float32", "cuda:1"):
            s = R.add(a, c)
            return s

    a = relax.Var("a", R.Tensor((128, 128), "float32", vdevices[3]))
    b = relax.Var("b", R.Tensor((128, 128), "float32", vdevices[0]))
    c = relax.Var("c", R.Tensor((128, 128), "float32", vdevices[3]))
    bb = relax.BlockBuilder()
    with bb.function("foo", (a, b, c)):
        out = bb.emit(relax.op.add(a, c))
        bb.emit_func_output(out)
    mod = bb.get()
    mod = mod.with_attr("attr", 10)
    mod.update_global_info("vdevice", vdevices)

    _check(TestModule, mod)


def test_direct_return():