[X86] Support tilezero intrinsic and c interface for AMX.

[pengfei]

llvm/lib/Target/X86/X86ExpandPseudo.cpp

499

Can change to i = 3?

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

4650

Why tilezero needs chain?

[LuoYuanke]

llvm/lib/Target/X86/X86ExpandPseudo.cpp

499

Yes, then remove line 498.

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

4650

I didn't declare "IntrNoMem" in "include/llvm/IR/IntrinsicsX86.td". I can revise the code to declare "IntrNoMem".

[LuoYuanke]

Address Pengfei's comments.

[LuoYuanke]

a. Change tilezero back to be chained.

1. It avoid copy if several tilezero are same. Copy is expensive in AMX.
2. To keep the original order of amx intrinsics.

b. Refactor __tilezero interface.

[LuoYuanke]

Rebase.

[LuoYuanke]

Add avx512f in test case.

[pengfei]

LGTM.

[LuoYuanke]

Rebase.

clang/include/clang/Basic/BuiltinsX86_64.def

 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_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")
 
 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

 /*===--------------- amxintrin.h - AMX intrinsics -*- C/C++ -*---------------===
  *
  * Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
  * See https://llvm.org/LICENSE.txt for license information.
  * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
  *
  *===------------------------------------------------------------------------===
  */
 
 #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 __DEFAULT_FN_ATTRS_TILE                                                \
   __attribute__((__always_inline__, __nodebug__, __target__("amx-tile")))
 
 /// 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.
 ///
 /// \headerfile <x86intrin.h>
 ///
 /// This intrinsic corresponds to the <c> LDTILECFG </c> instruction.
 ///
 /// \param __config
 ///    A pointer to 512-bits configuration
 static __inline__ void __DEFAULT_FN_ATTRS_TILE
 _tile_loadconfig(const void *__config) {
   __builtin_ia32_tile_loadconfig(__config);
 }
 
 /// Stores the current tile configuration to 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 tiles
 /// are not configured, all zeroes will be stored to memory.
 ///
 /// \headerfile <x86intrin.h>
 ///
 /// This intrinsic corresponds to the <c> STTILECFG </c> instruction.
 ///
 /// \param __config
 ///    A pointer to 512-bits configuration
 static __inline__ void __DEFAULT_FN_ATTRS_TILE
 _tile_storeconfig(void *__config) {
   __builtin_ia32_tile_storeconfig(__config);
 }
 
 /// Release the tile configuration to return to the init state, which
 /// releases all storage it currently holds.
 ///
 /// \headerfile <x86intrin.h>
 ///
 /// This intrinsic corresponds to the <c> TILERELEASE </c> instruction.
 static __inline__ void __DEFAULT_FN_ATTRS_TILE _tile_release(void) {
   __builtin_ia32_tilerelease();
 }
 
 /// Load tile rows from memory specifieid by "base" address and "stride" into
 /// destination tile "dst" using the tile configuration previously configured
 /// via "_tile_loadconfig".
 ///
 /// \headerfile <x86intrin.h>
 ///

 
 #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
 _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__ 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 {
   const unsigned short row;
   const unsigned short col;
   _tile1024i tile;
 } __tile1024i;
 
-__DEFAULT_FN_ATTRS_INT8
+__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_dpbsud(__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
+__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_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_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_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],
+                        []>;
 }
 
 //===----------------------------------------------------------------------===//
 // 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

bool X86ExpandPseudo::ExpandMI(MachineBasicBlock &MBB,
   }
   case X86::PTILESTOREDV: {
     MI.RemoveOperand(8); // Remove $tmmcfg
     for (int i = 1; i >= 0; --i)
       MI.RemoveOperand(i);
     MI.setDesc(TII->get(X86::TILESTORED));
     return true;
   }
