[X86] Pass to transform tdpbf16ps intrinsics to scalar operation.

[yubing]

Summary

In previous patch https://reviews.llvm.org/D93594, we only scalarize tilezero, tileload, tilestore and tiledpbssd. In this patch we scalarize tdpbf16ps intrinsic.

[yubing]

Rebase and add a testcase for dpbf16ps intrinsic.

[pengfei]

llvm/test/CodeGen/X86/AMX/amx-low-intrinsics.ll

213–214

Can we use a shuffle instruction?

[yubing]

Address comments above and refactor some code

[yubing]

Rebase and add a testcase.

[yubing]

Fix incorrect naming for dpbf16's bb

[yubing]

Modify some comments

[yubing]

just do a rebase

[pengfei]

llvm/lib/Target/X86/X86LowerAMXIntrinsics.cpp

406–412

Is it concise to use below?

template <Intrinsic::ID IntrID>
typename std::enable_if_t<
    IntrID == Intrinsic::x86_tdpbssd_internal ||
    IntrID == Intrinsic::x86_tdpbf16ps_internal, bool>
lowerTileDP(Instruction *TileDP);

407–408

Can we create vecC with <256 x float>?

430

better to use EltCF32 or CF32

437

ditto

438

Better to define a variable for it and reuse.

[yubing]

llvm/lib/Target/X86/X86LowerAMXIntrinsics.cpp

407–408

In fact, we are trying to find a bitcast whose operand is <256 x i32>, as shown in line229.

[yubing]

address Pengfei's comments

[pengfei]

LGTM.

[yubing]

Rebase after https://reviews.llvm.org/D98773 is merged.

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_tdpbf16ps_internal, "V256iUsUsUsV256iV256iV256i", "n", "amx-bf16")
 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

 /// \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")))
 
+#define __DEFAULT_FN_ATTRS_BF16                                                \
+  __attribute__((__always_inline__, __nodebug__, __target__("amx-bf16")))
+
 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__ _tile1024i __DEFAULT_FN_ATTRS_INT8
+_tile_dpbf16ps_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);
+}
+
 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_dpbf16ps(__tile1024i *dst, __tile1024i src1,
+                            __tile1024i src2) {
+  dst->tile = _tile_dpbf16ps_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 */

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_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], []>;
   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

bool X86ExpandPseudo::ExpandMI(MachineBasicBlock &MBB,
   case X86::PTDPBSSDV: {
     MI.untieRegOperand(4);
     for (unsigned i = 3; i > 0; --i)
       MI.RemoveOperand(i);
     MI.setDesc(TII->get(X86::TDPBSSD));
     MI.tieOperands(0, 1);
     return true;
   }
+  case X86::PTDPBF16PSV: {
+    MI.untieRegOperand(4);
+    for (unsigned i = 3; i > 0; --i)
+      MI.RemoveOperand(i);
+    MI.setDesc(TII->get(X86::TDPBF16PS));
+    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;
   }
   case X86::PTILEZEROV: {
     for (int i = 2; i > 0; --i) // Remove row, col

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

case Intrinsic::x86_tdpbssd_internal: {
                        Node->getOperand(6),
                        Node->getOperand(7),
                        Chain};
       MachineSDNode *CNode =
           CurDAG->getMachineNode(Opc, dl, {MVT::x86amx, MVT::Other}, Ops);
       ReplaceNode(Node, CNode);
       return;
     }
+    case Intrinsic::x86_tdpbf16ps_internal: {
+      if (!Subtarget->hasAMXTILE())
+        break;
+      SDValue Chain = Node->getOperand(0);
+      unsigned Opc = X86::PTDPBF16PSV;
+      SDValue Ops[] = {Node->getOperand(2),
+                       Node->getOperand(3),
+                       Node->getOperand(4),
+                       Node->getOperand(5),
+                       Node->getOperand(6),
+                       Node->getOperand(7),
+                       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 Ops[] = {Node->getOperand(2), Node->getOperand(3), Chain};
       MachineSDNode *CNode =
           CurDAG->getMachineNode(Opc, dl, {MVT::x86amx, MVT::Other}, Ops);

llvm/lib/Target/X86/X86InstrAMX.td

let SchedRW = [WriteSystem] in {
     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)]>;
     }
+    // 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), []>;
   }
 } // HasAMXTILE, HasAMXBF16

llvm/lib/Target/X86/X86LowerAMXIntrinsics.cpp

 //
 #include "X86.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/CodeGen/Passes.h"
-
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/ValueTypes.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/IntrinsicsX86.h"
if (IntrID == Intrinsic::x86_tileloadd64_internal) {
     B.SetInsertPoint(ColBody->getTerminator());
     Value *Elt = B.CreateExtractElement(Vec, Idx);
 
     B.CreateStore(Elt, EltPtr);
     return nullptr;
   }
 }
 
