ti.gpu runs the wrong answer for prime demo

[learningpro]

Describe the bug
I just try count_primes demo on my MacBookPro M1. It is interesting that ti.cpu and ti.gpu give out different answers.
The only difference is line 4. ti.gpu runs very quickly but the answer is wrong.

To Reproduce

[turbo0628]

Could you paste the full code in text to make it easier for others to reproduce?

[learningpro]

Could you paste the full code in text to make it easier for others to reproduce?

You can Ctrl-C and Ctrl-V from https://docs.taichi-lang.org/docs/accelerate_python

[tao-j]

The same for Vulkan backend.

import taichi as ti
ti.init(arch=ti.gpu)

@ti.func
def is_prime(n: int):
    result = True
    for k in range(2, int(n ** 0.5) + 1):
        if n % k == 0:
            result = False
            break
    return result

@ti.kernel
def count_primes(n: int) -> int:
    count = 0
    for k in range(2, n):
        if is_prime(k):
            count += 1

    return count

print(count_primes(100_000_000))
print(count_primes(10_000_000))

Output from AMD APU

[Taichi] version 1.3.0, llvm 15.0.4, commit 0f25b95e, linux, python 3.10.9
[W 01/15/23 21:08:20.539 107616] [cuda_driver.cpp:load_lib@36] libcuda.so lib not found.
[Taichi] Starting on arch=vulkan
[I 01/15/23 21:08:20.561 107616] [vulkan_device_creator.cpp:pick_physical_device@399] Found Vulkan Device 0 (AMD Radeon Graphics (RADV RENOIR))
[I 01/15/23 21:08:20.561 107616] [vulkan_device_creator.cpp:pick_physical_device@399] Found Vulkan Device 1 (llvmpipe (LLVM 15.0.6, 256 bits))
[I 01/15/23 21:08:20.561 107616] [vulkan_device_creator.cpp:create_logical_device@469] Vulkan Device "AMD Radeon Graphics (RADV RENOIR)" supports Vulkan 0 version 1.3.230
0
664579

Expected result:

5761455
664579

Btw, vulkan failed on Intel

[Taichi] Starting on arch=vulkan
[I 01/15/23 21:01:20.759 323725] [vulkan_device_creator.cpp:pick_physical_device@399] Found Vulkan Device 0 (Intel(R) Xe Graphics (TGL GT2))
[I 01/15/23 21:01:20.759 323725] [vulkan_device_creator.cpp:pick_physical_device@399] Found Vulkan Device 1 (llvmpipe (LLVM 15.0.6, 256 bits))
[I 01/15/23 21:01:20.759 323725] [vulkan_device_creator.cpp:create_logical_device@469] Vulkan Device "Intel(R) Xe Graphics (TGL GT2)" supports Vulkan 0 version 1.3.230
[E 01/15/23 21:01:44.704 323725] [vulkan_device.cpp:submit@1720] Vulkan Error : -4 : failed to submit command buffer


Traceback (most recent call last):
  File "/wks/Downloads/nxc/cbr/test.py", line 22, in <module>
    print(count_primes(100_000_000))
  File "/home/yuanzhang/venv/lib64/python3.10/site-packages/taichi/lang/kernel_impl.py", line 945, in wrapped
    return primal(*args, **kwargs)
  File "/home/yuanzhang/venv/lib64/python3.10/site-packages/taichi/lang/kernel_impl.py", line 872, in __call__
    return self.runtime.compiled_functions[key](*args)
  File "/home/yuanzhang/venv/lib64/python3.10/site-packages/taichi/lang/kernel_impl.py", line 804, in func__
    runtime_ops.sync()
  File "/home/yuanzhang/venv/lib64/python3.10/site-packages/taichi/lang/runtime_ops.py", line 8, in sync
    impl.get_runtime().sync()
  File "/home/yuanzhang/venv/lib64/python3.10/site-packages/taichi/lang/impl.py", line 489, in sync
    self.prog.synchronize()
RuntimeError: [vulkan_device.cpp:submit@1720] Vulkan Error : -4 : failed to submit command buffer
[E 01/15/23 21:01:44.732 323725] [vulkan_device.cpp:submit@1720] Vulkan Error : -4 : failed to submit command buffer


[E 01/15/23 21:01:44.733 323725] [vulkan_device.cpp:submit@1720] Vulkan Error : -4 : failed to submit command buffer


terminate called after throwing an instance of 'std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >'