[X86] Support amx-int8 intrinsic.

[LiuChen3]

Summary

Adding support for intrinsics of TDPBSUD/TDPBUSD/TDPBUUD.

[LiuChen3]

clang/include/clang/Basic/BuiltinsX86_64.def

106

Forgot to delete the blank line. I will delete these in next patch.

[pengfei]

LGTM.

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

4625

Nit: Maybe we can try to move the lowering of the AMX intrinsics to the td file as a follow up.

[LiuChen3]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

4625

Good idea. I will work on that.

clang/include/clang/Basic/BuiltinsX86_64.def

 TARGET_BUILTIN(__builtin_ia32_clui, "v", "n", "uintr")
 TARGET_BUILTIN(__builtin_ia32_stui, "v", "n", "uintr")
 TARGET_BUILTIN(__builtin_ia32_testui, "Uc", "n", "uintr")
 TARGET_BUILTIN(__builtin_ia32_senduipi, "vUWi", "n", "uintr")
 
 // 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")
 // 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")
 

clang/lib/Headers/amxintrin.h

 }
 
 static __inline__ _tile1024i __DEFAULT_FN_ATTRS_INT8
 _tile_dpbssd_internal(unsigned short m, unsigned short n, unsigned short k,
                       _tile1024i dst, _tile1024i src1, _tile1024i src2) {
   return __builtin_ia32_tdpbssd_internal(m, n, k, dst, src1, src2);
 }
 
+static __inline__ _tile1024i __DEFAULT_FN_ATTRS_INT8
+_tile_dpbsud_internal(unsigned short m, unsigned short n, unsigned short k,
+                      _tile1024i dst, _tile1024i src1, _tile1024i src2) {
+  return __builtin_ia32_tdpbsud_internal(m, n, k, dst, src1, src2);
+}
+
+static __inline__ _tile1024i __DEFAULT_FN_ATTRS_INT8
+_tile_dpbusd_internal(unsigned short m, unsigned short n, unsigned short k,
+                      _tile1024i dst, _tile1024i src1, _tile1024i src2) {
+  return __builtin_ia32_tdpbusd_internal(m, n, k, dst, src1, src2);
+}
+
+static __inline__ _tile1024i __DEFAULT_FN_ATTRS_INT8
+_tile_dpbuud_internal(unsigned short m, unsigned short n, unsigned short k,
+                      _tile1024i dst, _tile1024i src1, _tile1024i src2) {
+  return __builtin_ia32_tdpbuud_internal(m, n, k, dst, src1, src2);
+}
+
 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);
 }
 
 typedef struct __tile1024i_str {

 
 __DEFAULT_FN_ATTRS_INT8
 static void __tile_dpbssd(__tile1024i *dst, __tile1024i src1,
                           __tile1024i src2) {
   dst->tile = _tile_dpbssd_internal(src1.row, src2.col, src1.col, dst->tile,
                                     src1.tile, src2.tile);
 }
 
+__DEFAULT_FN_ATTRS_INT8
+static void __tile_dpbsud(__tile1024i *dst, __tile1024i src1,
+                          __tile1024i src2) {
+  dst->tile = _tile_dpbsud_internal(src1.row, src2.col, src1.col, dst->tile,
+                                    src1.tile, src2.tile);
+}
+
+__DEFAULT_FN_ATTRS_INT8
+static void __tile_dpbusd(__tile1024i *dst, __tile1024i src1,
+                          __tile1024i src2) {
+  dst->tile = _tile_dpbusd_internal(src1.row, src2.col, src1.col, dst->tile,
+                                    src1.tile, src2.tile);
+}
+
+__DEFAULT_FN_ATTRS_INT8
+static void __tile_dpbuud(__tile1024i *dst, __tile1024i src1,
+                          __tile1024i src2) {
+  dst->tile = _tile_dpbuud_internal(src1.row, src2.col, src1.col, dst->tile,
+                                    src1.tile, src2.tile);
+}
+
 __DEFAULT_FN_ATTRS_TILE
 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);
 }
 
 #endif /* __x86_64__ */
 #endif /* __AMXINTRIN_H */

clang/test/CodeGen/X86/amx_api.c

 
 void test_tile_dpbssd(__tile1024i a, __tile1024i b, __tile1024i c) {
   //CHECK-LABEL: @test_tile_dpbssd
   //CHECK: call x86_amx @llvm.x86.tdpbssd.internal
   //CHECK-NEXT: {{%.*}} = bitcast x86_amx {{%.*}} to <256 x i32>
   __tile_dpbssd(&c, a, b);
 }
 
