The AMX Vector Dialect is now ready for review. It is a rather large drop, so below is a brief description. Everything seems to work, including three full integration tests (running on a Sapphire Rapids emulator), but the documentation around AMX and its intrinsics is sometimes a bit “sparse” so please let me know if you encounter issues. It tried to make this a nice demonstration of the power of MLIR’s progressive lowering from pleasant high-level VV abstractions into efficient, but much more detailed low-level HWV code.
(1) The AMX dialect in MLIR itself is very concise, it consists of just a few new vector operations.
amx.tilezero
amx.mulf/muli
amx.tileload/store
All these operations work on concepts familiar within the MLIR framework: 2-d vectors and memrefs. For example, clearing a tile within an enveloping buffer can be done as follows.
%1 = amx.tilezero : vector<16x16xi32>
amx.tilestore %arg0[%i, %j], %1 : memref<?x?xi32>, vector<16x16xi32>
The idiomatic TMUL operations look as follows.
%0 = amx.tilemulf %a, %b, %c
: vector<16x32xbf16>, vector<16x32xbf16>, vector<16x16xf32>
The dialect verifies that the types and shapes are actually supported by AMX.
(2) The AMX dialect is lowered into an LLVMAMX dialect, which is closer to the compiler-oriented “internal” intrinsics of LLVM IR. The lowering takes care of some tedious details, such as providing tile parameters, stride computations, and instruction selection. The “amx.multf” is lowered to ‘tdpbf16ps’ and “amx.multi” to one of the ‘tdpbssd’, ‘tdpbsud’, ‘tdpbusd’, or ‘tdpbuud’ intrinsics. All this is completely transparent to the higher levels, though.
The example above lowers to the following.
%45 = "llvm_amx.tdpbf16ps"(%43, %42, %44, %40, %28, %28) : (i16, i16, i16, !llvm.array<16 x vector<16xf32>>, !llvm.array<16 x vector<32xbf16>>, !llvm.array<16 x vector<32xbf16>>) -> !llvm.array<16 x vector<16xf32>>
(3) This dialect is further lowered into LLVM IR dialect. At this point, the x86_amx type comes into play.
%47 = call x86_amx @llvm.x86.tdpbf16ps.internal(i16 16, i16 64, i16 64, x86_amx %46, x86_amx %37, x86_amx %37)
(4) The LLVM IR dialect is eventually handed off the LLVM. The backend converts the intrinsics into instructions and adds the necessary tile configuration instructions.
A sample configuration set up and actual kernel instructions are illustrated below.
vpxord %zmm0, %zmm0, %zmm0
vmovdqu64 %zmm0, -64(%rsp)
movb $1, -64(%rsp)
movb $16, -16(%rsp)
movw $64, -48(%rsp)
movb $16, -15(%rsp)
movw $64, -46(%rsp)
ldtilecfg -64(%rsp)
movq 24(%rsp), %rax
movq 48(%rsp), %rcx
.Ltmp2:
addq %r8, %r8
movw $64, %dx
movw $16, %di
tileloadd (%rsi,%r8), %tmm0
shlq $2, %rcx
tileloadd (%rax,%rcx), %tmm1
tdpbf16ps %tmm0, %tmm0, %tmm1
tilestored %tmm1, (%rax,%rcx)
tilerelease
vzeroupper
retq
I added several integration tests to make sure everything works as expected (on the emulator).
For example, storing a 16x16 vector into e.g. a 19x19 buffer accounts for the proper stride for the enveloping sizes:
( 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1 )
( 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 )
( 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 )
Likewise, a matrix multiply of pairwise bf16 into f32 works as expected.
( 124, 144 ) = ( ( 1.0, 2.0, 3.0, 4.0 ), (( 9.0, 10.0, 11.0, 12.0 ),
( 308, 360 ) ) ( 5.0, 6.0, 7.0, 8.0 )) x (13.0, 14.0, 15.0, 16.0 )):
Please let me know if you have any feedback.
Happy TMULing!