-static Value *createTileDPBSSDLoops(BasicBlock *Start, BasicBlock *End,
-                                    IRBuilderBase &B, DomTreeUpdater &DTU,
-                                    LoopInfo &LI, Value *Row, Value *Col,
-                                    Value *K, Value *Acc, Value *LHS,
-                                    Value *RHS) {
+template <Intrinsic::ID IntrID,
+          typename = typename std::enable_if<
+              IntrID == Intrinsic::x86_tdpbssd_internal ||
+              IntrID == Intrinsic::x86_tdpbf16ps_internal>::type>
+static Value *createTileDPLoops(BasicBlock *Start, BasicBlock *End,
+                                LoopInfo &LI, Value *Row, Value *Col, Value *K,
+                                Value *Acc, Value *LHS, Value *RHS) {
+  std::string IntrinName =
+      IntrID == Intrinsic::x86_tdpbssd_internal ? "tiledpbssd" : "tdpbf16ps";
   Loop *RowLoop = LI.AllocateLoop();
   Loop *ColLoop = LI.AllocateLoop();
   Loop *InnerLoop = LI.AllocateLoop();
   ColLoop->addChildLoop(InnerLoop);
   RowLoop->addChildLoop(ColLoop);
   if (Loop *ParentL = LI.getLoopFor(Start))
     ParentL->addChildLoop(RowLoop);
   else
     LI.addTopLevelLoop(RowLoop);
 
   BasicBlock *RowBody =
-      createLoop(Start, End, Row, B.getInt16(1), "tiledpbssd.scalarize.rows", B,
-                 DTU, RowLoop, LI);
+      createLoop(Start, End, Row, B.getInt16(1), IntrinName + ".scalarize.rows",
+                 B, DTU, RowLoop, LI);
   BasicBlock *RowLatch = RowBody->getSingleSuccessor();
 
   BasicBlock *ColBody =
       createLoop(RowBody, RowLatch, Col, B.getInt16(1),
-                 "tiledpbssd.scalarize.cols", B, DTU, ColLoop, LI);
+                 IntrinName + ".scalarize.cols", B, DTU, ColLoop, LI);
   BasicBlock *ColLoopLatch = ColBody->getSingleSuccessor();
 
   B.SetInsertPoint(ColBody->getTerminator());
   BasicBlock *InnerBody =
       createLoop(ColBody, ColLoopLatch, K, B.getInt16(1),
-                 "tiledpbssd.scalarize.inner", B, DTU, InnerLoop, LI);
+                 IntrinName + ".scalarize.inner", B, DTU, InnerLoop, LI);
 
   BasicBlock *ColumnLoopHeader = ColBody->getSinglePredecessor();
   BasicBlock *RowLoopHeader = RowBody->getSinglePredecessor();
   BasicBlock *InnerLoopHeader = InnerBody->getSinglePredecessor();
   BasicBlock *InnerLoopLatch = InnerBody->getSingleSuccessor();
   Value *CurrentRow = &*RowLoopHeader->begin();
   Value *CurrentCol = &*ColumnLoopHeader->begin();
   Value *CurrentInner = &*InnerLoopHeader->begin();
