[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

174–175 ↗

(On Diff #321675)

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

333–334

Can we create vecC with <256 x float>?

356

better to use EltCF32 or CF32

363

ditto

364

Better to define a variable for it and reuse.

427–430

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);

[yubing]

llvm/lib/Target/X86/X86LowerAMXIntrinsics.cpp

333–334

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

 /*===--------------- 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
  *
  *===------------------------------------------------------------------------===
  */

 /// \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

 //===------- X86ExpandPseudo.cpp - Expand pseudo instructions -------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
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

 //===- X86ISelDAGToDAG.cpp - A DAG pattern matching inst selector for X86 -===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
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

 //===- llvm/CodeGen/TileShapeInfo.h - ---------------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
static Value *createTileDPBSSDLoops(BasicBlock *Start, BasicBlock *End,
   Value *NewVecC = B.CreateInsertElement(VecCPhi, ResElt, IdxC);
   VecCPhi->addIncoming(NewVecC, InnerLoopLatch);
   VecPhi_Row_Loop->addIncoming(NewVecC, RowLatch);
   VecPhi_Col_Loop->addIncoming(NewVecC, ColLoopLatch);
 
   return NewVecC;
 }
 
+static Value *createTileDPBF16PSLoops(BasicBlock *Start, BasicBlock *End,
+                                    IRBuilderBase &B, DomTreeUpdater &DTU,
+                                    LoopInfo &LI, Value *Row, Value *Col,
+                                    Value *K, Value *Acc, Value *LHS,
+                                    Value *RHS) {
+  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), "tiledpbf16ps.unroll.rows", B,
+                 DTU, RowLoop, LI);
+  BasicBlock *RowLatch = RowBody->getSingleSuccessor();
+
+  BasicBlock *ColBody =
+      createLoop(RowBody, RowLatch, Col, B.getInt16(1),
+                 "tiledpbf16ps.unroll.cols", B, DTU, ColLoop, LI);
+  BasicBlock *ColLoopLatch = ColBody->getSingleSuccessor();
+
+  B.SetInsertPoint(ColBody->getTerminator());
+  BasicBlock *InnerBody =
+      createLoop(ColBody, ColLoopLatch, K, B.getInt16(1),
+                 "tiledpbf16ps.unroll.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();
+
+  FixedVectorType *V256I32Ty = FixedVectorType::get(B.getInt32Ty(), 256);
+  // Type *EltTy = V256I32Ty->getElementType();
+  Value *VecC, *VecA, *VecB;
+  if (auto BitCast = dyn_cast<BitCastInst>(Acc))
+    VecC = BitCast->getOperand(0);
+  assert(VecC->getType()->isVectorTy() && "bitcast from non-v256i32 to x86amx");
+  // TODO else create BitCast from x86amx to v256i32.
+  // Store x86amx to memory, and reload from memory
+  // to vector. However with -O0, it doesn't happen.
+  if (auto BitCast = dyn_cast<BitCastInst>(LHS))
+    VecA = BitCast->getOperand(0);
+  assert(VecA->getType()->isVectorTy() && "bitcast from non-v256i32 to x86amx");
+  if (auto BitCast = dyn_cast<BitCastInst>(RHS))
+    VecB = BitCast->getOperand(0);
+  assert(VecB->getType()->isVectorTy() && "bitcast from non-v256i32 to x86amx");
+
+  // tiledpbf16ps.unroll.rows.header:
+  // %vec.phi.rows = phi <256 x i32> [ %vec_c, %continue ], [ %NewVecC,
+  // %tiledpbf16ps.unroll.rows.latch ]
+  B.SetInsertPoint(RowLoopHeader->getTerminator());
+  PHINode *VecPhi_Row_Loop = B.CreatePHI(V256I32Ty, 2, "vec.phi.row");
+  VecPhi_Row_Loop->addIncoming(VecC, Start);
+
+  // tiledpbf16ps.unroll.cols.header:
+  // %vec.phi.cols = phi <256 x i32> [ %vec.phi.rows,
+  // %tiledpbf16ps.unroll.rows.body ], [ %NewVecC, %tiledpbf16ps.unroll.cols.latch ]
+  B.SetInsertPoint(ColumnLoopHeader->getTerminator());
+  PHINode *VecPhi_Col_Loop = B.CreatePHI(V256I32Ty, 2, "vec.phi.col");
+  VecPhi_Col_Loop->addIncoming(VecPhi_Row_Loop, RowBody);
+
+  // Generate PHI vector for C.
+  B.SetInsertPoint(InnerLoopHeader->getTerminator());
+  PHINode *VecCPhi = B.CreatePHI(V256I32Ty, 2, "vec.phi");
+  VecCPhi->addIncoming(VecPhi_Col_Loop, ColBody);
+
+  // Generate accmulate multiply in innerbody.
+  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 *V2I16Ty = FixedVectorType::get(B.getInt16Ty(), 2);
+  FixedVectorType *V2I32Ty = FixedVectorType::get(B.getInt32Ty(), 2);
+  FixedVectorType *V2F32Ty = FixedVectorType::get(B.getFloatTy(), 2);
+  FixedVectorType *V4I32Ty = FixedVectorType::get(B.getInt32Ty(), 4);
+  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 *A_V2F32 = B.CreateBitCast(B.CreateShl(B.CreateZExt(SubVecA, V2I32Ty), B.CreateVectorSplat(2,B.getInt32(16))), V2F32Ty);
+  Value *B_V2F32 = B.CreateBitCast(B.CreateShl(B.CreateZExt(SubVecB, V2I32Ty), B.CreateVectorSplat(2,B.getInt32(16))), V2F32Ty);
+  Value *SubVecR = B.CreateFAddReduce(C_F32, B.CreateFMul(A_V2F32, B_V2F32));
+  Value *ResElt = B.CreateBitCast(SubVecR, B.getInt32Ty());
+  Value *NewVecC = B.CreateInsertElement(VecCPhi, ResElt, IdxC);
+  VecCPhi->addIncoming(NewVecC, InnerLoopLatch);
+  VecPhi_Row_Loop->addIncoming(NewVecC, RowLatch);
+  VecPhi_Col_Loop->addIncoming(NewVecC, ColLoopLatch);
+
+  return NewVecC;
+}
+
 namespace {
 class X86LowerAMXIntrinsics {
   Function &Func;
 
 public:
   X86LowerAMXIntrinsics(Function &F, DominatorTree *DT, LoopInfo *LI)
       : Func(F), DT(DT), LI(LI) {}
   bool visit();
 
 private:
   DominatorTree *DT;
   LoopInfo *LI;
   bool lowerTileLoad(Instruction *TileLoad);
   bool lowerTileDPBSSD(Instruction *TileDPBSSD);
+  bool lowerTileDPBF16PS(Instruction *TileDPBSSD);
   bool lowerTileStore(Instruction *TileStore);
   bool lowerTileZero(Instruction *TileZero);
 };
 
 bool X86LowerAMXIntrinsics::lowerTileDPBSSD(Instruction *TileDPBSSD) {
   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)));
   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
