def METHOD_NAME():
    @T.prim_func
    def add_0(
        A: T.Buffer((2048, 2048, 2048), "float32"),
        B: T.Buffer((2048, 2048, 2048), "float32"),
    ) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main"})
        # body
        # with T.block("root")
        A_cached = T.alloc_buffer([2048, 2048, 2048], dtype="float32")
        for i0, j0, i1, j1, k0, i2 in T.grid(128, 64, 4, 4, 64, 4):
            for ax0, ax1, ax2 in T.grid(1, 8, 32):
                with T.block("move"):
                    vi = T.axis.spatial(2048, i0 * 16 + i1 * 4 + i2 + ax0)
                    vj = T.axis.spatial(2048, j0 * 32 + j1 * 8 + ax1)
                    vk = T.axis.spatial(2048, k0 * 32 + ax2)
                    T.reads(A[vi, vj, vk])
                    T.writes(A_cached[vi, vj, vk])
                    A_cached[vi, vj, vk] = A[vi, vj, vk]
            for j2, k1 in T.grid(8, 32):
                with T.block("add"):
                    vi = T.axis.spatial(2048, i0 * 16 + i1 * 4 + i2)
                    vj = T.axis.spatial(2048, j0 * 32 + j1 * 8 + j2)
                    vk = T.axis.spatial(2048, k0 * 32 + k1)
                    T.reads(A_cached[vi, vj, vk])
                    T.writes(B[vi, vj, vk])
                    B[vi, vj, vk] = A_cached[vi, vj, vk] + T.float32(1)
    decision_0 = [
        ("SampleComputeLocation", 5),
    ]
    mod = Add
    actual = generate_design_space(
        kind="llvm",
        mod=mod,
        target=Target("llvm"),
        types=None,
        sch_rules=[ms.schedule_rule.RandomComputeLocation()],
    )
    check_sketches(
        mod,
        sketches=actual,
        expected_mods=[add_0],
        expected_decisions=[decision_0],
    )