static Value *createTileDPLoops(BasicBlock *Start, BasicBlock *End,
   // %tiledpbssd.scalarize.cols.body ], [ %NewVecD,
   // %tiledpbssd.scalarize.inner.latch ]
   B.SetInsertPoint(InnerLoopHeader->getTerminator());
   PHINode *VecCPhi = B.CreatePHI(V256I32Ty, 2, "vec.c.inner.phi");
   VecCPhi->addIncoming(VecCPhiColLoop, ColBody);
   PHINode *VecDPhi = B.CreatePHI(V256I32Ty, 2, "vec.d.inner.phi");
   VecDPhi->addIncoming(VecDPhiColLoop, ColBody);
 
+  B.SetInsertPoint(InnerBody->getTerminator());
+  Value *IdxC =
+      B.CreateAdd(B.CreateMul(CurrentRow, B.getInt16(16)), CurrentCol);
+  Value *IdxA =
+      B.CreateAdd(B.CreateMul(CurrentRow, B.getInt16(16)), CurrentInner);
+  Value *IdxB =
+      B.CreateAdd(B.CreateMul(CurrentInner, B.getInt16(16)), CurrentCol);
+  Value *ResElt = nullptr;
+  if (IntrID == Intrinsic::x86_tdpbssd_internal) {
     // tiledpbssd.scalarize.inner.body:
     // calculate idxa, idxb, idxc
     // %eltc = extractelement <256 x i32> %vec.c.inner.phi, i16 %idxc
     // %elta = extractelement <256 x i32> %veca, i16 %idxa
     // %eltav4i8 = bitcast i32 %elta to <4 x i8>
     // %eltb = extractelement <256 x i32> %vecb, i16 %idxb
     // %eltbv4i8 = bitcast i32 %eltb to <4 x i8>
     // %eltav4i32 = sext <4 x i8> %eltav4i8 to <4 x i32>
     // %eltbv4i32 = sext <4 x i8> %eltbv4i8 to <4 x i32>
     // %mulab = mul <4 x i32> %eltbv4i32, %eltav4i32
     // %acc = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %131)
     // %neweltc = add i32 %elt, %acc
     // %NewVecC = insertelement <256 x i32> %vec.c.inner.phi, i32 %neweltc,
     // i16 %idxc
     // %NewVecD = insertelement <256 x i32> %vec.d.inner.phi, i32 %neweltc,
     // i16 %idxc
-  B.SetInsertPoint(InnerBody->getTerminator());
-  Value *IdxC =
-      B.CreateAdd(B.CreateMul(CurrentRow, B.getInt16(16)), CurrentCol);
-  Value *IdxA =
-      B.CreateAdd(B.CreateMul(CurrentRow, B.getInt16(16)), CurrentInner);
-  Value *IdxB =
-      B.CreateAdd(B.CreateMul(CurrentInner, B.getInt16(16)), CurrentCol);
-
     FixedVectorType *V4I8Ty = FixedVectorType::get(B.getInt8Ty(), 4);
     FixedVectorType *V4I32Ty = FixedVectorType::get(B.getInt32Ty(), 4);
     Value *EltC = B.CreateExtractElement(VecCPhi, IdxC);
     Value *EltA = B.CreateExtractElement(VecA, IdxA);
     Value *SubVecA = B.CreateBitCast(EltA, V4I8Ty);
     Value *EltB = B.CreateExtractElement(VecB, IdxB);
     Value *SubVecB = B.CreateBitCast(EltB, V4I8Ty);
     Value *SubVecR = B.CreateAddReduce(B.CreateMul(
         B.CreateSExt(SubVecA, V4I32Ty), B.CreateSExt(SubVecB, V4I32Ty)));
-  Value *ResElt = B.CreateAdd(EltC, SubVecR);
+    ResElt = B.CreateAdd(EltC, SubVecR);
+  } else if (IntrID == Intrinsic::x86_tdpbf16ps_internal) {
+    // tiledpbf16ps.scalarize.inner.body:
+    // calculate idxa, idxb, idxc
+    // %eltc = extractelement <256 x i32> %vec.c.inner.phi, i16 %idxc
+    // %elta = extractelement <256 x i32> %veca, i16 %idxa
+    // %eltav4i8 = bitcast i32 %elta to <4 x i8>
+    // %eltb = extractelement <256 x i32> %vecb, i16 %idxb
+    // %eltbv4i8 = bitcast i32 %eltb to <4 x i8>
+    // %eltav4i32 = sext <4 x i8> %eltav4i8 to <4 x i32>
+    // %eltbv4i32 = sext <4 x i8> %eltbv4i8 to <4 x i32>
+    // %mulab = mul <4 x i32> %eltbv4i32, %eltav4i32
+    // %acc = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %mulab)
+    // %neweltc = add i32 %elt, %acc
+    // %NewVecC = insertelement <256 x i32> %vec.c.inner.phi, i32 %neweltc,
+    // i16 %idxc
+    // %NewVecD = insertelement <256 x i32> %vec.d.inner.phi, i32 %neweltc,
+    // i16 %idxc
+
+    // %eltc = extractelement <256 x i32> %vec.c.inner.phi, i16 %idxc
+    // %eltcf32 = bitcast i32 %eltc to float
+    // %elta = extractelement <256 x i32> %veca, i16 %idxa
+    // %eltav2i16 = bitcast i32 %elta to <2 x i16>
+    // %eltb = extractelement <256 x i32> %vecb, i16 %idxb
+    // %eltbv2i16 = bitcast i32 %eltb to <2 x i16>
+    // %shufflea = shufflevector <2 x i16> %elta, <2 x i16> zeroinitializer, <4
+    // x i32> <i32 2, i32 0, i32 3, i32 1> %eltav2f32 = bitcast <4 x i16>
+    // %shufflea to <2 x float>
+    // %shuffleb = shufflevector <2 x i16> %eltb, <2 xi16> zeroinitializer, <4 x
+    // i32> <i32 2, i32 0, i32 3, i32 1> %eltbv2f32 = bitcast <4 x i16>
+    // %shuffleb to <2 x float> %mulab = fmul <2 x float> %eltav2f32, %eltbv2f32
+    // %acc = call float
+    // @llvm.vector.reduce.fadd.v2f32(float %eltcf32, <2 x float> %mulab)
+    // %neweltc = bitcast float %acc to i32
+    // %NewVecC = insertelement <256 x i32> %vec.c.inner.phi, i32 %neweltc,
+    // i16 %idxc
+    // %NewVecD = insertelement <256 x i32> %vec.d.inner.phi, i32 %neweltc,
+    // i16 %idxc
+    FixedVectorType *V2I16Ty = FixedVectorType::get(B.getInt16Ty(), 2);
+    FixedVectorType *V2F32Ty = FixedVectorType::get(B.getFloatTy(), 2);
+    Value *EltC = B.CreateExtractElement(VecCPhi, IdxC);
+    Value *C_F32 = B.CreateBitCast(EltC, B.getFloatTy());
+    Value *EltA = B.CreateExtractElement(VecA, IdxA);
+    Value *SubVecA = B.CreateBitCast(EltA, V2I16Ty);
+    Value *EltB = B.CreateExtractElement(VecB, IdxB);
+    Value *SubVecB = B.CreateBitCast(EltB, V2I16Ty);
+    Value *ZeroV2I16 = Constant::getNullValue(V2I16Ty);
+    int ShuffleMask[4] = {2, 0, 3, 1};
+    Value *A_V2F32 = B.CreateBitCast(
+        B.CreateShuffleVector(SubVecA, ZeroV2I16, makeArrayRef(ShuffleMask)),
+        V2F32Ty);
+    Value *B_V2F32 = B.CreateBitCast(
+        B.CreateShuffleVector(SubVecB, ZeroV2I16, makeArrayRef(ShuffleMask)),
+        V2F32Ty);
+    Value *SubVecR = B.CreateFAddReduce(C_F32, B.CreateFMul(A_V2F32, B_V2F32));
+    ResElt = B.CreateBitCast(SubVecR, B.getInt32Ty());
+  } else {
+    llvm_unreachable("it is not a tdpb intrinsic");
+  }
   Value *NewVecC = B.CreateInsertElement(VecCPhi, ResElt, IdxC);
   Value *NewVecD = B.CreateInsertElement(VecDPhi, ResElt, IdxC);
 
   VecCPhi->addIncoming(NewVecC, InnerLoopLatch);
   VecCPhiRowLoop->addIncoming(NewVecC, RowLatch);
   VecCPhiColLoop->addIncoming(NewVecC, ColLoopLatch);
   VecDPhi->addIncoming(NewVecD, InnerLoopLatch);
   VecDPhiRowLoop->addIncoming(NewVecD, RowLatch);
private:
   DominatorTree *DT;
   LoopInfo *LI;
   template <Intrinsic::ID IntrID,
             typename = typename std::enable_if<
                 IntrID == Intrinsic::x86_tileloadd64_internal ||
                 IntrID == Intrinsic::x86_tilestored64_internal>::type>
   bool lowerTileLoadStore(Instruction *TileLoad);
   bool lowerTileLoad(Instruction *TileLoad);
-  bool lowerTileDPBSSD(Instruction *TileDPBSSD);
+  template <Intrinsic::ID IntrID,
+            typename = typename std::enable_if<
+                IntrID == Intrinsic::x86_tdpbssd_internal ||
+                IntrID == Intrinsic::x86_tdpbf16ps_internal>::type>
+  bool lowerTileDP(Instruction *TileDP);
   bool lowerTileStore(Instruction *TileStore);
   bool lowerTileZero(Instruction *TileZero);
 };
 
