[Vulkan] ti.sync() changes program behavior on Vulkan

[AmesingFlank]

Example of when ti.sync() changes program behavior on Vulkan: #3790. In this example, the program only behaves correctly when the ti.sync() is added.

On vulkan, we aggressively batch consecutive kernel launches into a single command buffer. We probably need to add a buffer_barrier or something between consecutive compute dispatches. I will investigate this.

[bobcao3]

We do have barriers between them.

taichi/taichi/backends/vulkan/runtime.cpp

Line 336 in 9f0d9f5

cmdlist->memory_barrier();

[AmesingFlank]

Yeah, I saw it. So why is this happening :((

[AmesingFlank]

To repro:

import taichi as ti

ti.init(ti.vulkan)

def parallel_sort(x):
    N = x.shape[0]

    @ti.kernel
    def sort_stage(x:ti.template(), N:int, p:int, k:int, invocations:int):
        for inv in range(invocations):
            j = k%p + inv * 2 * k
            for i in range(0, min(k,N-j-k)):
                a = i+j
                b = i+j+k
                if int(a / (p*2)) == int(b / (p*2)):
                    val_a = x[a]
                    val_b = x[b]
                    if val_a > val_b:
                        x[a] = val_b
                        x[b] = val_a

    p = 1
    while p < N:
        k = p
        while k >= 1:
            invocations = int((N-k-k%p) / (2*k)) + 1
            sort_stage(x,N,p,k, invocations)
            # ti.sync()
            k = int(k/2)
        p = int(p * 2)
  
def test_sort():
    def test_sort_for_dtype(dtype,N):
        x = ti.field(dtype, N)

        @ti.kernel
        def fill():
            for i in x:
                x[i] = ti.random()*N
        
        fill()
        parallel_sort(x)

        x_host = x.to_numpy()
        
        for i in range(N-1):
            assert x_host[i] <= x_host[i+1]

    test_sort_for_dtype(ti.i32,1)
    test_sort_for_dtype(ti.i32,256)
    test_sort_for_dtype(ti.i32,100001)
    test_sort_for_dtype(ti.f32,1)
    test_sort_for_dtype(ti.f32,256)
    test_sort_for_dtype(ti.f32,100001)

test_sort()

[bobcao3]

What if we change

taichi/taichi/backends/vulkan/vulkan_device.cpp

Line 808 in 9f0d9f5

/*srcStageMask=*/

Change it to src_stage_mask=BOTTOM_OF_PIPE and dst_stage_mask=TOP_OF_PIPE

[AmesingFlank]

Doesn't seem to fix it :<