Function check_inf_nan

tests/python/codegen/test_target_codegen_rocm.py:31–50

(dev, n, value, dtype)

@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
def test_rocm_inf_nan():
    def check_inf_nan(dev, n, value, dtype):
        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
                T.func_attr({"tirx.noalias": True})
                for i_0 in T.thread_binding(1, thread="blockIdx.x"):
                    for i_1 in T.thread_binding(128, thread="threadIdx.x"):
                        with T.sblock("C"):
                            v_i = T.axis.spatial(1, i_0 * 128 + i_1)
                            T.where(i_0 * 128 + i_1 < 1)
                            T.reads()
                            T.writes(C[v_i])
                            C[v_i] = T.Cast(dtype, value)

        fun = tvm.compile(Module, "rocm")
        a = tvm.runtime.empty((n,), dtype, dev)
        c = tvm.runtime.empty((n,), dtype, dev)
        # Only need to test compiling here
        fun(a, c)

    dev = tvm.rocm(0)