[X86] Sink x86_amx load in AMX type lowering.

[LuoYuanke]

Summary

When transform x86_amx load to amx load intrinsics, compiler need
to know the shape the matrix. The shape can be deduced from the user
of the load. However the shape may be defined after the load
instruction, so we need to sink the load to avoid such issue. This patch
only support sink within a basic block.

[craig.topper]

llvm/lib/Target/X86/X86LowerAMXType.cpp

100

Does this prevent sinking across a call that may change the memory being loaded?

[craig.topper]

llvm/lib/Target/X86/X86LowerAMXType.cpp

100

Or atomics, or anything with side effects.

[LuoYuanke]

llvm/lib/Target/X86/X86LowerAMXType.cpp

100

Thank you for review. Yes. I want to prevent all the scenario that may change the memory being loaded. I'll check how many scenario we need to prevent from sinking.

[pengfei]

It's better to set parent since the patch is based on D91927.

llvm/lib/Target/X86/X86LowerAMXType.cpp

123–124

Why don't pass the pointer directly for LD and II?

[LuoYuanke]

llvm/lib/Target/X86/X86LowerAMXType.cpp

100

@craig.topper , I do some study, but I still don't understand why load can't sink across atomic instruction or side effect instruction. I notice in MergedLoadStoreMotion::isStoreSinkBarrierInRange(), it only check myThrow() and alias.

[pengfei]

llvm/lib/Target/X86/X86LowerAMXType.cpp

100

I only noticed we chain the memory operations during building the DAG and these side effect instructions. Do we have such rules for middle end passes?

[craig.topper]

llvm/lib/Target/X86/X86LowerAMXType.cpp

100

Doesn't isStoreSinkBarrierInRange call canInstructionRangeModRef not just isnoalias with a range including every instruction between the store and where it wants to be moved to?

"Release" ordering for an atomic means that no load/store can be moved below it, for example.

[LuoYuanke]

Address Craig and Pengfei's comments.

llvm/lib/Target/X86/X86LowerAMXType.cpp

 /// and only AMX intrinsics can operate on the type, we need transform
 /// load/store <256 x i32> instruction to AMX load/store. Otherwise we are not
 /// able to lower the bitcast instruction to X86 instruction.
 //
 //===----------------------------------------------------------------------===//
 //
 #include "X86.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/CodeGen/Passes.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"
case Intrinsic::x86_tdpbssd_internal: {
     }
     break;
   }
   }
 
   return std::make_pair(Row, Col);
 }
 
+// Start is load instruction. Check if the load can sink through
+// the following instruction. The sink don't across basic block.
+// TODO: improve the sink across basic block.
+static Instruction *findSafePointToSink(Instruction *LD, Instruction *User,
+                                        AliasAnalysis *AA) {
+  Instruction *Start = &*(++LD->getIterator());
+  Instruction *End = nullptr;
+  BasicBlock *BB = LD->getParent();
+  if (LD->getParent() == User->getParent())
+    End = User;
+  else
+    End = BB->getTerminator();
+
+  MemoryLocation Loc = MemoryLocation::get(LD);
+  for (Instruction &Inst :
+       make_range(Start->getIterator(), End->getIterator())) {
+    if (Inst.mayThrow())
+      return &Inst;
+
+    if (isModOrRefSet(
+            intersectModRef(AA->getModRefInfo(&Inst, Loc), ModRefInfo::Mod)))
+      return &Inst;
+  }
+
+  return End;
+}
+
 // %1 = load x86_amx, x86_amx* %0, align 64
 // %2 = call x86_amx @llvm.x86.tdpbssd.internal(%1, %1, %1, ...)
 // -->
 // %1 = call x86_amx @llvm.x86.tileloadd64.internal()
 // %2 = call x86_amx @llvm.x86.tdpbssd.internal(%1, %1, %1, ...)
