[X86][AMX] Let Store not be removed if combineCastStore failed

[pengfei]

llvm/lib/Target/X86/X86LowerAMXType.cpp

999

Given the condition here, there's no chance return false so far.

[LuoYuanke]

Is there any test case for the patch?

[yubing]

llvm/lib/Target/X86/X86LowerAMXType.cpp

999

but better we align combineCastStore with combineLoadCast. combineCastStore might return false in the future

[LuoYuanke]

LGTM.

llvm/lib/Target/X86/X86LowerAMXType.cpp

 namespace {
 
 class X86LowerAMXCast {
   Function &Func;
   std::unique_ptr<DominatorTree> DT;
 
 public:
   X86LowerAMXCast(Function &F) : Func(F), DT(nullptr) {}
-  void combineCastStore(IntrinsicInst *Cast, StoreInst *ST);
+  bool combineCastStore(IntrinsicInst *Cast, StoreInst *ST);
   bool combineLoadCast(IntrinsicInst *Cast, LoadInst *LD);
   bool combineLdSt(SmallVectorImpl<Instruction *> &Casts);
   bool combineAMXcast(TargetLibraryInfo *TLI);
   bool transformAMXCast(IntrinsicInst *AMXCast);
   bool transformAllAMXCast();
   bool optimizeAMXCastFromPhi(IntrinsicInst *CI, PHINode *PN,
                               SmallSetVector<Instruction *, 16> &DeadInst);
 };
bool X86LowerAMXCast::optimizeAMXCastFromPhi(
   return true;
 }
 
 // %43 = call <256 x i32> @llvm.x86.cast.tile.to.vector.v256i32(x86_amx %42)
 // store <256 x i32> %43, <256 x i32>* %p, align 64
 // -->
 // call void @llvm.x86.tilestored64.internal(i16 %row, i16 %col, i8* %p,
 //                                           i64 64, x86_amx %42)
-void X86LowerAMXCast::combineCastStore(IntrinsicInst *Cast, StoreInst *ST) {
+bool X86LowerAMXCast::combineCastStore(IntrinsicInst *Cast, StoreInst *ST) {
   Value *Tile = Cast->getOperand(0);
   // TODO: If it is cast intrinsic or phi node, we can propagate the
   // shape information through def-use chain.
   if (!isAMXIntrinsic(Tile))
-    return;
+    return false;
   auto *II = cast<IntrinsicInst>(Tile);
   // Tile is output from AMX intrinsic. The first operand of the
   // intrinsic is row, the second operand of the intrinsic is column.
   Value *Row = II->getOperand(0);
   Value *Col = II->getOperand(1);
   IRBuilder<> Builder(ST);
   // Use the maximum column as stride. It must be the same with load
   // stride.
   Value *Stride = Builder.getInt64(64);
   Value *I8Ptr =
       Builder.CreateBitCast(ST->getOperand(1), Builder.getInt8PtrTy());
   std::array<Value *, 5> Args = {Row, Col, I8Ptr, Stride, Tile};
   Builder.CreateIntrinsic(Intrinsic::x86_tilestored64_internal, std::nullopt,
                           Args);
+  return true;
 }
 
 // %65 = load <256 x i32>, <256 x i32>* %p, align 64
 // %66 = call x86_amx @llvm.x86.cast.vector.to.tile(<256 x i32> %65)
 // -->
 // %66 = call x86_amx @llvm.x86.tileloadd64.internal(i16 %row, i16 %col,
 //                                                   i8* %p, i64 64)
 bool X86LowerAMXCast::combineLoadCast(IntrinsicInst *Cast, LoadInst *LD) {
bool X86LowerAMXCast::combineLoadCast(IntrinsicInst *Cast, LoadInst *LD) {
 
   return EraseLoad;
 }
 
 bool X86LowerAMXCast::combineLdSt(SmallVectorImpl<Instruction *> &Casts) {
   bool Change = false;
   for (auto *Cast : Casts) {
     auto *II = cast<IntrinsicInst>(Cast);
     // %43 = call <256 x i32> @llvm.x86.cast.tile.to.vector(x86_amx %42)
     // store <256 x i32> %43, <256 x i32>* %p, align 64
     // -->
     // call void @llvm.x86.tilestored64.internal(i16 %row, i16 %col, i8* %p,
     //                                           i64 64, x86_amx %42)
     if (II->getIntrinsicID() == Intrinsic::x86_cast_tile_to_vector) {
       SmallVector<Instruction *, 2> DeadStores;
       for (User *U : Cast->users()) {
         StoreInst *Store = dyn_cast<StoreInst>(U);
         if (!Store)
           continue;
-        combineCastStore(cast<IntrinsicInst>(Cast), Store);
+        if (combineCastStore(cast<IntrinsicInst>(Cast), Store)) {
           DeadStores.push_back(Store);
           Change = true;
         }
+      }
       for (auto *Store : DeadStores)
         Store->eraseFromParent();
     } else { // x86_cast_vector_to_tile
       SmallVector<Instruction *, 2> DeadLoads;
       auto *Load = dyn_cast<LoadInst>(Cast->getOperand(0));
       if (!Load || !Load->hasOneUse())
         continue;
       // %65 = load <256 x i32>, <256 x i32>* %p, align 64