Function func

tests/python/codegen/test_target_codegen_cuda_fp8.py:821–826

(A: T.Buffer((4,), dtype))

def test_const(dtype):
    @T.prim_func(s_tir=True)
    def func(A: T.Buffer((4,), dtype)) -> None:
        A_local = T.sblock_alloc_buffer((4,), dtype=dtype, scope="local")
        for tx in T.thread_binding(0, 4, "threadIdx.x"):
            for i in T.vectorized(4):
                A_local[i] = T.float32(1.0).astype(dtype)
            A[tx] = A_local[tx]

    mod = tvm.IRModule({"main": func})
    tvm.compile(mod, target="cuda")