[X86][AMX] Prohibit pointer cast on load.

[LuoYuanke]

Summary

The load/store instruction will be transformed to amx intrinsics in the
pass of AMX type lowering. Prohibiting the pointer cast make that pass
happy.

[LuoYuanke]

Fix typo in the comments.

[pengfei]

to amx intrinsics the the

to amx intrinsics that?

llvm/test/Transforms/InstCombine/load.ll

427 ↗

(On Diff #315650)

Missing check results?

[LuoYuanke]

llvm/test/Transforms/InstCombine/load.ll

427 ↗

(On Diff #315650)

Oops. Thank Pengfei. :)

[LuoYuanke]

Address Pengfei's comments.

[pengfei]

llvm/test/Transforms/InstCombine/load.ll

452 ↗

(On Diff #315660)

I think we had disabled the chance that x86_amx* been generated by compiler. It may mean somthing wrong if we see such load. So I think it's better just raise an assert fail for it here.
Besides, it looks safe and better for performance to combine

%vec = load x86_amx, x86_amx* %src, align 64
%bc = bitcast x86_amx %vec to <256 x i32>

to

%bc = load <256 x i32>, <256 x i32>* %src, align 64

[lebedev.ri]

This comment was removed by lebedev.ri.

[lebedev.ri]

What specifically is having problems after that transformation?

[LuoYuanke]

In D94372#2489048, @lebedev.ri wrote:

What specifically is having problems after that transformation?

The the amx load require shape and stride. We need transform "%vec = load x86_amx, x86_amx* %src" to "%vec = @llvm.x86.tileloadd64.internal(i16 %row, i16 %col, i8* %src, i64 %stride)". The %row, %col need to be deduced from the context. Prohibit the transform make pass "lower-amx-type" happy.

[LuoYuanke]

llvm/test/Transforms/InstCombine/load.ll

452 ↗

(On Diff #315660)

So we may delete this case, as we think "x86_amx*" is invalid. I can take x86_amx* as invalid in IR verification in another patch.

[LuoYuanke]

Address Pengfei's comments.

[pengfei]

LGTM, but I'd like to see if @lebedev.ri has any objections.

[LuoYuanke]

In D94372#2491970, @pengfei wrote:

LGTM, but I'd like to see if @lebedev.ri has any objections.

Ok, I'll wait 1 day to see if @lebedev.ri has any objection for the patch.

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

static Instruction *combineLoadToOperationType(InstCombinerImpl &IC,
   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())
+  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())
+        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());
           CI->replaceAllUsesWith(NewLoad);
           IC.eraseInstFromFunction(*CI);
           return &LI;
         }
+  }
 
   // FIXME: We should also canonicalize loads of vectors when their elements are
   // cast to other types.
   return nullptr;
 }
 
 static Instruction *unpackLoadToAggregate(InstCombinerImpl &IC, LoadInst &LI) {
   // FIXME: We could probably with some care handle both volatile and atomic
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 make the pass that lower
+    // Don't transform when the type is x86_amx, it makes the pass that lower
     // x86_amx type happy.
-    if (BC->getType()->isX86_AMXTy() || V->getType()->isX86_AMXTy())
+    if (V->getType()->isX86_AMXTy())
       return false;
     if (!SI.isAtomic() || isSupportedAtomicType(V->getType())) {
       combineStoreToNewValue(IC, SI, V);
       return true;
     }
   }
 
   if (Value *U = likeBitCastFromVector(IC, V))

llvm/test/Transforms/InstCombine/X86/x86-amx-load-store.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -instcombine -S < %s | FileCheck %s
+; RUN: opt -passes=instcombine -S < %s | FileCheck %s
+
+; Prohibit poiter cast for amx.
+define dso_local void @test_amx_load_store(<256 x i32>* %src, i8* %dst) {
+; CHECK-LABEL: @test_amx_load_store(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[VEC:%.*]] = load <256 x i32>, <256 x i32>* [[SRC:%.*]], align 64
+; CHECK-NEXT:    [[BC:%.*]] = bitcast <256 x i32> [[VEC]] to x86_amx
+; CHECK-NEXT:    tail call void @llvm.x86.tilestored64.internal(i16 16, i16 16, i8* [[DST:%.*]], i64 64, x86_amx [[BC]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %vec = load <256 x i32>, <256 x i32>* %src, align 64
+  %bc = bitcast <256 x i32> %vec to x86_amx
+  tail call void @llvm.x86.tilestored64.internal(i16 16, i16 16, i8* %dst, i64 64, x86_amx %bc)
+  ret void
+}
+
+; Prohibit poiter cast for amx.
+define dso_local void @test_amx_load_store2(<256 x i32>* %dst, i8* %src) {
+; CHECK-LABEL: @test_amx_load_store2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AMX:%.*]] = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 16, i16 16, i8* [[SRC:%.*]], i64 64)
+; CHECK-NEXT:    [[BC:%.*]] = bitcast x86_amx [[AMX]] to <256 x i32>
+; CHECK-NEXT:    store <256 x i32> [[BC]], <256 x i32>* [[DST:%.*]], align 1024
+; CHECK-NEXT:    ret void
+;
+entry:
+  %amx = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 16, i16 16, i8* %src, i64 64)
+  %bc = bitcast x86_amx %amx to <256 x i32>
+  store <256 x i32> %bc, <256 x i32>* %dst
+  ret void
+}
+
+declare x86_amx @llvm.x86.tileloadd64.internal(i16, i16, i8*, i64)
+declare void @llvm.x86.tilestored64.internal(i16, i16, i8*, i64, x86_amx)