[X86_64] [AMX] Wrong tile register reload instruction is generated

[ienkovich]

amx-reload-bug.txt

I build LLC from the main branch and compile the attached test case with -mcpu=sapphirerapids. The resulting assembly code has the following instructions:

        tilestored      %tmm0, 6016(%rsp,%rax)  # 1024-byte Folded Spill
        tileloadd       6016(%rsp), %tmm7       # 1024-byte Folded Reload

Here tile reload misses index register with a stride, therefore zero stride is used and the first row of the stored tile is broadcasted. Also, it looks like this spill/reload pair can be avoided.

[llvmbot]

@llvm/issue-subscribers-backend-x86

Author: Ilya Enkovich (ienkovich)

[amx-reload-bug.txt](https://github.com/user-attachments/files/17157112/amx-reload-bug.txt)

I build LLC from the main branch and compile the attached test case with -mcpu=sapphirerapids. The resulting assembly code has the following instructions:

        tilestored      %tmm0, 6016(%rsp,%rax)  # 1024-byte Folded Spill
        tileloadd       6016(%rsp), %tmm7       # 1024-byte Folded Reload

Here tile reload misses index register with a stride, therefore zero stride is used and the first row of the stored tile is broadcasted. Also, it looks like this spill/reload pair can be avoided.

[phoebewang]

@ienkovich Thanks for the report! Should be fixed by #110226

I was astonished this bug existing since the beginning and feel ashamed for not finding this obvious mistake when modifying code around it 🤦‍♀️

[phoebewang]

Also, it looks like this spill/reload pair can be avoided.

It cannot be avoided. We don't have copy instruction from one tile to another, so we have to do the spill/reload. I think the code was unrolled too many times. It is not a good idea to do more unroll for AMX code, because we only have 8 tile registers. The frequent spill/reload till registers will reduce performance greatly. So we control the unroll times in practice, which explains why the bug exists so long time :)

[ienkovich]

@phoebewang Thanks a lot for the quick fix!

It cannot be avoided. We don't have copy instruction from one tile to another, so we have to do the spill/reload.

The fact that we can reload spilled tmm0 to tmm7 means these tiles have the same config. So simple tmm0 <-> tmm7 renaming down the code should help to remove this copy through load+store. I assumed that the register allocator didn't pay too much attention to this, because for other registers it's not a big issue and you don't need to copy them through memory. Is there some reason forcing us to make this copy?

So we control the unroll times in practice, which explains why the bug exists so long time :)

Absolutely. I've found this problem because of the wrong config used in my codegen. I allowed 16 live tiles at a time instead of 8. Sometimes bugs help to find other bugs :)

[phoebewang]

The fact that we can reload spilled tmm0 to tmm7 means these tiles have the same config. So simple tmm0 <-> tmm7 renaming down the code should help to remove this copy through load+store. I assumed that the register allocator didn't pay too much attention to this, because for other registers it's not a big issue and you don't need to copy them through memory. Is there some reason forcing us to make this copy?

It is not so easy, RA is a global thing. Actually, the registers are renamable to the RA. If you take a deep look at its work, it has already reduced 3 COPYs: https://godbolt.org/z/4qxoe8z1G. I believe the last one must not due to it slacks off :)

[ienkovich]

BTW I also had a case where allocator couldn't allocate registers at all. Here is a snippet:

    %1588 = "amx.tileloadd64"(%5, %4, %1322, %3) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<16xi32>> loc(#loc11)
    %1589 = "amx.tileloadd64"(%5, %4, %1386, %3) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<16xi32>> loc(#loc11)
    %1590 = "amx.tileloadd64"(%5, %4, %1450, %3) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<16xi32>> loc(#loc11)
    %1591 = "amx.tileloadd64"(%5, %4, %1514, %3) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<16xi32>> loc(#loc11)
    %1592 = "amx.tileloadd64"(%5, %2, %951, %1) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<32xi8>> loc(#loc11)
    %1593 = "amx.tileloadd64"(%5, %2, %967, %1) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<32xi8>> loc(#loc11)
    %1594 = "amx.tileloadd64"(%5, %2, %983, %1) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<32xi8>> loc(#loc11)
    %1595 = "amx.tileloadd64"(%5, %2, %999, %1) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<32xi8>> loc(#loc11)
    %1596 = "amx.tileloadd64"(%0, %4, %1021, %3) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<8 x vector<64xi8>> loc(#loc11)

Here I load 4 tiles 16x16xi32, 4 tiles 16x32xi8, and 1 tile 8x64xi8. I understand that I have 9 live tiles in total and I've already fixed my codegen. But for the original code I got:

error: ran out of registers during register allocation

Should it be considered a bug or it's expected behavior?

[phoebewang]

BTW I also had a case where allocator couldn't allocate registers at all. Here is a snippet:

    %1588 = "amx.tileloadd64"(%5, %4, %1322, %3) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<16xi32>> loc(#loc11)
    %1589 = "amx.tileloadd64"(%5, %4, %1386, %3) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<16xi32>> loc(#loc11)
    %1590 = "amx.tileloadd64"(%5, %4, %1450, %3) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<16xi32>> loc(#loc11)
    %1591 = "amx.tileloadd64"(%5, %4, %1514, %3) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<16xi32>> loc(#loc11)
    %1592 = "amx.tileloadd64"(%5, %2, %951, %1) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<32xi8>> loc(#loc11)
    %1593 = "amx.tileloadd64"(%5, %2, %967, %1) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<32xi8>> loc(#loc11)
    %1594 = "amx.tileloadd64"(%5, %2, %983, %1) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<32xi8>> loc(#loc11)
    %1595 = "amx.tileloadd64"(%5, %2, %999, %1) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<16 x vector<32xi8>> loc(#loc11)
    %1596 = "amx.tileloadd64"(%0, %4, %1021, %3) : (i16, i16, !llvm.ptr, i64) -> !llvm.array<8 x vector<64xi8>> loc(#loc11)

Here I load 4 tiles 16x16xi32, 4 tiles 16x32xi8, and 1 tile 8x64xi8. I understand that I have 9 live tiles in total and I've already fixed my codegen. But for the original code I got:

error: ran out of registers during register allocation

Should it be considered a bug or it's expected behavior?

It's expected behavior. When you have different matrix shapes, make the unrolling even less, because registers cannot be reused for different shapes.