+void test_tile_dpbsud(__tile1024i a, __tile1024i b, __tile1024i c) {
+  //CHECK-LABEL: @test_tile_dpbsud
+  //CHECK: call x86_amx @llvm.x86.tdpbsud.internal
+  //CHECK-NEXT: {{%.*}} = bitcast x86_amx {{%.*}} to <256 x i32>
+  __tile_dpbsud(&c, a, b);
+}
+
+void test_tile_dpbusd(__tile1024i a, __tile1024i b, __tile1024i c) {
+  //CHECK-LABEL: @test_tile_dpbusd
+  //CHECK: call x86_amx @llvm.x86.tdpbusd.internal
+  //CHECK-NEXT: {{%.*}} = bitcast x86_amx {{%.*}} to <256 x i32>
+  __tile_dpbusd(&c, a, b);
+}
+
+void test_tile_dpbuud(__tile1024i a, __tile1024i b, __tile1024i c) {
+  //CHECK-LABEL: @test_tile_dpbuud
+  //CHECK: call x86_amx @llvm.x86.tdpbuud.internal
+  //CHECK-NEXT: {{%.*}} = bitcast x86_amx {{%.*}} to <256 x i32>
+  __tile_dpbuud(&c, a, b);
+}
+
 void test_tile_stored(__tile1024i c) {
   //CHECK-LABEL: @test_tile_stored
   //CHECK: {{%.*}} = bitcast <256 x i32> {{%.*}} to x86_amx
   //CHECK-NEXT: call void @llvm.x86.tilestored64.internal
   __tile_stored(buf, STRIDE, c);
 }
 
 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);
 }

llvm/include/llvm/IR/IntrinsicsX86.td

def int_x86_tileloadd64_internal :
                         [llvm_i16_ty, llvm_i16_ty, llvm_ptr_ty, llvm_i64_ty],
                         []>;
   def int_x86_tdpbssd_internal :
               GCCBuiltin<"__builtin_ia32_tdpbssd_internal">,
               Intrinsic<[llvm_x86amx_ty],
                         [llvm_i16_ty, llvm_i16_ty, llvm_i16_ty,
                          llvm_x86amx_ty, llvm_x86amx_ty,
                          llvm_x86amx_ty], []>;
+  def int_x86_tdpbsud_internal :
+              GCCBuiltin<"__builtin_ia32_tdpbsud_internal">,
+              Intrinsic<[llvm_x86amx_ty],
+                        [llvm_i16_ty, llvm_i16_ty, llvm_i16_ty,
+                         llvm_x86amx_ty, llvm_x86amx_ty,
+                         llvm_x86amx_ty], []>;
+  def int_x86_tdpbusd_internal :
+              GCCBuiltin<"__builtin_ia32_tdpbusd_internal">,
+              Intrinsic<[llvm_x86amx_ty],
+                        [llvm_i16_ty, llvm_i16_ty, llvm_i16_ty,
+                         llvm_x86amx_ty, llvm_x86amx_ty,
+                         llvm_x86amx_ty], []>;
+  def int_x86_tdpbuud_internal :
+              GCCBuiltin<"__builtin_ia32_tdpbuud_internal">,
+              Intrinsic<[llvm_x86amx_ty],
+                        [llvm_i16_ty, llvm_i16_ty, llvm_i16_ty,
+                         llvm_x86amx_ty, llvm_x86amx_ty,
+                         llvm_x86amx_ty], []>;
   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],
                         []>;

llvm/lib/Target/X86/X86ExpandPseudo.cpp

case TargetOpcode::ICALL_BRANCH_FUNNEL:
     ExpandICallBranchFunnel(&MBB, MBBI);
     return true;
   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::PTDPBSSDV:
+  case X86::PTDPBSUDV:
+  case X86::PTDPBUSDV:
+  case X86::PTDPBUUDV: {
     MI.untieRegOperand(4);
     for (unsigned i = 3; i > 0; --i)
       MI.RemoveOperand(i);
-    MI.setDesc(TII->get(X86::TDPBSSD));
+    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;
+    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)
       MI.RemoveOperand(i);
     MI.setDesc(TII->get(X86::TILESTORED));
     return true;

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

case Intrinsic::x86_tileloadd64_internal: {
                        Index,
                        Disp,
                        Segment,
                        Chain};
       CNode = CurDAG->getMachineNode(Opc, dl, {MVT::x86amx, MVT::Other}, Ops);
       ReplaceNode(Node, CNode);
       return;
     }