-static void transformTileLoad(LoadInst *LD) {
+static void transformTileLoad(LoadInst *LD, AliasAnalysis *AA) {
   Value *Row = nullptr, *Col = nullptr;
   Use &U = *(LD->use_begin());
   unsigned OpNo = U.getOperandNo();
   auto *II = cast<IntrinsicInst>(U.getUser());
   std::tie(Row, Col) = getShape(II, OpNo);
-  IRBuilder<> Builder(LD);
+  auto *InsertPt = findSafePointToSink(LD, II, AA);
+  IRBuilder<> Builder(InsertPt);
   // Use the maximun column as stride.
   Value *Stride = Builder.getInt64(64);
   Value *I8Ptr =
       Builder.CreateBitCast(LD->getOperand(0), Builder.getInt8PtrTy());
   std::array<Value *, 4> Args = {Row, Col, I8Ptr, Stride};
 
   Value *NewInst = Builder.CreateIntrinsic(
       Intrinsic::x86_tileloadd64_internal, None, Args);
std::array<Value *, 5> Args = {Row, Col, I8Ptr, Stride,
                                  ST->getValueOperand()};
   Builder.CreateIntrinsic(Intrinsic::x86_tilestored64_internal, None,
                           Args);
 }
 
 namespace {
 class X86LowerAMXType {
   Function &Func;
+  AliasAnalysis *AA;
 
 public:
-  X86LowerAMXType(Function &F) : Func(F) {}
+  X86LowerAMXType(Function &F, AliasAnalysis *AA) : Func(F), AA(AA) {}
   bool visit();
 };
 
 bool X86LowerAMXType::visit() {
   bool C;
   SmallVector<Instruction *, 8> DeadInsts;
 
   for (BasicBlock &BB : Func) {
     for (Instruction &Inst : BB) {
       if (!dyn_cast<BitCastInst>(&Inst))
         continue;
       auto *Src = Inst.getOperand(0);
       Type *Ty = Inst.getType();
       if (Ty->isPointerTy() &&
           cast<PointerType>(Ty)->getElementType()->isX86_AMXTy()) {
         for (auto UI = Inst.use_begin(), UE = Inst.use_end(); UI != UE;) {
           Value *I = (UI++)->getUser();
           auto *LD = dyn_cast<LoadInst>(I);
           // %0 = bitcast <256 x i32>* %tile to x86_amx*
           // %1 = load x86_amx, x86_amx* %0, align 64
           if (LD) {
-            transformTileLoad(LD);
+            transformTileLoad(LD, AA);
             DeadInsts.push_back(LD);
           }
           auto *ST = dyn_cast<StoreInst>(I);
           if (ST) {
             // %addr = bitcast <256 x i32>* %tile to x86_amx*
             // store x86_amx %9, x86_amx* %addr, align 64
             // -->
             // call void @llvm.x86.tilestored64.internal(%row, %col, %addr,

 public:
   static char ID;
 
   X86LowerAMXTypeLegacyPass() : FunctionPass(ID) {
     initializeX86LowerAMXTypeLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnFunction(Function &F) override {
-    X86LowerAMXType LAT(F);
+    auto *AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
+    X86LowerAMXType LAT(F, AA);
     bool C = LAT.visit();
     return C;
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<AAResultsWrapperPass>();
     AU.setPreservesCFG();
   }
 };
 
 } // anonymous namespace
 
 static const char PassName[] = "Lower AMX type for load/store";
 char X86LowerAMXTypeLegacyPass::ID = 0;
 INITIALIZE_PASS_BEGIN(X86LowerAMXTypeLegacyPass, DEBUG_TYPE, PassName, false,
                       false)
+INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
 INITIALIZE_PASS_END(X86LowerAMXTypeLegacyPass, DEBUG_TYPE, PassName, false,
                     false)
 
 FunctionPass *llvm::createX86LowerAMXTypePass() {
   return new X86LowerAMXTypeLegacyPass();
 }

llvm/test/CodeGen/X86/AMX/amx-type.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -lower-amx-type %s -S | FileCheck %s
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 %struct.__tile_str = type { i16, i16, <256 x i32> }
 
 @buf = dso_local global [1024 x i8] zeroinitializer, align 16
 @buf2 = dso_local global [1024 x i8] zeroinitializer, align 16
 
+define dso_local void @test_tile_sink(%struct.__tile_str* %a, %struct.__tile_str* %b, %struct.__tile_str* %c) #0 {
+; CHECK-LABEL: @test_tile_sink(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[B_COLPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR:%.*]], %struct.__tile_str* [[B:%.*]], i64 0, i32 1
+; CHECK-NEXT:    [[B_COL:%.*]] = load i16, i16* [[B_COLPTR]], align 2
+; CHECK-NEXT:    [[B_VPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[B]], i64 0, i32 2
+; CHECK-NEXT:    [[B_AMXPTR:%.*]] = bitcast <256 x i32>* [[B_VPTR]] to x86_amx*
+; CHECK-NEXT:    [[A_ROWPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[A:%.*]], i64 0, i32 0
+; CHECK-NEXT:    [[A_ROW:%.*]] = load i16, i16* [[A_ROWPTR]], align 64
+; CHECK-NEXT:    [[A_COLPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[A]], i64 0, i32 1
+; CHECK-NEXT:    [[A_COL:%.*]] = load i16, i16* [[A_COLPTR]], align 2
+; CHECK-NEXT:    [[A_VPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[A]], i64 0, i32 2
+; CHECK-NEXT:    [[A_AMXPTR:%.*]] = bitcast <256 x i32>* [[A_VPTR]] to x86_amx*
+; CHECK-NEXT:    [[C_VPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[C:%.*]], i64 0, i32 2
+; CHECK-NEXT:    [[C_AMXPTR:%.*]] = bitcast <256 x i32>* [[C_VPTR]] to x86_amx*
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast x86_amx* [[B_AMXPTR]] to i8*
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast x86_amx* [[A_AMXPTR]] to i8*
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast x86_amx* [[C_AMXPTR]] to i8*
+; CHECK-NEXT:    [[TMP3:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[A_ROW]], i16 [[B_COL]], i8* [[TMP2]], i64 64)
+; CHECK-NEXT:    [[TMP4:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[A_ROW]], i16 [[A_COL]], i8* [[TMP1]], i64 64)
+; CHECK-NEXT:    [[TMP5:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[A_COL]], i16 [[B_COL]], i8* [[TMP0]], i64 64)
+; CHECK-NEXT:    [[RES:%.*]] = tail call x86_amx @llvm.x86.tdpbssd.internal(i16 [[A_ROW]], i16 [[B_COL]], i16 [[A_COL]], x86_amx [[TMP3]], x86_amx [[TMP4]], x86_amx [[TMP5]]) [[ATTR1:#.*]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %b.colptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %b, i64 0, i32 1
+  %b.col = load i16, i16* %b.colptr, align 2
+  %b.vptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %b, i64 0, i32 2
+  %b.amxptr  = bitcast <256 x i32>* %b.vptr to x86_amx*
+  %b.tile = load x86_amx, x86_amx* %b.amxptr, align 64
+  %a.rowptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %a, i64 0, i32 0
+  %a.row = load i16, i16* %a.rowptr, align 64
+  %a.colptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %a, i64 0, i32 1
+  %a.col = load i16, i16* %a.colptr, align 2
+  %a.vptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %a, i64 0, i32 2
+  %a.amxptr = bitcast <256 x i32>* %a.vptr to x86_amx*
+  %a.tile = load x86_amx, x86_amx* %a.amxptr, align 64
+  %c.vptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %c, i64 0, i32 2
+  %c.amxptr = bitcast <256 x i32>* %c.vptr to x86_amx*
+  %c.tile = load x86_amx, x86_amx* %c.amxptr, align 64, !tbaa !2
+  %res = tail call x86_amx @llvm.x86.tdpbssd.internal(i16 %a.row, i16 %b.col, i16 %a.col, x86_amx %c.tile, x86_amx %a.tile, x86_amx %b.tile) #2
+  ret void
+}
+
+define dso_local void @test_tile_sink_noalias(%struct.__tile_str* %a, %struct.__tile_str* %b, %struct.__tile_str* %c) #0 {
+; CHECK-LABEL: @test_tile_sink_noalias(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[B_COLPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR:%.*]], %struct.__tile_str* [[B:%.*]], i64 0, i32 1
+; CHECK-NEXT:    [[B_COL:%.*]] = load i16, i16* [[B_COLPTR]], align 2
+; CHECK-NEXT:    [[B_VPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[B]], i64 0, i32 2
+; CHECK-NEXT:    [[B_AMXPTR:%.*]] = bitcast <256 x i32>* [[B_VPTR]] to x86_amx*
+; CHECK-NEXT:    store i16 8, i16* [[B_COLPTR]], align 2
+; CHECK-NEXT:    [[A_ROWPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[A:%.*]], i64 0, i32 0
+; CHECK-NEXT:    [[A_ROW:%.*]] = load i16, i16* [[A_ROWPTR]], align 64
+; CHECK-NEXT:    [[A_COLPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[A]], i64 0, i32 1
+; CHECK-NEXT:    [[A_COL:%.*]] = load i16, i16* [[A_COLPTR]], align 2
+; CHECK-NEXT:    [[A_VPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[A]], i64 0, i32 2
+; CHECK-NEXT:    [[A_AMXPTR:%.*]] = bitcast <256 x i32>* [[A_VPTR]] to x86_amx*
+; CHECK-NEXT:    [[C_VPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[C:%.*]], i64 0, i32 2
+; CHECK-NEXT:    [[C_AMXPTR:%.*]] = bitcast <256 x i32>* [[C_VPTR]] to x86_amx*
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast x86_amx* [[B_AMXPTR]] to i8*
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast x86_amx* [[A_AMXPTR]] to i8*
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast x86_amx* [[C_AMXPTR]] to i8*
+; CHECK-NEXT:    [[TMP3:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[A_ROW]], i16 [[B_COL]], i8* [[TMP2]], i64 64)
+; CHECK-NEXT:    [[TMP4:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[A_ROW]], i16 [[A_COL]], i8* [[TMP1]], i64 64)
+; CHECK-NEXT:    [[TMP5:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[A_COL]], i16 [[B_COL]], i8* [[TMP0]], i64 64)
+; CHECK-NEXT:    [[RES:%.*]] = tail call x86_amx @llvm.x86.tdpbssd.internal(i16 [[A_ROW]], i16 [[B_COL]], i16 [[A_COL]], x86_amx [[TMP3]], x86_amx [[TMP4]], x86_amx [[TMP5]]) [[ATTR1]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %b.colptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %b, i64 0, i32 1
+  %b.col = load i16, i16* %b.colptr, align 2
+  %b.vptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %b, i64 0, i32 2
+  %b.amxptr  = bitcast <256 x i32>* %b.vptr to x86_amx*
+  %b.tile = load x86_amx, x86_amx* %b.amxptr, align 64
+  ; test noalias
+  store i16 8, i16* %b.colptr
+  %a.rowptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %a, i64 0, i32 0
+  %a.row = load i16, i16* %a.rowptr, align 64
+  %a.colptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %a, i64 0, i32 1
+  %a.col = load i16, i16* %a.colptr, align 2
+  %a.vptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %a, i64 0, i32 2
+  %a.amxptr = bitcast <256 x i32>* %a.vptr to x86_amx*
+  %a.tile = load x86_amx, x86_amx* %a.amxptr, align 64
+  %c.vptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %c, i64 0, i32 2
+  %c.amxptr = bitcast <256 x i32>* %c.vptr to x86_amx*
+  %c.tile = load x86_amx, x86_amx* %c.amxptr, align 64, !tbaa !2
+  %res = tail call x86_amx @llvm.x86.tdpbssd.internal(i16 %a.row, i16 %b.col, i16 %a.col, x86_amx %c.tile, x86_amx %a.tile, x86_amx %b.tile) #2
+  ret void
+}
+
+define dso_local void @test_tile_sink_alias(%struct.__tile_str* %a, %struct.__tile_str* %b, %struct.__tile_str* %c) #0 {
+; CHECK-LABEL: @test_tile_sink_alias(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[B_COLPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR:%.*]], %struct.__tile_str* [[B:%.*]], i64 0, i32 1
+; CHECK-NEXT:    [[B_COL:%.*]] = load i16, i16* [[B_COLPTR]], align 2
+; CHECK-NEXT:    [[B_VPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[B]], i64 0, i32 2
+; CHECK-NEXT:    [[B_AMXPTR:%.*]] = bitcast <256 x i32>* [[B_VPTR]] to x86_amx*
+; CHECK-NEXT:    [[A_COLPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[A:%.*]], i64 0, i32 1
+; CHECK-NEXT:    [[A_COL:%.*]] = load i16, i16* [[A_COLPTR]], align 2
+; CHECK-NEXT:    [[B_SCALARPTR:%.*]] = bitcast <256 x i32>* [[B_VPTR]] to i32*
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast x86_amx* [[B_AMXPTR]] to i8*
+; CHECK-NEXT:    [[TMP1:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[A_COL]], i16 [[B_COL]], i8* [[TMP0]], i64 64)
+; CHECK-NEXT:    store i32 8, i32* [[B_SCALARPTR]], align 4
+; CHECK-NEXT:    [[A_ROWPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[A]], i64 0, i32 0
+; CHECK-NEXT:    [[A_ROW:%.*]] = load i16, i16* [[A_ROWPTR]], align 64
+; CHECK-NEXT:    [[A_VPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[A]], i64 0, i32 2
+; CHECK-NEXT:    [[A_AMXPTR:%.*]] = bitcast <256 x i32>* [[A_VPTR]] to x86_amx*
+; CHECK-NEXT:    [[C_VPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[C:%.*]], i64 0, i32 2
+; CHECK-NEXT:    [[C_AMXPTR:%.*]] = bitcast <256 x i32>* [[C_VPTR]] to x86_amx*
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast x86_amx* [[A_AMXPTR]] to i8*
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast x86_amx* [[C_AMXPTR]] to i8*
+; CHECK-NEXT:    [[TMP4:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[A_ROW]], i16 [[B_COL]], i8* [[TMP3]], i64 64)
+; CHECK-NEXT:    [[TMP5:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[A_ROW]], i16 [[A_COL]], i8* [[TMP2]], i64 64)
+; CHECK-NEXT:    [[RES:%.*]] = tail call x86_amx @llvm.x86.tdpbssd.internal(i16 [[A_ROW]], i16 [[B_COL]], i16 [[A_COL]], x86_amx [[TMP4]], x86_amx [[TMP5]], x86_amx [[TMP1]]) [[ATTR1]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %b.colptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %b, i64 0, i32 1
+  %b.col = load i16, i16* %b.colptr, align 2
+  %b.vptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %b, i64 0, i32 2
+  %b.amxptr  = bitcast <256 x i32>* %b.vptr to x86_amx*
+  %b.tile = load x86_amx, x86_amx* %b.amxptr, align 64
+  %a.colptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %a, i64 0, i32 1
+  %a.col = load i16, i16* %a.colptr, align 2
+  ; test alias
+  %b.scalarptr  = bitcast <256 x i32>* %b.vptr to i32*
+  store i32 8, i32* %b.scalarptr
+  %a.rowptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %a, i64 0, i32 0
+  %a.row = load i16, i16* %a.rowptr, align 64
+  %a.vptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %a, i64 0, i32 2
+  %a.amxptr = bitcast <256 x i32>* %a.vptr to x86_amx*
+  %a.tile = load x86_amx, x86_amx* %a.amxptr, align 64
+  %c.vptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %c, i64 0, i32 2
+  %c.amxptr = bitcast <256 x i32>* %c.vptr to x86_amx*
+  %c.tile = load x86_amx, x86_amx* %c.amxptr, align 64, !tbaa !2
+  %res = tail call x86_amx @llvm.x86.tdpbssd.internal(i16 %a.row, i16 %b.col, i16 %a.col, x86_amx %c.tile, x86_amx %a.tile, x86_amx %b.tile) #2
+  ret void
+}
+
+define dso_local void @test_tile_sink_across_bb(%struct.__tile_str* %a, %struct.__tile_str* %b, %struct.__tile_str* %c) #0 {
+; CHECK-LABEL: @test_tile_sink_across_bb(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[B_COLPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR:%.*]], %struct.__tile_str* [[B:%.*]], i64 0, i32 1
+; CHECK-NEXT:    [[B_COL:%.*]] = load i16, i16* [[B_COLPTR]], align 2
+; CHECK-NEXT:    [[B_VPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[B]], i64 0, i32 2
+; CHECK-NEXT:    [[B_AMXPTR:%.*]] = bitcast <256 x i32>* [[B_VPTR]] to x86_amx*
+; CHECK-NEXT:    [[A_COLPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[A:%.*]], i64 0, i32 1
+; CHECK-NEXT:    [[A_COL:%.*]] = load i16, i16* [[A_COLPTR]], align 2
+; CHECK-NEXT:    [[A_ROWPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[A]], i64 0, i32 0
+; CHECK-NEXT:    [[A_ROW:%.*]] = load i16, i16* [[A_ROWPTR]], align 64
+; CHECK-NEXT:    [[A_VPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[A]], i64 0, i32 2
+; CHECK-NEXT:    [[A_AMXPTR:%.*]] = bitcast <256 x i32>* [[A_VPTR]] to x86_amx*
+; CHECK-NEXT:    [[C_VPTR:%.*]] = getelementptr inbounds [[STRUCT___TILE_STR]], %struct.__tile_str* [[C:%.*]], i64 0, i32 2
+; CHECK-NEXT:    [[C_AMXPTR:%.*]] = bitcast <256 x i32>* [[C_VPTR]] to x86_amx*
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast x86_amx* [[B_AMXPTR]] to i8*
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast x86_amx* [[A_AMXPTR]] to i8*
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast x86_amx* [[C_AMXPTR]] to i8*
+; CHECK-NEXT:    [[TMP3:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[A_ROW]], i16 [[B_COL]], i8* [[TMP2]], i64 64)
+; CHECK-NEXT:    [[TMP4:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[A_ROW]], i16 [[A_COL]], i8* [[TMP1]], i64 64)
+; CHECK-NEXT:    [[TMP5:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[A_COL]], i16 [[B_COL]], i8* [[TMP0]], i64 64)
+; CHECK-NEXT:    br label [[DOTPROD:%.*]]
+; CHECK:       dotprod:
+; CHECK-NEXT:    [[RES:%.*]] = tail call x86_amx @llvm.x86.tdpbssd.internal(i16 [[A_ROW]], i16 [[B_COL]], i16 [[A_COL]], x86_amx [[TMP3]], x86_amx [[TMP4]], x86_amx [[TMP5]]) [[ATTR1]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %b.colptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %b, i64 0, i32 1
+  %b.col = load i16, i16* %b.colptr, align 2
+  %b.vptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %b, i64 0, i32 2
+  %b.amxptr  = bitcast <256 x i32>* %b.vptr to x86_amx*
+  %b.tile = load x86_amx, x86_amx* %b.amxptr, align 64
+  %a.colptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %a, i64 0, i32 1
+  %a.col = load i16, i16* %a.colptr, align 2
+  %a.rowptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %a, i64 0, i32 0
+  %a.row = load i16, i16* %a.rowptr, align 64
+  %a.vptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %a, i64 0, i32 2
+  %a.amxptr = bitcast <256 x i32>* %a.vptr to x86_amx*
+  %a.tile = load x86_amx, x86_amx* %a.amxptr, align 64
+  %c.vptr = getelementptr inbounds %struct.__tile_str, %struct.__tile_str* %c, i64 0, i32 2
+  %c.amxptr = bitcast <256 x i32>* %c.vptr to x86_amx*
+  %c.tile = load x86_amx, x86_amx* %c.amxptr, align 64, !tbaa !2
+  br label %dotprod
+dotprod:                                           ; preds = %if.else, %if.then
+  %res = tail call x86_amx @llvm.x86.tdpbssd.internal(i16 %a.row, i16 %b.col, i16 %a.col, x86_amx %c.tile, x86_amx %a.tile, x86_amx %b.tile) #2
+  ret void
+}
+
 define dso_local void @test_amx_store(<256 x i32>* %in, i16 %m, i16 %n, i8 *%buf, i64 %s) #2 {
 ; CHECK-LABEL: @test_amx_store(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[T0:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[M:%.*]], i16 [[N:%.*]], i8* [[BUF:%.*]], i64 [[S:%.*]]) [[ATTR3:#.*]]
 ; CHECK-NEXT:    [[ADDR:%.*]] = bitcast <256 x i32>* [[IN:%.*]] to x86_amx*
 ; CHECK-NEXT:    [[TMP0:%.*]] = bitcast x86_amx* [[ADDR]] to i8*
 ; CHECK-NEXT:    call void @llvm.x86.tilestored64.internal(i16 [[M]], i16 [[N]], i8* [[TMP0]], i64 64, x86_amx [[T0]])
 ; CHECK-NEXT:    ret void