-bool X86LowerAMXIntrinsics::lowerTileDPBSSD(Instruction *TileDPBSSD) {
+template <Intrinsic::ID IntrID,
+          typename = typename std::enable_if<
+              IntrID == Intrinsic::x86_tdpbssd_internal ||
+              IntrID == Intrinsic::x86_tdpbf16ps_internal>::type>
+bool X86LowerAMXIntrinsics::lowerTileDP(Instruction *TileDP) {
   Value *M, *N, *K, *C, *A, *B;
-  match(TileDPBSSD, m_Intrinsic<Intrinsic::x86_tdpbssd_internal>(
-                        m_Value(M), m_Value(N), m_Value(K), m_Value(C),
-                        m_Value(A), m_Value(B)));
+  match(TileDP, m_Intrinsic<IntrID>(m_Value(M), m_Value(N), m_Value(K),
+                                    m_Value(C), m_Value(A), m_Value(B)));
   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
-  Instruction *InsertI = TileDPBSSD;
-  IRBuilder<> PreBuilder(TileDPBSSD);
-  PreBuilder.SetInsertPoint(TileDPBSSD);
+  Instruction *InsertI = TileDP;
+  IRBuilder<> PreBuilder(TileDP);
+  PreBuilder.SetInsertPoint(TileDP);
   // We visit the loop with (m, n/4, k/4):
   // %n_dword = udiv i16 %n, 4
   // %k_dword = udiv i16 %k, 4
   Value *NDWord = PreBuilder.CreateLShr(N, PreBuilder.getInt16(2));
   Value *KDWord = PreBuilder.CreateLShr(K, PreBuilder.getInt16(2));
   BasicBlock *Start = InsertI->getParent();
   BasicBlock *End =
       SplitBlock(InsertI->getParent(), InsertI, DT, LI, nullptr, "continue");
-  IRBuilder<> Builder(TileDPBSSD);
-  Value *ResVec = createTileDPBSSDLoops(Start, End, Builder, DTU, *LI, M,
+  IRBuilder<> Builder(TileDP);
+  Value *ResVec = createTileDPLoops<IntrID>(Start, End, Builder, DTU, *LI, M,
                                             NDWord, KDWord, C, A, B);
-  // we cannot assume there always be bitcast after tiledpbssd. So we need to
+  // we cannot assume there always be bitcast after TileDP. So we need to
   // insert one bitcast as required
   Builder.SetInsertPoint(End->getFirstNonPHI());
   Value *ResAMX =
       Builder.CreateBitCast(ResVec, Type::getX86_AMXTy(Builder.getContext()));
-  // Delete tiledpbssd intrinsic and do some clean-up.
-  for (auto UI = TileDPBSSD->use_begin(), UE = TileDPBSSD->use_end();
-       UI != UE;) {
+  // Delete TileDP intrinsic and do some clean-up.
+  for (auto UI = TileDP->use_begin(), UE = TileDP->use_end(); UI != UE;) {
     Instruction *I = cast<Instruction>((UI++)->getUser());
     Value *Vec;
     if (match(I, m_BitCast(m_Value(Vec)))) {
       I->replaceAllUsesWith(ResVec);
       I->eraseFromParent();
     }
   }
-  TileDPBSSD->replaceAllUsesWith(ResAMX);
-  TileDPBSSD->eraseFromParent();
+  TileDP->replaceAllUsesWith(ResAMX);
+  TileDP->eraseFromParent();
   return true;
 }
 
 template <Intrinsic::ID IntrID,
           typename = typename std::enable_if<
               IntrID == Intrinsic::x86_tileloadd64_internal ||
               IntrID == Intrinsic::x86_tilestored64_internal>::type>
 bool X86LowerAMXIntrinsics::lowerTileLoadStore(Instruction *TileLoad) {
bool X86LowerAMXIntrinsics::visit() {
   for (BasicBlock *BB : depth_first(&Func)) {
     for (BasicBlock::iterator II = BB->begin(), IE = BB->end(); II != IE;) {
       if (auto *Inst = dyn_cast<IntrinsicInst>(&*II++)) {
         switch (Inst->getIntrinsicID()) {
         case Intrinsic::x86_tdpbssd_internal:
         case Intrinsic::x86_tileloadd64_internal:
         case Intrinsic::x86_tilestored64_internal:
         case Intrinsic::x86_tilezero_internal:
+        case Intrinsic::x86_tdpbf16ps_internal:
           WorkList.push_back(Inst);
           break;
         default:
           break;
         }
       }
     }
   }
 
   for (auto *Inst : WorkList) {
     switch (Inst->getIntrinsicID()) {
     case Intrinsic::x86_tdpbssd_internal:
-      C = lowerTileDPBSSD(Inst) || C;
+      C = lowerTileDP<Intrinsic::x86_tdpbssd_internal>(Inst) || C;
+      break;
+    case Intrinsic::x86_tdpbf16ps_internal:
+      C = lowerTileDP<Intrinsic::x86_tdpbf16ps_internal>(Inst) || C;
       break;
     case Intrinsic::x86_tileloadd64_internal:
       C = lowerTileLoadStore<Intrinsic::x86_tileloadd64_internal>(Inst) || C;
       break;
     case Intrinsic::x86_tilestored64_internal:
       C = lowerTileLoadStore<Intrinsic::x86_tilestored64_internal>(Inst) || C;
       break;
     case Intrinsic::x86_tilezero_internal:

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_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) {
   unsigned Opcode = MI.getOpcode();
   switch (Opcode) {
   default:
     llvm_unreachable("Unexpected machine instruction on tile");
   case X86::PTILELOADDV:
   case X86::PTDPBSSDV:
+  case X86::PTDPBF16PSV:
   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::PTDPBF16PSV:
   case X86::PTILEZEROV:
     return true;
   }
 }
 
 struct BBInfo {
   bool HasAMX = false;
   bool HasCallBeforeAMX = false;

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::PTDPBF16PSV:
   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-low-intrinsics.ll

 ;
 entry:
   %amx = call x86_amx @llvm.x86.tileloadd64.internal(i16 %row, i16 %col, i8* %ptr, i64 %stride)
   %vec = bitcast x86_amx %amx to <256 x i32>
   store <256 x i32> %vec, <256 x i32>* %vptr, align 64
   ret void
 }
 
-define dso_local void @test_amx_dp(i16 signext %row, i16 signext %col, i16 signext %k, <256 x i32> %c, <256 x i32> %a, <256 x i32> %b, <256 x i32>* %vptr) #0 {
-; CHECK-LABEL: @test_amx_dp(
+define dso_local void @test_amx_dpbssd(i16 signext %row, i16 signext %col, i16 signext %k, <256 x i32> %c, <256 x i32> %a, <256 x i32> %b, <256 x i32>* %vptr) #0 {
+; CHECK-LABEL: @test_amx_dpbssd(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[A_AMX:%.*]] = bitcast <256 x i32> [[A:%.*]] to x86_amx
 ; CHECK-NEXT:    [[B_AMX:%.*]] = bitcast <256 x i32> [[B:%.*]] to x86_amx
 ; CHECK-NEXT:    [[C_AMX:%.*]] = bitcast <256 x i32> [[C:%.*]] to x86_amx
 ; CHECK-NEXT:    [[TMP0:%.*]] = lshr i16 [[COL:%.*]], 2
 ; CHECK-NEXT:    [[TMP1:%.*]] = lshr i16 [[K:%.*]], 2
 ; CHECK-NEXT:    br label [[TILEDPBSSD_SCALARIZE_ROWS_HEADER:%.*]]
 ; CHECK:       tiledpbssd.scalarize.rows.header:
entry:
   %b.amx = bitcast <256 x i32> %b to x86_amx
   %c.amx = bitcast <256 x i32> %c to x86_amx
   %acc = call x86_amx @llvm.x86.tdpbssd.internal(i16 %row, i16 %col, i16 %k, x86_amx %c.amx, x86_amx %a.amx, x86_amx %b.amx)
   %vec = bitcast x86_amx %acc to <256 x i32>
   store <256 x i32> %vec, <256 x i32>* %vptr, align 64
   ret void
 }
 
+define dso_local void @test_amx_dpbf16ps(i16 signext %row, i16 signext %col, i16 signext %k, <256 x i32> %c, <256 x i32> %a, <256 x i32> %b, <256 x i32>* %vptr) #0 {
+; CHECK-LABEL: @test_amx_dpbf16ps(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A_AMX:%.*]] = bitcast <256 x i32> [[A:%.*]] to x86_amx
+; CHECK-NEXT:    [[B_AMX:%.*]] = bitcast <256 x i32> [[B:%.*]] to x86_amx
+; CHECK-NEXT:    [[C_AMX:%.*]] = bitcast <256 x i32> [[C:%.*]] to x86_amx
+; CHECK-NEXT:    [[TMP0:%.*]] = lshr i16 [[COL:%.*]], 2
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i16 [[K:%.*]], 2
+; CHECK-NEXT:    br label [[TDPBF16PS_SCALARIZE_ROWS_HEADER:%.*]]
+; CHECK:       tdpbf16ps.scalarize.rows.header:
+; CHECK-NEXT:    [[TDPBF16PS_SCALARIZE_ROWS_IV:%.*]] = phi i16 [ 0, [[ENTRY:%.*]] ], [ [[TDPBF16PS_SCALARIZE_ROWS_STEP:%.*]], [[TDPBF16PS_SCALARIZE_ROWS_LATCH:%.*]] ]
+; CHECK-NEXT:    [[VEC_C_PHI_ROW:%.*]] = phi <256 x i32> [ [[C]], [[ENTRY]] ], [ [[TMP21:%.*]], [[TDPBF16PS_SCALARIZE_ROWS_LATCH]] ]
+; CHECK-NEXT:    [[VEC_D_PHI_ROW:%.*]] = phi <256 x i32> [ zeroinitializer, [[ENTRY]] ], [ [[TMP22:%.*]], [[TDPBF16PS_SCALARIZE_ROWS_LATCH]] ]
+; CHECK-NEXT:    br label [[TDPBF16PS_SCALARIZE_ROWS_BODY:%.*]]
+; CHECK:       tdpbf16ps.scalarize.rows.body:
+; CHECK-NEXT:    br label [[TDPBF16PS_SCALARIZE_COLS_HEADER:%.*]]
+; CHECK:       tdpbf16ps.scalarize.cols.header:
+; CHECK-NEXT:    [[TDPBF16PS_SCALARIZE_COLS_IV:%.*]] = phi i16 [ 0, [[TDPBF16PS_SCALARIZE_ROWS_BODY]] ], [ [[TDPBF16PS_SCALARIZE_COLS_STEP:%.*]], [[TDPBF16PS_SCALARIZE_COLS_LATCH:%.*]] ]
+; CHECK-NEXT:    [[VEC_C_PHI_COL:%.*]] = phi <256 x i32> [ [[VEC_C_PHI_ROW]], [[TDPBF16PS_SCALARIZE_ROWS_BODY]] ], [ [[TMP21]], [[TDPBF16PS_SCALARIZE_COLS_LATCH]] ]
+; CHECK-NEXT:    [[VEC_D_PHI_COL:%.*]] = phi <256 x i32> [ [[VEC_D_PHI_ROW]], [[TDPBF16PS_SCALARIZE_ROWS_BODY]] ], [ [[TMP22]], [[TDPBF16PS_SCALARIZE_COLS_LATCH]] ]
+; CHECK-NEXT:    br label [[TDPBF16PS_SCALARIZE_COLS_BODY:%.*]]
+; CHECK:       tdpbf16ps.scalarize.cols.body:
+; CHECK-NEXT:    br label [[TDPBF16PS_SCALARIZE_INNER_HEADER:%.*]]
+; CHECK:       tdpbf16ps.scalarize.inner.header:
+; CHECK-NEXT:    [[TDPBF16PS_SCALARIZE_INNER_IV:%.*]] = phi i16 [ 0, [[TDPBF16PS_SCALARIZE_COLS_BODY]] ], [ [[TDPBF16PS_SCALARIZE_INNER_STEP:%.*]], [[TDPBF16PS_SCALARIZE_INNER_LATCH:%.*]] ]
+; CHECK-NEXT:    [[VEC_C_INNER_PHI:%.*]] = phi <256 x i32> [ [[VEC_C_PHI_COL]], [[TDPBF16PS_SCALARIZE_COLS_BODY]] ], [ [[TMP21]], [[TDPBF16PS_SCALARIZE_INNER_LATCH]] ]
+; CHECK-NEXT:    [[VEC_D_INNER_PHI:%.*]] = phi <256 x i32> [ [[VEC_D_PHI_COL]], [[TDPBF16PS_SCALARIZE_COLS_BODY]] ], [ [[TMP22]], [[TDPBF16PS_SCALARIZE_INNER_LATCH]] ]
+; CHECK-NEXT:    br label [[TDPBF16PS_SCALARIZE_INNER_BODY:%.*]]
+; CHECK:       tdpbf16ps.scalarize.inner.body:
+; CHECK-NEXT:    [[TMP2:%.*]] = mul i16 [[TDPBF16PS_SCALARIZE_ROWS_IV]], 16
+; CHECK-NEXT:    [[TMP3:%.*]] = add i16 [[TMP2]], [[TDPBF16PS_SCALARIZE_COLS_IV]]
+; CHECK-NEXT:    [[TMP4:%.*]] = mul i16 [[TDPBF16PS_SCALARIZE_ROWS_IV]], 16
+; CHECK-NEXT:    [[TMP5:%.*]] = add i16 [[TMP4]], [[TDPBF16PS_SCALARIZE_INNER_IV]]
+; CHECK-NEXT:    [[TMP6:%.*]] = mul i16 [[TDPBF16PS_SCALARIZE_INNER_IV]], 16
+; CHECK-NEXT:    [[TMP7:%.*]] = add i16 [[TMP6]], [[TDPBF16PS_SCALARIZE_COLS_IV]]
+; CHECK-NEXT:    [[TMP8:%.*]] = extractelement <256 x i32> [[VEC_C_INNER_PHI]], i16 [[TMP3]]
+; CHECK-NEXT:    [[TMP9:%.*]] = bitcast i32 [[TMP8]] to float
+; CHECK-NEXT:    [[TMP10:%.*]] = extractelement <256 x i32> [[A]], i16 [[TMP5]]
+; CHECK-NEXT:    [[TMP11:%.*]] = bitcast i32 [[TMP10]] to <2 x i16>
+; CHECK-NEXT:    [[TMP12:%.*]] = extractelement <256 x i32> [[B]], i16 [[TMP7]]
+; CHECK-NEXT:    [[TMP13:%.*]] = bitcast i32 [[TMP12]] to <2 x i16>
+; CHECK-NEXT:    [[TMP14:%.*]] = shufflevector <2 x i16> [[TMP11]], <2 x i16> zeroinitializer, <4 x i32> <i32 2, i32 0, i32 3, i32 1>
+; CHECK-NEXT:    [[TMP15:%.*]] = bitcast <4 x i16> [[TMP14]] to <2 x float>
+; CHECK-NEXT:    [[TMP16:%.*]] = shufflevector <2 x i16> [[TMP13]], <2 x i16> zeroinitializer, <4 x i32> <i32 2, i32 0, i32 3, i32 1>
+; CHECK-NEXT:    [[TMP17:%.*]] = bitcast <4 x i16> [[TMP16]] to <2 x float>
+; CHECK-NEXT:    [[TMP18:%.*]] = fmul <2 x float> [[TMP15]], [[TMP17]]
+; CHECK-NEXT:    [[TMP19:%.*]] = call float @llvm.vector.reduce.fadd.v2f32(float [[TMP9]], <2 x float> [[TMP18]])
+; CHECK-NEXT:    [[TMP20:%.*]] = bitcast float [[TMP19]] to i32
+; CHECK-NEXT:    [[TMP21]] = insertelement <256 x i32> [[VEC_C_INNER_PHI]], i32 [[TMP20]], i16 [[TMP3]]
+; CHECK-NEXT:    [[TMP22]] = insertelement <256 x i32> [[VEC_D_INNER_PHI]], i32 [[TMP20]], i16 [[TMP3]]
+; CHECK-NEXT:    br label [[TDPBF16PS_SCALARIZE_INNER_LATCH]]
+; CHECK:       tdpbf16ps.scalarize.inner.latch:
+; CHECK-NEXT:    [[TDPBF16PS_SCALARIZE_INNER_STEP]] = add i16 [[TDPBF16PS_SCALARIZE_INNER_IV]], 1
+; CHECK-NEXT:    [[TDPBF16PS_SCALARIZE_INNER_COND:%.*]] = icmp ne i16 [[TDPBF16PS_SCALARIZE_INNER_STEP]], [[TMP1]]
+; CHECK-NEXT:    br i1 [[TDPBF16PS_SCALARIZE_INNER_COND]], label [[TDPBF16PS_SCALARIZE_INNER_HEADER]], label [[TDPBF16PS_SCALARIZE_COLS_LATCH]]
+; CHECK:       tdpbf16ps.scalarize.cols.latch:
+; CHECK-NEXT:    [[TDPBF16PS_SCALARIZE_COLS_STEP]] = add i16 [[TDPBF16PS_SCALARIZE_COLS_IV]], 1
+; CHECK-NEXT:    [[TDPBF16PS_SCALARIZE_COLS_COND:%.*]] = icmp ne i16 [[TDPBF16PS_SCALARIZE_COLS_STEP]], [[TMP0]]
+; CHECK-NEXT:    br i1 [[TDPBF16PS_SCALARIZE_COLS_COND]], label [[TDPBF16PS_SCALARIZE_COLS_HEADER]], label [[TDPBF16PS_SCALARIZE_ROWS_LATCH]]
+; CHECK:       tdpbf16ps.scalarize.rows.latch:
+; CHECK-NEXT:    [[TDPBF16PS_SCALARIZE_ROWS_STEP]] = add i16 [[TDPBF16PS_SCALARIZE_ROWS_IV]], 1
+; CHECK-NEXT:    [[TDPBF16PS_SCALARIZE_ROWS_COND:%.*]] = icmp ne i16 [[TDPBF16PS_SCALARIZE_ROWS_STEP]], [[ROW:%.*]]
+; CHECK-NEXT:    br i1 [[TDPBF16PS_SCALARIZE_ROWS_COND]], label [[TDPBF16PS_SCALARIZE_ROWS_HEADER]], label [[CONTINUE:%.*]]
+; CHECK:       continue:
+; CHECK-NEXT:    [[TMP23:%.*]] = bitcast <256 x i32> [[TMP22]] to x86_amx
+; CHECK-NEXT:    store <256 x i32> [[TMP22]], <256 x i32>* [[VPTR:%.*]], align 64
+; CHECK-NEXT:    ret void
+;
+entry:
+  %a.amx = bitcast <256 x i32> %a to x86_amx
+  %b.amx = bitcast <256 x i32> %b to x86_amx
+  %c.amx = bitcast <256 x i32> %c to x86_amx
+  %acc = call x86_amx @llvm.x86.tdpbf16ps.internal(i16 %row, i16 %col, i16 %k, x86_amx %c.amx, x86_amx %a.amx, x86_amx %b.amx)
+  %vec = bitcast x86_amx %acc to <256 x i32>
+  store <256 x i32> %vec, <256 x i32>* %vptr, align 64
+  ret void
+}
+
 define dso_local void @test_amx_store(i16 signext %row, i16 signext %col, i8 *%ptr, i64 %stride, <256 x i32>* %vptr, <256 x i32> %vec) #0 {
 ; CHECK-LABEL: @test_amx_store(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[AMX:%.*]] = bitcast <256 x i32> [[VEC:%.*]] to x86_amx
 ; CHECK-NEXT:    [[TMP0:%.*]] = lshr i16 [[COL:%.*]], 2
 ; CHECK-NEXT:    [[TMP1:%.*]] = lshr i64 [[STRIDE:%.*]], 2
 ; CHECK-NEXT:    br label [[TILESTORE_SCALARIZE_ROWS_HEADER:%.*]]
 ; CHECK:       tilestore.scalarize.rows.header:
entry:
   %vec = bitcast x86_amx %amx to <256 x i32>
   store <256 x i32> %vec, <256 x i32>* %vptr, align 64
   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.tdpbf16ps.internal(i16, i16, i16, x86_amx, x86_amx, x86_amx)
 declare void @llvm.x86.tilestored64.internal(i16, i16, i8*, i64, x86_amx)
 
 attributes #0 = { noinline nounwind optnone }