[X86] Support amx-bf16 intrinsic.

[LiuChen3]

Summary

Adding support for intrinsics of AMX-BF16.
This patch alse fix a bug that AMX-INT8 instructions will be selected with wrong
predicate.

[LiuChen3]

Adding back 'avx512f' to amx-tile-basic.ll

[pengfei]

clang/lib/Headers/amxintrin.h

281

Is there much value to differentiate the type? We are using the same AMX type in the builtins. What do you think? @LuoYuanke

[LuoYuanke]

clang/lib/Headers/amxintrin.h

281

My first though is that we can reuse __tile1024i for bf16 tile for 2 reasons.

1. We don't access the element of the tile.
2. The destination element of amx-int8 is int32 and the destination element of amx-bf16 is float32, the element size is the same.

[LiuChen3]

clang/lib/Headers/amxintrin.h

281

Does this means that user need to do explicitly type conversion?

[pengfei]

clang/lib/Headers/amxintrin.h

281

We don't allow user to use the tile in the structure. User should always use load/store intrinsic to pass there own data.

[LiuChen3]

Address Pengfei and Yuanke's comments. We don't need more tile type.

[pengfei]

LGTM.

[LiuChen3]

I don't know why pre-merge-checks failed. I can check-all successfully locally in redhat8. I don't have debian mainchine to reproduce this problem.

[yubing]

clang/lib/Headers/amxintrin.h

326

Should we align this with "tile_dpbssd" by renaming it wth "tile_dpbf16ps"?

[xiangzhangllvm]

+1 first, didn't see key problems.

clang/lib/Headers/amxintrin.h

326

Yes, "t" already means "tile"

clang/include/clang/Basic/BuiltinsX86_64.def

 // AMX internal builtin
 TARGET_BUILTIN(__builtin_ia32_tileloadd64_internal, "V256iUsUsvC*z", "n", "amx-tile")
 TARGET_BUILTIN(__builtin_ia32_tdpbssd_internal, "V256iUsUsUsV256iV256iV256i", "n", "amx-int8")
 TARGET_BUILTIN(__builtin_ia32_tdpbsud_internal, "V256iUsUsUsV256iV256iV256i", "n", "amx-int8")
 TARGET_BUILTIN(__builtin_ia32_tdpbusd_internal, "V256iUsUsUsV256iV256iV256i", "n", "amx-int8")
 TARGET_BUILTIN(__builtin_ia32_tdpbuud_internal, "V256iUsUsUsV256iV256iV256i", "n", "amx-int8")
 TARGET_BUILTIN(__builtin_ia32_tilestored64_internal, "vUsUsv*zV256i", "n", "amx-tile")
 TARGET_BUILTIN(__builtin_ia32_tilezero_internal, "V256iUsUs", "n", "amx-tile")
+TARGET_BUILTIN(__builtin_ia32_tdpbf16ps_internal, "V256iUsUsUsV256iV256iV256i", "n", "amx-bf16")
 // AMX
 TARGET_BUILTIN(__builtin_ia32_tile_loadconfig, "vvC*", "n", "amx-tile")
 TARGET_BUILTIN(__builtin_ia32_tile_storeconfig, "vvC*", "n", "amx-tile")
 TARGET_BUILTIN(__builtin_ia32_tilerelease, "v", "n", "amx-tile")
 TARGET_BUILTIN(__builtin_ia32_tilezero, "vUc", "n", "amx-tile")
 
 TARGET_BUILTIN(__builtin_ia32_tileloadd64, "vIUcvC*z", "n", "amx-tile")
 TARGET_BUILTIN(__builtin_ia32_tileloaddt164, "vIUcvC*z", "n", "amx-tile")

clang/lib/Headers/amxintrin.h

 #ifndef __IMMINTRIN_H
 #error "Never use <amxintrin.h> directly; include <immintrin.h> instead."
 #endif /* __IMMINTRIN_H */
 
 #ifndef __AMXINTRIN_H
 #define __AMXINTRIN_H
 #ifdef __x86_64__
 
+/* Define the default attributes for the functions in this file. */
 #define __DEFAULT_FN_ATTRS_TILE                                                \
   __attribute__((__always_inline__, __nodebug__, __target__("amx-tile")))
