Metal: 4Gb memory limit causes quiet overwriting of field values

[kitizz]

Describe the bug

When the cumulative size of ti.fields in Metal memory exceeds 4Gb, it appears there is an address overflow that causes reads/writes to overflow back to the start of the memory block.

To Reproduce

import numpy as np
import taichi as ti

ti.init(arch=ti.gpu, device_memory_GB=10)

N = 1073741824  # x 4 bytes == 4GB
M = 1

a = ti.field(ti.i32, shape=N)
b = ti.field(ti.i32, shape=M)

arr_before = np.zeros(N, dtype=np.int32)
print("Megabytes:", arr_before.nbytes / 1024 / 1024)

a.from_numpy(arr_before)
b.from_numpy(np.ones(M, dtype=np.int32))

arr_after = a.to_numpy()

num_overwritten = np.sum(arr_after)
assert num_overwritten == 0, f"Memory bug! {num_overwritten} elements overwritten!"
print("All good!")

Log/Screenshots

$ python taichi_memory_bug.py
[Taichi] version 1.7.4, llvm 15.0.7, commit b4b956fd, osx, python 3.13.7
[Taichi] Starting on arch=metal
Megabytes: 4096.0
Traceback (most recent call last):
  File "/Users/kit/projects/python/play/taichi_memory_bug.py", line 24, in <module>
    assert num_overwritten == 0, f"Memory bug! {num_overwritten} elements overwritten!"
           ^^^^^^^^^^^^^^^^^^^^
AssertionError: Memory bug! 1 elements overwritten!

Additional comments

Arch: Apple M1 Max

This does not occur when arch=cpu.

I've observed that this issue can be side-stepped by using FieldsBuilder if each individual field is under 4Gb. But this remains an issue for single fields that exceed this size.