Function while_loop

tests/python/tvmscript/test_tvmscript_roundtrip.py:2061–2075

()

def while_loop():
    @T.prim_func(s_tir=True)
    def while_loop(a: T.handle, b: T.handle) -> None:
        A = T.match_buffer(a, (16,), "float32")
        B = T.match_buffer(b, (16,), "float32")
        i = T.sblock_alloc_buffer((), "int32", scope="local")
        for ii in range(16):
            with T.sblock():
                vi = T.axis.S(16, ii)
                B[vi] = 0
            while i[()] < 10:
                for j in range(16):
                    B[j] += A[j]

    return while_loop


# fmt: off