+#define __DEFAULT_FN_ATTRS_INT8                                                \
+  __attribute__((__always_inline__, __nodebug__, __target__("amx-int8")))
+#define __DEFAULT_FN_ATTRS_BF16                                                \
+  __attribute__((__always_inline__, __nodebug__, __target__("amx-bf16")))
 
 /// Load tile configuration from a 64-byte memory location specified by
 /// "mem_addr". The tile configuration includes the tile type palette, the
 /// number of bytes per row, and the number of rows. If the specified
 /// palette_id is zero, that signifies the init state for both the tile
 /// config and the tile data, and the tiles are zeroed. Any invalid
 /// configurations will result in #GP fault.
 ///

 ///    The destination tile. Max size is 1024 Bytes.
 /// \param src0
 ///    The 1st source tile. Max size is 1024 Bytes.
 /// \param src1
 ///    The 2nd source tile. Max size is 1024 Bytes.
 #define _tile_dpbf16ps(dst, src0, src1)                                        \
   __builtin_ia32_tdpbf16ps((dst), (src0), (src1))
 
-#define __DEFAULT_FN_ATTRS_INT8                                                \
-  __attribute__((__always_inline__, __nodebug__, __target__("amx-int8")))
-
 typedef int _tile1024i __attribute__((__vector_size__(1024), __aligned__(64)));
+
 static __inline__ _tile1024i __DEFAULT_FN_ATTRS_INT8
 _tile_loadd_internal(unsigned short m, unsigned short n, const void *base,
                      __SIZE_TYPE__ stride) {
   return __builtin_ia32_tileloadd64_internal(m, n, base,
                                              (__SIZE_TYPE__)(stride));
 }
 
 static __inline__ _tile1024i __DEFAULT_FN_ATTRS_INT8

 
 static __inline__ void __DEFAULT_FN_ATTRS_INT8
 _tile_stored_internal(unsigned short m, unsigned short n, void *base,
                       __SIZE_TYPE__ stride, _tile1024i tile) {
   return __builtin_ia32_tilestored64_internal(m, n, base,
                                               (__SIZE_TYPE__)(stride), tile);
 }
 
+static __inline__ _tile1024i __DEFAULT_FN_ATTRS_BF16
+_tile_tdpbf16ps_internal(unsigned short m, unsigned short n, unsigned short k,
+                         _tile1024i dst, _tile1024i src1, _tile1024i src2) {
+  return __builtin_ia32_tdpbf16ps_internal(m, n, k, dst, src1, src2);
+}
+
 typedef struct __tile1024i_str {
   const unsigned short row;
   const unsigned short col;
   _tile1024i tile;
 } __tile1024i;
 
 __DEFAULT_FN_ATTRS_TILE
 static void __tile_loadd(__tile1024i *dst, const void *base,
                          __SIZE_TYPE__ stride) {
   dst->tile = _tile_loadd_internal(dst->row, dst->col, base, stride);
 }
 
 __DEFAULT_FN_ATTRS_INT8
 static void __tile_dpbssd(__tile1024i *dst, __tile1024i src1,
                           __tile1024i src2) {
static void __tile_stored(void *base, __SIZE_TYPE__ stride, __tile1024i src) {
   _tile_stored_internal(src.row, src.col, base, stride, src.tile);
 }
 
 __DEFAULT_FN_ATTRS_TILE
 static void __tile_zero(__tile1024i *dst) {
   dst->tile = __builtin_ia32_tilezero_internal(dst->row, dst->col);
 }
 
+__DEFAULT_FN_ATTRS_BF16
+static void __tile_tdpbf16ps(__tile1024i *dst, __tile1024i src1,
+                             __tile1024i src2) {
+  dst->tile = _tile_tdpbf16ps_internal(src1.row, src2.col, src1.col, dst->tile,
+                                       src1.tile, src2.tile);
+}
+
+#undef __DEFAULT_FN_ATTRS_TILE
+#undef __DEFAULT_FN_ATTRS_INT8
+#undef __DEFAULT_FN_ATTRS_BF16
+
 #endif /* __x86_64__ */
 #endif /* __AMXINTRIN_H */

clang/test/CodeGen/X86/amx_api.c

-// RUN: %clang_cc1 %s -ffreestanding -triple=x86_64-unknown-unknown  -target-feature +avx512f  -target-feature +amx-int8  \
+// RUN: %clang_cc1 %s -flax-vector-conversions=none -ffreestanding -triple=x86_64-unknown-unknown  -target-feature +avx512f  -target-feature +amx-int8  \
 // RUN: -target-feature +amx-bf16 -emit-llvm -o - -Werror -pedantic | FileCheck %s --check-prefixes=CHECK
 
 #include <immintrin.h>
 
 char buf[1024];
 #define STRIDE 32
 
 char buf2[1024];

 }
 
 void test_tile_zero(__tile1024i c) {
   //CHECK-LABEL: @test_tile_zero
   //CHECK: call x86_amx @llvm.x86.tilezero.internal
   //CHECK-NEXT bitcast x86_amx {{%.*}} to <256 x i32>
   __tile_zero(&c);
 }
