[X86] Support __tile_stream_loadd intrinsic for new AMX interface

[yubing]

Summary

Adding support for __tile_stream_loadd intrinsic.

[LuoYuanke]

llvm/lib/Target/X86/X86FastTileConfig.cpp

124

Also add the stream load for X86PreAMXConfig.cpp: isTileLoad().

llvm/lib/Target/X86/X86InstrAMX.td

57

indent.

[yubing]

Address yuanke's comments

[yubing]

Rebase

[LuoYuanke]

LGTM. Thank you!

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_tile_loadconfig_internal, "vvC*", "n", "amx-tile")
 TARGET_BUILTIN(__builtin_ia32_tileloadd64_internal, "V256iUsUsvC*z", "n", "amx-tile")
+TARGET_BUILTIN(__builtin_ia32_tileloaddt164_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

clang/lib/Headers/amxintrin.h

 _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));
 }
 
 /// This is internal intrinsic. C/C++ user should avoid calling it directly.
 static __inline__ _tile1024i __DEFAULT_FN_ATTRS_INT8
+_tile_loaddt1_internal(unsigned short m, unsigned short n, const void *base,
+                       __SIZE_TYPE__ stride) {
+  return __builtin_ia32_tileloaddt164_internal(m, n, base,
+                                               (__SIZE_TYPE__)(stride));
+}
+
+/// This is internal intrinsic. C/C++ user should avoid calling it directly.
+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);
 }
 
 /// This is internal intrinsic. C/C++ user should avoid calling it directly.
 static __inline__ _tile1024i __DEFAULT_FN_ATTRS_INT8
 _tile_dpbsud_internal(unsigned short m, unsigned short n, unsigned short k,

 /// \param stride
 ///    The stride between the rows' data to be loaded in memory.
 __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);
 }
 
+/// Load tile rows from memory specifieid by "base" address and "stride" into
+/// destination tile "dst". This intrinsic provides a hint to the implementation
+/// that the data will likely not be reused in the near future and the data
+/// caching can be optimized accordingly.
+///
+/// \headerfile <immintrin.h>
+///
+/// This intrinsic corresponds to the <c> TILELOADDT1 </c> instruction.
+///
+/// \param dst
+///    A destination tile. Max size is 1024 Bytes.
+/// \param base
+///    A pointer to base address.
+/// \param stride
+///    The stride between the rows' data to be loaded in memory.
+__DEFAULT_FN_ATTRS_TILE
+static void __tile_stream_loadd(__tile1024i *dst, const void *base,
+                                __SIZE_TYPE__ stride) {
+  dst->tile = _tile_loaddt1_internal(dst->row, dst->col, base, stride);
+}
+
 /// Compute dot-product of bytes in tiles with a source/destination accumulator.
 /// Multiply groups of 4 adjacent pairs of signed 8-bit integers in src0 with
 /// corresponding signed 8-bit integers in src1, producing 4 intermediate 32-bit
 /// results. Sum these 4 results with the corresponding 32-bit integer in "dst",
 /// and store the 32-bit result back to tile "dst".
 ///
 /// \headerfile <immintrin.h>
 ///

clang/test/CodeGen/X86/amx_api.c

 void test_tile_loadd(short row, short col) {
   //CHECK-LABEL: @test_tile_loadd
   //CHECK: call x86_amx @llvm.x86.tileloadd64.internal
   //CHECK-NEXT: {{%.*}} = bitcast x86_amx {{%.*}} to <256 x i32>
   __tile1024i a = {row, col};
   __tile_loadd(&a, buf, STRIDE);
 }
 