if (match(I, m_BitCast(m_Value(Vec)))) {
       I->replaceAllUsesWith(ResVec);
       I->eraseFromParent();
     }
   }
   TileDPBSSD->eraseFromParent();
   return true;
 }
 
+bool X86LowerAMXIntrinsics::lowerTileDPBF16PS(Instruction *TileDPBF16PS) {
+  Value *M, *N, *K, *C, *A, *B;
+  match(TileDPBF16PS, m_Intrinsic<Intrinsic::x86_tdpbf16ps_internal>(
+                        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 = TileDPBF16PS;
+  IRBuilder<> Builder_Prepare(TileDPBF16PS);
+  Builder_Prepare.SetInsertPoint(TileDPBF16PS);
+  // We visit the loop with (m, n/4, k/4):
+  // %n_dword = udiv i16 %n, 4
+  // %k_dword = udiv i16 %k, 4
+  Value *N_DWord = Builder_Prepare.CreateUDiv(N, Builder_Prepare.getInt16(4));
+  Value *K_DWord = Builder_Prepare.CreateUDiv(K, Builder_Prepare.getInt16(4));
+  BasicBlock *Start = InsertI->getParent();
+  BasicBlock *End =
+      SplitBlock(InsertI->getParent(), InsertI, DT, LI, nullptr, "continue");
+  IRBuilder<> Builder(TileDPBF16PS);
+  Value *ResVec = createTileDPBF16PSLoops(Start, End, Builder, DTU, *LI, M,
+                                        N_DWord, K_DWord, C, A, B);
+
+  // Delete tileloadd6 intrinsic and bitcast instruction.
+  for (auto UI = TileDPBF16PS->use_begin(), UE = TileDPBF16PS->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();
+    }
+  }
+  TileDPBF16PS->eraseFromParent();
+  return true;
+}
+
 bool X86LowerAMXIntrinsics::lowerTileLoad(Instruction *TileLoad) {
   Value *M, *N, *Ptr, *Stride;
   match(TileLoad, m_Intrinsic<Intrinsic::x86_tileloadd64_internal>(
                       m_Value(M), m_Value(N), m_Value(Ptr), m_Value(Stride)));
   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
   Instruction *InsertI = TileLoad;
   IRBuilder<> Builder_Prepare(TileLoad);
   Builder_Prepare.SetInsertPoint(TileLoad);
bool X86LowerAMXIntrinsics::lowerTileZero(Instruction *TileZero) {
   }
   TileZero->eraseFromParent();
   return true;
 }
 
 bool X86LowerAMXIntrinsics::visit() {
   bool C;
   SmallVector<Instruction *, 8> TileDPBSSDs;
+  SmallVector<Instruction *, 8> TileDPBF16PSs;
   SmallVector<Instruction *, 8> TileLoads;
   SmallVector<Instruction *, 8> TileStores;
   SmallVector<Instruction *, 8> TileZeros;
 
   for (BasicBlock *BB : post_order(&Func)) {
     for (BasicBlock::reverse_iterator II = BB->rbegin(), IE = BB->rend();
          II != IE;) {
       Instruction &Inst = *II++;
       if (match(&Inst, m_Intrinsic<Intrinsic::x86_tdpbssd_internal>())) {
         // %amx1 = bitcast <256 x i32> %vec to x86_amx
         // %res = call x86_amx @llvm.x86.tdpbssd.internal(i16 m, i16 n, i16 k,
         //                                                x86_amx, %amx1, ...)
         // %vec2 = bitcast x86_amx %res to <256 x i32>
         TileDPBSSDs.push_back(&Inst);
+      } else if (match(&Inst, m_Intrinsic<Intrinsic::x86_tdpbf16ps_internal>())) {
+        // %amx1 = bitcast <256 x i32> %vec to x86_amx
+        // %res = call x86_amx @llvm.x86.tdpbf16ps.internal(i16 m, i16 n, i16 k,
+        //                                                x86_amx, %amx1, ...)
+        // %vec2 = bitcast x86_amx %res to <256 x i32>
+        TileDPBF16PSs.push_back(&Inst);
       } else if (match(&Inst,
                        m_Intrinsic<Intrinsic::x86_tileloadd64_internal>())) {
         // %17 = call x86_amx @llvm.x86.tileloadd64.internal(i16 %13, i16 %14,
         //                                                   i8* %15, i64 %16)
         // %18 = bitcast x86_amx %17 to <256 x i32>
         TileLoads.push_back(&Inst);
       } else if (match(&Inst,
                        m_Intrinsic<Intrinsic::x86_tilestored64_internal>())) {
bool X86LowerAMXIntrinsics::visit() {
   }
 
   for (auto *Inst : TileLoads) {
     C |= lowerTileLoad(Inst);
   }
   for (auto *Inst : TileDPBSSDs) {
     C |= lowerTileDPBSSD(Inst);
   }
+  for (auto *Inst : TileDPBF16PSs) {
+    C |= lowerTileDPBF16PS(Inst);
+  }
   for (auto *Inst : TileStores) {
     C |= lowerTileStore(Inst);
   }
   for (auto *Inst : TileZeros) {
     C |= lowerTileZero(Inst);
   }
 
   return C;

llvm/lib/Target/X86/X86LowerAMXType.cpp

 //===- llvm/CodeGen/TileShapeInfo.h - ---------------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
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

 //===-- X86PreTileConfig.cpp - Tile Register Configure---------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //

 
 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

 //===-- X86RegisterInfo.cpp - X86 Register Information --------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
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;
   }
 }