+
+void test_tile_tdpbf16ps(__tile1024i a, __tile1024i b, __tile1024i c) {
+  //CHECK-LABEL: @test_tile_tdpbf16ps
+  //CHECK: call x86_amx @llvm.x86.tdpbf16ps.internal
+  //CHECK-NEXT: {{%.*}} = bitcast x86_amx {{%.*}} to <256 x i32>
+  __tile_tdpbf16ps(&a, b, c);
+}

llvm/include/llvm/IR/IntrinsicsX86.td

let TargetPrefix = "x86" in {
   def int_x86_tilestored64_internal :
               GCCBuiltin<"__builtin_ia32_tilestored64_internal">,
               Intrinsic<[], [llvm_i16_ty, llvm_i16_ty, llvm_ptr_ty,
                              llvm_i64_ty, llvm_x86amx_ty], []>;
   def int_x86_tilezero_internal :
               GCCBuiltin<"__builtin_ia32_tilezero_internal">,
               Intrinsic<[llvm_x86amx_ty], [llvm_i16_ty, llvm_i16_ty],
                         []>;
+  def int_x86_tdpbf16ps_internal :
+              GCCBuiltin<"__builtin_ia32_tdpbf16ps_internal">,
+              Intrinsic<[llvm_x86amx_ty],
+                        [llvm_i16_ty, llvm_i16_ty, llvm_i16_ty,
+                         llvm_x86amx_ty, llvm_x86amx_ty,
+                         llvm_x86amx_ty], []>;
 }
 
 //===----------------------------------------------------------------------===//
 // UINTR - User Level Interrupt
 
 let TargetPrefix = "x86" in {
   def int_x86_clui : GCCBuiltin<"__builtin_ia32_clui">,
               Intrinsic<[], [], []>;
   def int_x86_stui : GCCBuiltin<"__builtin_ia32_stui">,
               Intrinsic<[], [], []>;
   def int_x86_testui : GCCBuiltin<"__builtin_ia32_testui">,
               Intrinsic<[llvm_i8_ty], [], []>;
   def int_x86_senduipi : GCCBuiltin<"__builtin_ia32_senduipi">,
               Intrinsic<[], [llvm_i64_ty], []>;
 }

llvm/lib/Target/X86/X86ExpandPseudo.cpp

case X86::PTILELOADDV: {
     for (unsigned i = 2; i > 0; --i)
       MI.RemoveOperand(i);
     MI.setDesc(TII->get(X86::TILELOADD));
     return true;
   }
   case X86::PTDPBSSDV:
   case X86::PTDPBSUDV:
   case X86::PTDPBUSDV:
