Kategori: Waktu Kompilasi: Penyesuaian Akses Memori Mosaic yang Belum Terbukti
Error ini terjadi saat compiler menganalisis operasi akses memori (seperti
vector.load, vector.store, tpu.load, atau tpu.store) dan tidak dapat
membuktikan secara statis bahwa indeks dinamis yang digunakan untuk dimensi tertentu adalah
kelipatan ukuran pengtile yang diperlukan.
Contoh Pesan Error:
INTERNAL: Mosaic failed to compile TPU kernel: cannot statically prove that index in dimension 1 is a multiple of 128
at location: ...
The MLIR operation involved:
%14372 = "vector.load"(%14371, %93, %14363) : (memref<4x256xf32, #tpu.memory_space<vmem>>, index, index) -> vector<1x32xf32>
Backend XLA: TPU
Ringkasan
Saat kernel Anda memuat atau menyimpan vektor, alamat memori
(dihitung dari pointer dasar ditambah indeks dinamis) harus selaras dengan
ukuran pengtile vektor di hardware. Misalnya, jika dimensi diatur secara berpetak
oleh 128 elemen, indeks dinamis yang digunakan untuk mengaksesnya harus 0, 128, 256,
dll. Perhatikan bahwa banyak operasi (seperti pemuatan dan penyimpanan vektor) tidak memiliki
persyaratan tersebut untuk indeks statis.
Compiler menerapkan persyaratan ini menggunakan analisis statis. Alat ini melacak histori variabel indeks kembali melalui operasi aritmatika yang menghasilkannya (misalnya, perkalian, penambahan). Jika kompiler tidak dapat menjamin (pada waktu kompilasi) bahwa nilai yang dihasilkan akan selalu dapat dibagi dengan ukuran petak, error ini akan muncul.
Compiler memperlakukan "ketidakselarasan yang terbukti" dan "ketidakselarasan yang tidak diketahui" secara identik.
Jadi, jika Anda menggunakan indeks yang secara matematis dijamin tidak selaras (misalnya,
i * 128 + 32), compiler akan memunculkan error yang sama.
Jadi, error ini dapat terjadi jika
- Anda menggunakan variabel runtime (indeks dinamis) untuk mengakses memori.
- Logika penghitungan indeks terlalu rumit untuk dianalisis oleh compiler.
- Indeks valid secara matematis, tetapi tidak memiliki bukti eksplisit dalam kode.
- Analisis statis menentukan "ketidakselarasan yang terbukti".
Proses debug
Untuk mengatasi error ini, Anda memiliki opsi berikut:
1. Menyatakan Keselarasan Secara Eksplisit
Jika Anda tahu indeks Anda valid, tetapi compiler tidak dapat membuktikannya, gunakan operasi
tpu.assume_multiple. Hal ini berfungsi sebagai janji kepada compiler bahwa
nilai dapat dibagi dengan faktor tertentu.
2. Menggunakan Muatan yang Sejajar dan Putar
Dalam skenario ketika ketidakselarasan disengaja, daripada memuat segmen vektor kecil yang tidak selaras:
- Muat petak yang lebih besar dan selaras sepenuhnya, lalu putar nilai dengan jumlah dinamis untuk memindahkan data yang diinginkan ke posisinya (karena irisan vektor dengan indeks awal dinamis tidak didukung). atau
- Ubah bentuk atau tambahkan padding pada tensor sehingga data dimulai pada indeks 0 dan langkah
di antara akses cocok dengan penyelarasan hardware.
- Contoh: Jika Anda melakukan iterasi pada potongan berukuran 32 yang dimulai pada offset 1, offset Anda adalah 1, 33, 65... (tidak selaras).
- Perbaikan: Kemas ulang data ke dalam tensor baru dengan potongan pertama di 0 dan dimensi yang diisi hingga 128. Offset Anda menjadi 0, 128, 256..., yang memenuhi persyaratan perataan.
Metode ini menggunakan lebih banyak memori, tetapi sering kali menyederhanakan logika kernel dan menghilangkan kebutuhan akan pernyataan penyelarasan manual.