Function test_global_info_vdevice

tests/python/relax/test_tvmscript_parser.py:315–362

()

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

    @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({"kind": "cuda", "arch": "sm_80"}, 0),
                    I.vdevice("metal", 0, "global"),
                ]
            }
        )

        @T.prim_func(private=True, s_tir=True)
        def tir_func(
            x: T.Buffer((T.int64(128), T.int64(128)), "float32"),
            y: T.Buffer((T.int64(128), T.int64(128)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            for i, j in T.grid(T.int64(128), T.int64(128)):
                with T.sblock():
                    vi, vj = T.axis.remap("SS", [i, j])
                    y[vi, vj] = x[vi, vj] + 1.0

        @R.function
        def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
            cls = TestModule
            gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128, 128), dtype="float32"))
            return gv0

    x = relax.Var("x", R.Tensor((128, 128), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,)):
        out = bb.emit_te(lambda x: x + 1, x, primfunc_name_hint="tir_func")
        bb.emit_func_output(out)
    mod = bb.get()
    mod.update_global_info("vdevice", vdevices)
    mod = mod.with_attr("attr", 10)
    _check(TestModule, mod)


def test_relax_tensor_op():