+void test_tile_stream_loadd(short row, short col) {
+  //CHECK-LABEL: @test_tile_stream_loadd
+  //CHECK: call x86_amx @llvm.x86.tileloaddt164.internal
+  //CHECK-NEXT: {{%.*}} = bitcast x86_amx {{%.*}} to <256 x i32>
+  __tile1024i a = {row, col};
+  __tile_stream_loadd(&a, buf, STRIDE);
+}
+
 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) {

llvm/include/llvm/IR/IntrinsicsX86.td

let TargetPrefix = "x86" in {
   def int_x86_ldtilecfg_internal :
               GCCBuiltin<"__builtin_ia32_tile_loadconfig_internal">,
               Intrinsic<[], [llvm_ptr_ty], []>;
   def int_x86_tileloadd64_internal :
               GCCBuiltin<"__builtin_ia32_tileloadd64_internal">,
               Intrinsic<[llvm_x86amx_ty],
                         [llvm_i16_ty, llvm_i16_ty, llvm_ptr_ty, llvm_i64_ty],
                         []>;
+  def int_x86_tileloaddt164_internal :
+              GCCBuiltin<"__builtin_ia32_tileloaddt164_internal">,
+              Intrinsic<[llvm_x86amx_ty],
+                        [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">,

llvm/lib/Target/X86/X86ExpandPseudo.cpp

bool X86ExpandPseudo::ExpandMI(MachineBasicBlock &MBB,
   }
   case TargetOpcode::ICALL_BRANCH_FUNNEL:
     ExpandICallBranchFunnel(&MBB, MBBI);
     return true;
   case X86::PLDTILECFGV: {
     MI.setDesc(TII->get(X86::LDTILECFG));
     return true;
   }
-  case X86::PTILELOADDV: {
+  case X86::PTILELOADDV:
+  case X86::PTILELOADDT1V: {
     for (unsigned i = 2; i > 0; --i)
       MI.RemoveOperand(i);
-    MI.setDesc(TII->get(X86::TILELOADD));
+    unsigned Opc =
+        Opcode == X86::PTILELOADDV ? X86::TILELOADD : X86::TILELOADDT1;
+    MI.setDesc(TII->get(Opc));
     return true;
   }
   case X86::PTDPBSSDV:
   case X86::PTDPBSUDV:
   case X86::PTDPBUSDV:
   case X86::PTDPBUUDV:
   case X86::PTDPBF16PSV: {
     MI.untieRegOperand(4);

llvm/lib/Target/X86/X86FastTileConfig.cpp

static inline void adjustRowCfg(unsigned TIdx, MachineInstr *MI) {
   MI->getOperand(3).ChangeToImmediate(Offset);
 }
 
 static inline void adjustColCfg(unsigned TIdx, MachineInstr *MI) {
   unsigned Offset = 16 + TIdx * 2;
   MI->getOperand(3).ChangeToImmediate(Offset);
 }
 
 bool X86FastTileConfig::isTileLoad(MachineInstr &MI) {
-  return MI.getOpcode() == X86::PTILELOADDV;
+  return MI.getOpcode() == X86::PTILELOADDV ||
+         MI.getOpcode() == X86::PTILELOADDT1V;
 }
 bool X86FastTileConfig::isTileStore(MachineInstr &MI) {
   return MI.getOpcode() == X86::PTILESTOREDV;
 }
 bool X86FastTileConfig::isAMXInstr(MachineInstr &MI) {
   // TODO: May need to handle some special nontile amx instrucion.
   if (MI.getOpcode() == X86::PLDTILECFGV || MI.isDebugInstr())
     return false;

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

case Intrinsic::x86_encodekey256: {
                                      Chain.getValue(1));
 
       MachineSDNode *Res = CurDAG->getMachineNode(
           Opcode, dl, Node->getVTList(),
           {Node->getOperand(2), Chain, Chain.getValue(1)});
       ReplaceNode(Node, Res);
       return;
     }
-    case Intrinsic::x86_tileloadd64_internal: {
+    case Intrinsic::x86_tileloadd64_internal:
+    case Intrinsic::x86_tileloaddt164_internal: {
       if (!Subtarget->hasAMXTILE())
         break;
-      unsigned Opc = X86::PTILELOADDV;
+      unsigned Opc = IntNo == Intrinsic::x86_tileloadd64_internal
+                         ? X86::PTILELOADDV
+                         : X86::PTILELOADDT1V;
       // _tile_loadd_internal(row, col, buf, STRIDE)
       SDValue Base = Node->getOperand(4);
       SDValue Scale = getI8Imm(1, dl);
       SDValue Index = Node->getOperand(5);
       SDValue Disp = CurDAG->getTargetConstant(0, dl, MVT::i32);
       SDValue Segment = CurDAG->getRegister(0, MVT::i16);
       SDValue Chain = Node->getOperand(0);
       MachineSDNode *CNode;

llvm/lib/Target/X86/X86InstrAMX.td

def TILEZERO : I<0x49, MRMr0, (outs TILE:$dst), (ins),
                      VEX, T8XD;
 
     // Pseduo instruction for RA.
     def PLDTILECFGV : PseudoI<(outs), (ins opaquemem:$src),
                               [(int_x86_ldtilecfg_internal addr:$src)]>;
     def PTILELOADDV : PseudoI<(outs TILE:$dst), (ins GR16:$src1,
                                                      GR16:$src2,
                                                      opaquemem:$src3), []>;
+    def PTILELOADDT1V : PseudoI<(outs TILE:$dst), (ins GR16:$src1,
+                                                       GR16:$src2,
+                                                       opaquemem:$src3), []>;
     def PTILESTOREDV : PseudoI<(outs), (ins GR16:$src1,
                                             GR16:$src2, opaquemem:$src3,
                                             TILE:$src4), []>;
     let isReMaterializable = 1, isAsCheapAsAMove = 1, canFoldAsLoad = 1 in
       def PTILEZEROV : PseudoI<(outs TILE:$dst), (ins GR16:$src1, GR16:$src2),
                                 [(set TILE:$dst, (int_x86_tilezero_internal
                                   GR16:$src1, GR16:$src2))]>;
 

llvm/lib/Target/X86/X86LowerAMXType.cpp

 
 std::pair<Value *, Value *> X86LowerAMXType::getShape(IntrinsicInst *II,
                                                       unsigned OpNo) {
   Value *Row = nullptr, *Col = nullptr;
   switch (II->getIntrinsicID()) {
   default:
     llvm_unreachable("Expect amx intrinsics");
   case Intrinsic::x86_tileloadd64_internal:
+  case Intrinsic::x86_tileloaddt164_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:

llvm/lib/Target/X86/X86PreAMXConfig.cpp

 static bool isAMXIntrinsic(IntrinsicInst *II) {
   for (Value *Operand : II->operands())
     if (Operand->getType()->isX86_AMXTy())
       return true;
   return II->getType()->isX86_AMXTy();
 }
 
 static bool isTileLoad(IntrinsicInst *II) {
-  return II->getIntrinsicID() == Intrinsic::x86_tileloadd64_internal;
+  return II->getIntrinsicID() == Intrinsic::x86_tileloadd64_internal ||
+         II->getIntrinsicID() == Intrinsic::x86_tileloaddt164_internal;
 }
 
 static bool isTileStore(IntrinsicInst *II) {
   return II->getIntrinsicID() == Intrinsic::x86_tilestored64_internal;
 }
 
 #ifndef NDEBUG
 static bool onlyTileDef(IntrinsicInst *II) {

llvm/lib/Target/X86/X86RegisterInfo.cpp

static ShapeT getTileShape(Register VirtReg, VirtRegMap *VRM,
   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::PTILELOADDT1V:
   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);

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

 ; 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:    tileloaddt1 (%rsi,%rdx), %tmm1
 ; 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)
   %d4 = call x86_amx @llvm.x86.tdpbf16ps.internal(i16 8, i16 8, i16 8, x86_amx %d3, x86_amx %a, x86_amx %b)
+  %e = call x86_amx @llvm.x86.tileloaddt164.internal(i16 8, i16 8, i8* %base, i64 %stride)
   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.tileloaddt164.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)