Function _create_mod

tests/python/codegen/test_target_codegen_cuda_fp8.py:51–71

(dtype)

    dtype, nv_dtype = input

    def _create_mod(dtype):
        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(
                A: T.Buffer((64,), dtype),
                B: T.Buffer((64,), dtype),
                C: T.Buffer((64,), dtype),
            ):
                T.func_attr({"tirx.noalias": True})
                for i_0 in T.thread_binding(2, thread="blockIdx.x"):
                    for i_1 in T.thread_binding(32, thread="threadIdx.x"):
                        with T.sblock("C"):
                            v_i = T.axis.spatial(64, i_0 * 32 + i_1)
                            T.reads(A[v_i], B[v_i])
                            T.writes(C[v_i])
                            C[v_i] = T.Cast(
                                dtype, T.Cast("float16", A[v_i]) + T.Cast("float16", B[v_i])
                            )

        return Module

    mod = _create_mod(dtype)
    target = "cuda"