类别：编译时：Mosaic 未经证实的内存访问对齐
当编译器分析内存访问操作（例如 vector.load、vector.store、tpu.load 或 tpu.store）时，如果无法静态证明用于特定维度的动态索引是所需平铺大小的倍数，就会发生此错误。
示例错误消息：
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>
XLA 后端：TPU
概览
当内核加载或存储向量时，内存地址（根据基指针加上动态索引计算得出）必须与硬件上的向量平铺大小保持一致。例如，如果某个维度按 128 个元素进行平铺，则用于访问该维度的动态索引必须为 0、128、256 等。请注意，许多操作（例如向量加载和存储）对静态索引没有此类要求。
编译器使用静态分析来强制执行此要求。它会回溯索引变量的历史记录，查找生成该变量的算术运算（例如乘法、加法）。如果编译器无法保证（在编译时）结果值始终可被平铺大小整除，则会引发此错误。
编译器会以相同方式处理“已证实的不对齐”和“未知对齐”。因此，如果您使用的指数在数学上保证会错位（例如，i * 128 + 32），编译器将引发相同的错误。
因此，在以下情况下可能会发生此错误
- 您可以使用运行时变量（动态索引）来访问内存。
- 指数计算逻辑过于复杂，编译器无法分析。
- 该指数在数学上有效，但在代码中缺少明确的证明。
- 静态分析可确定“已证实的不一致”。
调试
如需解决此错误，您可以尝试以下方法：
1. 明确断言对齐
如果您知道自己的索引有效，但编译器无法证明这一点，请使用 tpu.assume_multiple 操作。这相当于向编译器承诺某个值可被特定因子整除。
2. 使用对齐的加载和旋转
如果错位是故意的，则不加载较小的未对齐向量段，而是：
- 加载一个较大且完全对齐的功能块，然后按动态量旋转值，将所需数据移到相应位置（因为不支持具有动态起始索引的矢量切片）。或者
- 重塑张量或填充张量，使数据从索引 0 开始，并且访问之间的步长与硬件对齐一致。
- 示例：如果您从偏移量 1 开始迭代大小为 32 的块，则偏移量为 1、33、65... （未对齐）。
- 修复：将数据重新打包到新张量中，其中第一个块位于 0，维度填充为 128。您的偏移量将变为 0、128、256...，从而满足对齐要求。
这些方法会消耗更多内存，但通常可以简化内核逻辑，并消除对人工对齐断言的需求。