[NFC] Refactor the code to avoid check target type.

[LuoYuanke]

Fix the comments.

[lebedev.ri]

IMHO being explicit about the ugliness of the hack (LHS of the diff) is better than inventing a confusing misnomer (how is bitcast between pointers with same address space not lossless?) and hiding behind it.

[RKSimon]

Is this really NFC? If it is, then as @lebedev.ri commented - it feels like you're hiding an AMX issue. If it isn't NFC we need more test coverage.

[LuoYuanke]

In D98595#2626127, @RKSimon wrote:

Is this really NFC? If it is, then as @lebedev.ri commented - it feels like you're hiding an AMX issue. If it isn't NFC we need more test coverage.

I think it should be NFC. The test case is already at test/Transforms/InstCombine/X86/x86-amx-load-store.ll.

[LuoYuanke]

I'll look into this further to see if there is a better solution.

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

if (LI.getPointerOperand()->isSwiftError())
     return nullptr;
 
   const DataLayout &DL = IC.getDataLayout();
 
   // Fold away bit casts of the loaded value by loading the desired type.
   // Note that we should not do this for pointer<->integer casts,
   // because that would result in type punning.
   if (LI.hasOneUse()) {
-    // Don't transform when the type is x86_amx, it makes the pass that lower
-    // x86_amx type happy.
     if (auto *BC = dyn_cast<BitCastInst>(LI.user_back())) {
-      assert(!LI.getType()->isX86_AMXTy() &&
-             "load from x86_amx* should not happen!");
-      if (BC->getType()->isX86_AMXTy())
+      // Prevent to bitcast the pointer if the cast is not lossless.
+      auto *DestPtrTy =
+          BC->getType()->getPointerTo(LI.getPointerAddressSpace());
+      auto *SrcPtrTy = LI.getPointerOperandType();
+      if (!SrcPtrTy->canLosslesslyBitCastTo(DestPtrTy))
         return nullptr;
     }
 
     if (auto* CI = dyn_cast<CastInst>(LI.user_back()))
       if (CI->isNoopCast(DL) && LI.getType()->isPtrOrPtrVectorTy() ==
                                     CI->getDestTy()->isPtrOrPtrVectorTy())
         if (!LI.isAtomic() || isSupportedAtomicType(CI->getDestTy())) {
           LoadInst *NewLoad = IC.combineLoadToNewType(LI, CI->getDestTy());
static bool combineStoreToValueType(InstCombinerImpl &IC, StoreInst &SI) {
   // swifterror values can't be bitcasted.
   if (SI.getPointerOperand()->isSwiftError())
     return false;
 
   Value *V = SI.getValueOperand();
 
   // Fold away bit casts of the stored value by storing the original type.
   if (auto *BC = dyn_cast<BitCastInst>(V)) {
-    assert(!BC->getType()->isX86_AMXTy() &&
-           "store to x86_amx* should not happen!");
     V = BC->getOperand(0);
-    // Don't transform when the type is x86_amx, it makes the pass that lower
-    // x86_amx type happy.
-    if (V->getType()->isX86_AMXTy())
+    auto *DestPtrTy = V->getType()->getPointerTo(SI.getPointerAddressSpace());
+    auto *SrcPtrTy = SI.getPointerOperandType();
+    if (!SrcPtrTy->canLosslesslyBitCastTo(DestPtrTy))
       return false;
     if (!SI.isAtomic() || isSupportedAtomicType(V->getType())) {
       combineStoreToNewValue(IC, SI, V);
       return true;
     }
   }
 
   if (Value *U = likeBitCastFromVector(IC, V))