-  case X86::PTDPBUUDV: {
+  case X86::PTDPBUUDV:
+  case X86::PTDPBF16PSV: {
     MI.untieRegOperand(4);
     for (unsigned i = 3; i > 0; --i)
       MI.RemoveOperand(i);
     unsigned Opc;
     switch (Opcode) {
-    case X86::PTDPBSSDV: Opc = X86::TDPBSSD; break;
-    case X86::PTDPBSUDV: Opc = X86::TDPBSUD; break;
-    case X86::PTDPBUSDV: Opc = X86::TDPBUSD; break;
-    case X86::PTDPBUUDV: Opc = X86::TDPBUUD; break;
+    case X86::PTDPBSSDV:   Opc = X86::TDPBSSD; break;
+    case X86::PTDPBSUDV:   Opc = X86::TDPBSUD; break;
+    case X86::PTDPBUSDV:   Opc = X86::TDPBUSD; break;
+    case X86::PTDPBUUDV:   Opc = X86::TDPBUUD; break;
+    case X86::PTDPBF16PSV: Opc = X86::TDPBF16PS; break;
     default: llvm_unreachable("Impossible Opcode!");
     }
     MI.setDesc(TII->get(Opc));
     MI.tieOperands(0, 1);
     return true;
   }
   case X86::PTILESTOREDV: {
     for (int i = 1; i >= 0; --i)

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

case Intrinsic::x86_tileloadd64_internal: {
       ReplaceNode(Node, CNode);
       return;
     }
 
     case Intrinsic::x86_tdpbssd_internal:
     case Intrinsic::x86_tdpbsud_internal:
     case Intrinsic::x86_tdpbusd_internal:
     case Intrinsic::x86_tdpbuud_internal: {
-      if (!Subtarget->hasAMXTILE())
+      if (!Subtarget->hasAMXINT8())
         break;
       SDValue Chain = Node->getOperand(0);
       unsigned Opc;
       switch (IntNo) {
       case Intrinsic::x86_tdpbssd_internal: Opc = X86::PTDPBSSDV; break;
       case Intrinsic::x86_tdpbsud_internal: Opc = X86::PTDPBSUDV; break;
       case Intrinsic::x86_tdpbusd_internal: Opc = X86::PTDPBUSDV; break;
       case Intrinsic::x86_tdpbuud_internal: Opc = X86::PTDPBUUDV; break;

llvm/lib/Target/X86/X86InstrAMX.td

 let Predicates = [HasAMXBF16, In64BitMode] in {
   let SchedRW = [WriteSystem] in {
     let Constraints = "$src1 = $dst" in
     def TDPBF16PS : I<0x5c, MRMSrcReg4VOp3, (outs TILE:$dst),
                       (ins TILE:$src1, TILE:$src2, TILE:$src3),
                       "tdpbf16ps\t{$src3, $src2, $dst|$dst, $src2, $src3}",
                       []>, VEX_4V, T8XS;
 
+    // Pseduo instruction for RA.
+    let Constraints = "$src4 = $dst" in
+      def PTDPBF16PSV : PseudoI<(outs TILE: $dst), (ins GR16:$src1,
+                                 GR16:$src2, GR16:$src3, TILE:$src4,
+                                 TILE:$src5, TILE:$src6),
+                                 [(set TILE: $dst,
+                                  (int_x86_tdpbf16ps_internal GR16:$src1,
+                                   GR16:$src2, GR16:$src3, TILE:$src4,
+                                   TILE:$src5, TILE:$src6))]>;
+
     let usesCustomInserter = 1 in {
       // Pseudo instructions, using immediates instead of tile registers.
       // To be translated to the actual instructions in X86ISelLowering.cpp
       def PTDPBF16PS : PseudoI<(outs), (ins u8imm:$src1,
                                u8imm:$src2, u8imm:$src3),
                                [(int_x86_tdpbf16ps timm:$src1,
                                  timm:$src2, timm:$src3)]>;
     }
   }
 } // HasAMXTILE, HasAMXBF16

llvm/lib/Target/X86/X86LowerAMXType.cpp

case Intrinsic::x86_tilestored64_internal: {
     Col = II->getArgOperand(1);
     break;
   }
   // a * b + c
   // The shape depends on which operand.
   case Intrinsic::x86_tdpbssd_internal:
   case Intrinsic::x86_tdpbsud_internal:
   case Intrinsic::x86_tdpbusd_internal:
-  case Intrinsic::x86_tdpbuud_internal: {
+  case Intrinsic::x86_tdpbuud_internal:
+  case Intrinsic::x86_tdpbf16ps_internal: {
     switch (OpNo) {
     case 3:
       Row = II->getArgOperand(0);
       Col = II->getArgOperand(1);
       break;
     case 4:
       Row = II->getArgOperand(0);
       Col = II->getArgOperand(2);

llvm/lib/Target/X86/X86PreTileConfig.cpp

static ShapeT getShape(const MachineInstr &MI, MachineRegisterInfo *MRI) {
   default:
     llvm_unreachable("Unexpected machine instruction on tile");
   case X86::PTILELOADDV:
   case X86::PTDPBSSDV:
   case X86::PTDPBSUDV:
   case X86::PTDPBUSDV:
   case X86::PTDPBUUDV:
   case X86::PTILEZEROV:
+  case X86::PTDPBF16PSV:
     MachineOperand &MO1 = const_cast<MachineOperand &>(MI.getOperand(1));
     MachineOperand &MO2 = const_cast<MachineOperand &>(MI.getOperand(2));
     ShapeT Shape(&MO1, &MO2, MRI);
     return Shape;
   }
 }
 
 MachineInstr *X86PreTileConfig::getTileConfigPoint() {
default:
     return false;
   case X86::PTILELOADDV:
   case X86::PTILESTOREDV:
   case X86::PTDPBSSDV:
   case X86::PTDPBSUDV:
   case X86::PTDPBUSDV:
   case X86::PTDPBUUDV:
   case X86::PTILEZEROV:
+  case X86::PTDPBF16PSV:
     return true;
   }
 }
 
 struct BBInfo {
   bool HasAMX = false;
   bool HasCallBeforeAMX = false;
   bool HasAMXBeforeCallInSuccs = false;

llvm/lib/Target/X86/X86RegisterInfo.cpp

static ShapeT getTileShape(Register VirtReg, VirtRegMap *VRM,
   }
   // We only collect the tile shape that is defined.
   case X86::PTILELOADDV:
   case X86::PTDPBSSDV:
   case X86::PTDPBSUDV:
   case X86::PTDPBUSDV:
   case X86::PTDPBUUDV:
   case X86::PTILEZEROV:
+  case X86::PTDPBF16PSV:
     MachineOperand &MO1 = MI->getOperand(1);
     MachineOperand &MO2 = MI->getOperand(2);
     ShapeT Shape(&MO1, &MO2, MRI);
     VRM->assignVirt2Shape(VirtReg, Shape);
     return Shape;
   }
 }
 

llvm/test/CodeGen/X86/AMX/amx-tile-basic.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+amx-tile -mattr=+avx512f -verify-machineinstrs | FileCheck %s
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+amx-tile,+amx-int8,+amx-bf16,+avx512f -verify-machineinstrs | FileCheck %s
 
 define void @test_amx(i8* %pointer, i8* %base, i64 %stride) {
 ; CHECK-LABEL: test_amx:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vpxord %zmm0, %zmm0, %zmm0
 ; CHECK-NEXT:    vmovdqu64 %zmm0, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movb $1, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movb $8, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movw $8, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movb $8, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movw $8, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movb $8, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movw $8, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    ldtilecfg -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movw $8, %ax
 ; CHECK-NEXT:    tilezero %tmm0
 ; CHECK-NEXT:    tileloadd (%rsi,%rdx), %tmm1
 ; CHECK-NEXT:    tileloadd (%rsi,%rdx), %tmm2
 ; CHECK-NEXT:    tdpbssd %tmm2, %tmm1, %tmm0
 ; CHECK-NEXT:    tdpbsud %tmm2, %tmm1, %tmm0
 ; CHECK-NEXT:    tdpbusd %tmm2, %tmm1, %tmm0
 ; CHECK-NEXT:    tdpbuud %tmm2, %tmm1, %tmm0
+; CHECK-NEXT:    tdpbf16ps %tmm2, %tmm1, %tmm0
 ; CHECK-NEXT:    tilestored %tmm0, (%rdi,%rdx)
 ; CHECK-NEXT:    tilerelease
 ; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    retq
   %c = call x86_amx @llvm.x86.tilezero.internal(i16 8, i16 8)
   %a = call x86_amx @llvm.x86.tileloadd64.internal(i16 8, i16 8, i8* %base, i64 %stride)
   %b = call x86_amx @llvm.x86.tileloadd64.internal(i16 8, i16 8, i8* %base, i64 %stride)
   %d0 = call x86_amx @llvm.x86.tdpbssd.internal(i16 8, i16 8, i16 8, x86_amx %c, x86_amx %a, x86_amx %b)
   %d1 = call x86_amx @llvm.x86.tdpbsud.internal(i16 8, i16 8, i16 8, x86_amx %d0, x86_amx %a, x86_amx %b)
   %d2 = call x86_amx @llvm.x86.tdpbusd.internal(i16 8, i16 8, i16 8, x86_amx %d1, x86_amx %a, x86_amx %b)
   %d3 = call x86_amx @llvm.x86.tdpbuud.internal(i16 8, i16 8, i16 8, x86_amx %d2, x86_amx %a, x86_amx %b)
-  call void @llvm.x86.tilestored64.internal(i16 8, i16 8, i8* %pointer, i64 %stride, x86_amx %d3)
+  %d4 = call x86_amx @llvm.x86.tdpbf16ps.internal(i16 8, i16 8, i16 8, x86_amx %d3, x86_amx %a, x86_amx %b)
+  call void @llvm.x86.tilestored64.internal(i16 8, i16 8, i8* %pointer, i64 %stride, x86_amx %d4)
 
   ret void
 }
 
 declare x86_amx @llvm.x86.tilezero.internal(i16, i16)
 declare x86_amx @llvm.x86.tileloadd64.internal(i16, i16, i8*, i64)
 declare x86_amx @llvm.x86.tdpbssd.internal(i16, i16, i16, x86_amx, x86_amx, x86_amx)
 declare x86_amx @llvm.x86.tdpbsud.internal(i16, i16, i16, x86_amx, x86_amx, x86_amx)
 declare x86_amx @llvm.x86.tdpbusd.internal(i16, i16, i16, x86_amx, x86_amx, x86_amx)
 declare x86_amx @llvm.x86.tdpbuud.internal(i16, i16, i16, x86_amx, x86_amx, x86_amx)
+declare x86_amx @llvm.x86.tdpbf16ps.internal(i16, i16, i16, x86_amx, x86_amx, x86_amx)
 declare void @llvm.x86.tilestored64.internal(i16, i16, i8*, i64, x86_amx)