[X86] Canonicalize AMX bitcast instruction.

[LuoYuanke]

Summary

Canonicalize AMX bitcast instruction to delete unnecessary bitcast instruction.

[LuoYuanke]

Move comments.

llvm/lib/Target/X86/X86LowerAMXType.cpp

 
 public:
   X86LowerAMXType(Function &F) : Func(F) {}
   bool visit();
 };
 
 bool X86LowerAMXType::visit() {
   SmallVector<Instruction *, 8> DeadInsts;
+  SmallVector<Instruction *, 2> DeadBitcasts;
 
   for (BasicBlock *BB : post_order(&Func)) {
     for (BasicBlock::reverse_iterator II = BB->rbegin(), IE = BB->rend();
          II != IE;) {
       Instruction &Inst = *II++;
       auto *Bitcast = dyn_cast<BitCastInst>(&Inst);
       if (!Bitcast)
         continue;
 
       Value *Src = Bitcast->getOperand(0);
       Type *Ty = Bitcast->getType();
+      auto CanonicalizeBitcast = [&]() {
+        if (Bitcast->user_empty()) {
+          DeadBitcasts.push_back(Bitcast);
+          return true;
+        }
+        bool Change = false;
+        Value *DstV = Src, *PreDst = Bitcast, *SrcV;
+        while (match(DstV, m_BitCast(m_Value(SrcV))) &&
+               SrcV->getType()->getTypeID() == PreDst->getType()->getTypeID()) {
+          PreDst->replaceAllUsesWith(SrcV);
+          DeadBitcasts.push_back(cast<Instruction>(PreDst));
+          PreDst = DstV;
+          DstV = SrcV;
+          Change = true;
+        }
+        return Change;
+      };
 
       if (Ty->isPointerTy() &&
           cast<PointerType>(Ty)->getElementType()->isX86_AMXTy()) {
         for (auto UI = Bitcast->use_begin(), UE = Bitcast->use_end();
              UI != UE;) {
           Value *I = (UI++)->getUser();
           auto *LD = dyn_cast<LoadInst>(I);
           // %0 = bitcast <256 x i32>* %tile to x86_amx*
for (BasicBlock::reverse_iterator II = BB->rbegin(), IE = BB->rend();
           }
         }
         // If the dst type is <256 x i32>*, it is valid intruction.
         // %0 = bitcast x86_amx* %tile to <256 x i32>*
         // %1 = load <256 x i32>, <256 x i32>* %0, align 64
         // store <256 x i32> %2, <256 x i32>* %0, align 64
       }
       if (Bitcast->getType()->isX86_AMXTy()) {
-        if (Bitcast->user_empty()) {
-          DeadInsts.push_back(Bitcast);
+        if (CanonicalizeBitcast())
           continue;
-        }
         LoadInst *LD = dyn_cast<LoadInst>(Src);
         if (!LD) {
           if (transformBitcast(Bitcast))
             DeadInsts.push_back(Bitcast);
           continue;
         }
         // If load has mutli-user, duplicate a amx load.
         // %src = load <256 x i32>, <256 x i32>* %addr, align 64
for (BasicBlock::reverse_iterator II = BB->rbegin(), IE = BB->rend();
         // -->
         // %2 = call x86_amx @llvm.x86.tileloadd64.internal(i16 %row, i16 %col,
         //                                            i8* %addr, i64 %stride64)
         combineLoadBitcast(LD, Bitcast);
         DeadInsts.push_back(Bitcast);
         if (LD->hasOneUse())
           DeadInsts.push_back(LD);
       } else if (Src->getType()->isX86_AMXTy()) {
-        if (Bitcast->user_empty()) {
-          DeadInsts.push_back(Bitcast);
+        if (CanonicalizeBitcast())
           continue;
-        }
         StoreInst *ST = nullptr;
         for (auto UI = Bitcast->use_begin(), UE = Bitcast->use_end();
              UI != UE;) {
           Value *I = (UI++)->getUser();
           ST = dyn_cast<StoreInst>(I);
           if (ST)
             break;
         }
for (BasicBlock::reverse_iterator II = BB->rbegin(), IE = BB->rend();
         combineBitcastStore(Bitcast, ST);
         // Delete user first.
         DeadInsts.push_back(ST);
         DeadInsts.push_back(Bitcast);
       }
     }
   }
 
-  bool C = !DeadInsts.empty();
+  bool C = !DeadInsts.empty() || !DeadBitcasts.empty();
 
   SmallSet<Instruction *, 8> DeletedInst;
   auto DeleteInst = [&](Instruction *Inst) {
     SmallVector<Instruction *, 4> DeadIs;
     DeadIs.push_back(Inst);
     while (!DeadIs.empty()) {
       auto *Inst = DeadIs.back();
       DeadIs.pop_back();
       if (DeletedInst.count(Inst))
         continue;
       for (auto I = Inst->op_begin(), E = Inst->op_end(); I != E;) {
         Instruction *Op = dyn_cast<Instruction>(*I);
         if (Op && Op->hasOneUse())
           DeadIs.push_back(Op);
         ++I;
       }
       Inst->eraseFromParent();
       DeletedInst.insert(Inst);
     }
   };
   for (auto *Inst : DeadInsts)
     DeleteInst(Inst);
+  for (auto *Inst : DeadBitcasts)
+    DeleteInst(Inst);
 
   return C;
 }
 } // anonymous namespace
 
 namespace {
 
 class X86LowerAMXTypeLegacyPass : public FunctionPass {

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 <256 x i32> @test_amx_bitcast(<256 x i32> %in) #2 {
+; CHECK-LABEL: @test_amx_bitcast(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    ret <256 x i32> [[IN:%.*]]
+;
+entry:
+  %amx = bitcast <256 x i32> %in to x86_amx
+  %vec = bitcast x86_amx %amx to <256 x i32>
+  ret <256 x i32> %vec
+}
+
 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