+  case X86::PTILEZEROV: {
+    for (int i = 3; i > 0; --i) // Remove row, col, $tmmcfg
+      MI.RemoveOperand(i);
+    MI.setDesc(TII->get(X86::TILEZERO));
+    return true;
+  }
   }
   llvm_unreachable("Previous switch has a fallthrough?");
 }
 
 /// Expand all pseudo instructions contained in \p MBB.
 /// \returns true if any expansion occurred for \p MBB.
 bool X86ExpandPseudo::ExpandMBB(MachineBasicBlock &MBB) {
   bool Modified = false;

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

case Intrinsic::x86_tdpbssd_internal: {
                        Node->getOperand(7),
                        CFG,
                        Chain};
       MachineSDNode *CNode =
           CurDAG->getMachineNode(Opc, dl, {MVT::x86amx, MVT::Other}, Ops);
       ReplaceNode(Node, CNode);
       return;
     }
+    case Intrinsic::x86_tilezero_internal: {
+      if (!Subtarget->hasAMXTILE())
+        break;
+      unsigned Opc = X86::PTILEZEROV;
+      SDValue Chain = Node->getOperand(0);
+      SDValue CFG = CurDAG->getRegister(0, MVT::Untyped);
+      SDValue Ops[] = {Node->getOperand(2), Node->getOperand(3), CFG, Chain};
+      MachineSDNode *CNode =
+          CurDAG->getMachineNode(Opc, dl, {MVT::x86amx, MVT::Other}, Ops);
+      ReplaceNode(Node, CNode);
+      return;
+    }
     }
     break;
   }
   case ISD::INTRINSIC_VOID: {
     unsigned IntNo = Node->getConstantOperandVal(1);
     switch (IntNo) {
     default: break;
     case Intrinsic::x86_sse3_monitor:

llvm/lib/Target/X86/X86InstrAMX.td

let SchedRW = [WriteSystem] in {
 
     def PTILELOADDV : PseudoI<(outs TILE: $dst), (ins GR16:$src1,
                                                       GR16:$src2,
                                                       opaquemem:$src3,
                                                       TILECFG:$cfg), []>;
     def PTILESTOREDV : PseudoI<(outs), (ins GR16:$src1,
                                             GR16:$src2, opaquemem:$src3,
                                             TILE:$src4, TILECFG:$cfg), []>;
+    def PTILEZEROV : PseudoI<(outs TILE: $dst), (ins GR16:$src1,
+                                                     GR16:$src2,
+                                                     TILECFG:$cfg), []>;
 
     let usesCustomInserter = 1 in {
       // Pseudo instructions, using immediates instead of tile registers.
       // To be translated to the actual instructions in X86ISelLowering.cpp
       def PTILELOADD : PseudoI<(outs), (ins u8imm:$src1, sibmem:$src2), []>;
       def PTILELOADDT1 : PseudoI<(outs), (ins u8imm:$src1,
                                           sibmem:$src2), []>;
       def PTILESTORED : PseudoI<(outs), (ins i8mem:$dst, u8imm:$src), []>;

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::PTILEZEROV:
     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() {
for (MachineBasicBlock &MBB : MF) {
     for (MachineInstr &MI : MBB) {
       unsigned Opcode = MI.getOpcode();
       switch (Opcode) {
       default:
         break;
       case X86::PTILELOADDV:
       case X86::PTILESTOREDV:
       case X86::PTDPBSSDV:
+      case X86::PTILEZEROV:
         unsigned NumOperands = MI.getNumOperands();
         MI.RemoveOperand(NumOperands - 1);
         MI.addOperand(MF, MachineOperand::CreateReg(CFG, false));
         break;
       }
     }
   }
 }

llvm/lib/Target/X86/X86RegisterInfo.cpp

static ShapeT getTileShape(Register VirtReg, VirtRegMap *VRM,
   unsigned OpCode = MI->getOpcode();
   switch (OpCode) {
   default:
     llvm_unreachable("Unexpected machine instruction on tile register!");
     break;
   // We only collect the tile shape that is defined.
   case X86::PTILELOADDV:
   case X86::PTDPBSSDV:
+  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

+; 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
+
+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:    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)
+
+  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 void @llvm.x86.tilestored64.internal(i16, i16, i8*, i64, x86_amx)