def METHOD_NAME(A: T.Buffer((128, 128), "float32")):
    # Buffer B should be remapped
    B = T.alloc_buffer((128, 128), "float32")
    # i, j should be remapped
    for i, j in T.grid(128, 128):
        with T.block("B"):
            # vi, vj should be remapped
            vi, vj = T.axis.remap("SS", [i, j])
            T.reads(A[vi, vj])
            T.writes(B[vi, vj])
            B[vi, vj] = A[vi, vj] * 2.0