-    case Intrinsic::x86_tdpbssd_internal: {
+
+    case Intrinsic::x86_tdpbssd_internal:
+    case Intrinsic::x86_tdpbsud_internal:
+    case Intrinsic::x86_tdpbusd_internal:
+    case Intrinsic::x86_tdpbuud_internal: {
       if (!Subtarget->hasAMXTILE())
         break;
       SDValue Chain = Node->getOperand(0);
-      unsigned Opc = X86::PTDPBSSDV;
+      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;
+      default: llvm_unreachable("Impossible intrinsic");
+      }
       SDValue Ops[] = {Node->getOperand(2),
                        Node->getOperand(3),
                        Node->getOperand(4),
                        Node->getOperand(5),
                        Node->getOperand(6),
                        Node->getOperand(7),
                        Chain};
       MachineSDNode *CNode =

llvm/lib/Target/X86/X86InstrAMX.td

let Constraints = "$src1 = $dst" in {
                       VEX_4V, T8PD;
       def TDPBUUD : I<0x5e, MRMSrcReg4VOp3, (outs TILE:$dst),
                       (ins TILE:$src1, TILE:$src2, TILE:$src3),
                       "tdpbuud\t{$src3, $src2, $dst|$dst, $src2, $src3}", []>,
                       VEX_4V, T8PS;
     }
 
     // Pseduo instruction for RA.
-    let Constraints = "$src4 = $dst" in
+    let Constraints = "$src4 = $dst" in {
       def PTDPBSSDV : PseudoI<(outs TILE: $dst), (ins GR16:$src1,
                               GR16:$src2, GR16:$src3, TILE:$src4,
                               TILE:$src5, TILE:$src6), []>;
+      def PTDPBSUDV : PseudoI<(outs TILE: $dst), (ins GR16:$src1,
+                              GR16:$src2, GR16:$src3, TILE:$src4,
+                              TILE:$src5, TILE:$src6), []>;
+      def PTDPBUSDV : PseudoI<(outs TILE: $dst), (ins GR16:$src1,
+                              GR16:$src2, GR16:$src3, TILE:$src4,
+                              TILE:$src5, TILE:$src6), []>;
+      def PTDPBUUDV : PseudoI<(outs TILE: $dst), (ins 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 PTDPBSSD : PseudoI<(outs), (ins u8imm:$src1,
                              u8imm:$src2, u8imm:$src3),
                              [(int_x86_tdpbssd timm:$src1,
                                timm:$src2, timm:$src3)]>;

llvm/lib/Target/X86/X86LowerAMXType.cpp

static std::pair<Value *, Value *> getShape(IntrinsicInst *II, unsigned OpNo) {
   case Intrinsic::x86_tileloadd64_internal:
   case Intrinsic::x86_tilestored64_internal: {
     Row = II->getArgOperand(0);
     Col = II->getArgOperand(1);
     break;
   }
   // a * b + c
   // The shape depends on which operand.
-  case Intrinsic::x86_tdpbssd_internal: {
+  case Intrinsic::x86_tdpbssd_internal:
+  case Intrinsic::x86_tdpbsud_internal:
+  case Intrinsic::x86_tdpbusd_internal:
+  case Intrinsic::x86_tdpbuud_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) {
   unsigned Opcode = MI.getOpcode();
   switch (Opcode) {
   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:
     MachineOperand &MO1 = const_cast<MachineOperand &>(MI.getOperand(1));
     MachineOperand &MO2 = const_cast<MachineOperand &>(MI.getOperand(2));
     ShapeT Shape(&MO1, &MO2, MRI);
     return Shape;
   }
 }
 

 
 static bool isAMXInstruction(MachineBasicBlock::iterator MII) {
   switch (MII->getOpcode()) {
   default:
     return false;
   case X86::PTILELOADDV:
   case X86::PTILESTOREDV:
   case X86::PTDPBSSDV:
+  case X86::PTDPBSUDV:
+  case X86::PTDPBUSDV:
+  case X86::PTDPBUUDV:
   case X86::PTILEZEROV:
     return true;
   }
 }
 
 struct BBInfo {
   bool HasAMX = false;
   bool HasCallBeforeAMX = false;

llvm/lib/Target/X86/X86RegisterInfo.cpp

case X86::COPY: {
     Register SrcReg = MI->getOperand(1).getReg();
     ShapeT Shape = getTileShape(SrcReg, VRM, MRI);
     VRM->assignVirt2Shape(VirtReg, Shape);
     return Shape;
   }
   // 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:
     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

 ; 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:    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)
-  %d = call x86_amx @llvm.x86.tdpbssd.internal(i16 8, i16 8, i16 8, x86_amx %c, x86_amx %a, x86_amx %b)
-  call void @llvm.x86.tilestored64.internal(i16 8, i16 8, i8* %pointer, i64 %stride, x86_amx %d)
+  %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)
 
   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 void @llvm.x86.tilestored64.internal(i16, i16, i8*, i64, x86_amx)