[X86][AMX] Prevent transforming load pointer from <256 x i32>* to x86_amx*.

[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.

[lebedev.ri]

I'm not really convinced this is moving in the right direction.
other side of the problem should be fixed.

[LuoYuanke]

In D98247#2613648, @lebedev.ri wrote:

I'm not really convinced this is moving in the right direction.
other side of the problem should be fixed.

Thanks for review. Prohibiting the transform of pointer cast from <256 x i32>* to x86_amx* looks straight forward. I also try another way in patch https://reviews.llvm.org/D93788, but it seems not as good as prohibit pointer cast. Do you have any suggestions? I would like to try if you have any good idea on it.

[lebedev.ri]

In D98247#2613689, @LuoYuanke wrote:

In D98247#2613648, @lebedev.ri wrote:

I'm not really convinced this is moving in the right direction.
other side of the problem should be fixed.

Thanks for review. Prohibiting the transform of pointer cast from <256 x i32>* to x86_amx* looks straight forward.

Hm. Two things: pointer types are going away, and we can't make any optimizations based on the pointer type regardless.

I also try another way in patch https://reviews.llvm.org/D93788, but it seems not as good as prohibit pointer cast. Do you have any suggestions? I would like to try if you have any good idea on it.

[LuoYuanke]

Hm. Two things: pointer types are going away, and we can't make any optimizations based on the pointer type regardless.

Would you post any linkage for the information that pointer types are going away? I'd like to investigate how it will impact AMX.
This patch is not for optimization. Since there is no HW instruction corresponding to "load x86_amx, x86_amx*", I'd like keep the load as "load <256 x i32>, <256 x i32>*" form.

[lebedev.ri]

In D98247#2613716, @LuoYuanke wrote:

Hm. Two things: pointer types are going away, and we can't make any optimizations based on the pointer type regardless.

Would you post any linkage for the information that pointer types are going away? I'd like to investigate how it will impact AMX.

See e.g. https://lists.llvm.org/pipermail/llvm-dev/2019-December/137684.html

This patch is not for optimization. Since there is no HW instruction corresponding to "load x86_amx, x86_amx*", I'd like keep the load as "load <256 x i32>, <256 x i32>*" form.

Let me reword: we can not use pointer element type to decide whether or not to perform a change.

[LuoYuanke]

See e.g. https://lists.llvm.org/pipermail/llvm-dev/2019-December/137684.html

There are already several pointer cast transform in LLVM. Does opaque pointer require change of all pointer cast in LLVM? Nevertheless, it seems opaque pointer is not related with this patch.

This patch is not for optimization. Since there is no HW instruction corresponding to "load x86_amx, x86_amx*", I'd like keep the load as "load <256 x i32>, <256 x i32>*" form.

Let me reword: we can not use pointer element type to decide whether or not to perform a change.

Do you have any proposal for my case? The transform break the lowering of "load x86_amx, x86_amx*". I'm appreciated if you have good ideas.

[pengfei]

Let me reword: we can not use pointer element type to decide whether or not to perform a change.

It's not precise to say we are using pointer element type. The DestTy is the dest type of a bitcast after phi. See the comments in line 2364.
What we are doing is to block the phi+bitcast optimization when it will change loading to amx* type. I don't think the opaque ptr solution can solve the problem. The amx type load is quite different from other loads. You can think its expensive. We need to reduce it, or at least stop optimizing to it.

[pengfei]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

2398

Better to add a comment for the reason.

[LuoYuanke]

Address Pengfei's comments.

[lebedev.ri]

I do not understand how load x86_amx can cause trouble while load <256 x i32> works.
More generally, if this proceeds, what will happen if i directly feed the current (bad) IR to llc?
It will still have issues, will it not?
If they are interchangeable like that, why the lowering can not be extended to be agnostic of the actual type?

[LuoYuanke]

In D98247#2615789, @lebedev.ri wrote:

I do not understand how load x86_amx can cause trouble while load <256 x i32> works.
More generally, if this proceeds, what will happen if i directly feed the current (bad) IR to llc?
It will still have issues, will it not?
If they are interchangeable like that, why the lowering can not be extended to be agnostic of the actual type?

It is a long story, but let me try to explain.

1. We have dedicate x86_amx type for AMX hardware.
2. LLVM provide several intrinsics for AMX, and Clang have corresponding builtin that map to those intrinsics.
3. We don't have a new type in C language for AMX, we use vector type instead.
4. When Clang front-end create LLVM AMX intrinsics from AMX buildin, it cast the vector type <256 x i32> to x86_amx and feed x86_amx to AMX intrinsics. However there are still load/store <256 x i32> instructions generated in O0 build. So we create a pass to transform load/store <256 x i32> if the input or output of load/store is propagated from AMX intrinsics. We do the transform start from bitcast of <256 x i32> or x86_amx. The bitcast of <256 x i32>* to x86_amx is unexpected. And it cause much work to support it. So we want to prevent generating x86_amx* in our IR.

[LuoYuanke]

@lebedev.ri, may I commit the patch? I think the risk of this patch is pretty low. We can revise our solution if we have any better ideas in the future.

[pengfei]

LGTM anyway :)

[LuoYuanke]

Rebase.

[lebedev.ri]

I still believe this is a wrong approach and the code
that is having trouble with the 'optimized' should be fixed,
but i'll defer to @spatel.

[spatel]

In D98247#2618631, @lebedev.ri wrote:

I still believe this is a wrong approach and the code
that is having trouble with the 'optimized' should be fixed,
but i'll defer to @spatel.

I don't know anything about the AMX type / functionality, so I'm probably not the best judge.
I agree that we want to avoid type-based hacks (but what does it mean that we even have target-specific types in IR?)...
OTOH there is already precedent for AMX exceptions in instcombine (and MMX before that). I think we managed to make some of the MMX hacks less obviously bad by excluding all target-specific types from a given transform. Is that a possibility here? That is, could we limit the transform using isIntOrIntVectorTy() or similar?

[LuoYuanke]

Address Sanjay's comments.

[LuoYuanke]

I don't know anything about the AMX type / functionality, so I'm probably not the best judge.
I agree that we want to avoid type-based hacks (but what does it mean that we even have target-specific types in IR?)...
OTOH there is already precedent for AMX exceptions in instcombine (and MMX before that). I think we managed to make some of the MMX hacks less obviously bad by excluding all target-specific types from a given transform. Is that a possibility here? That is, could we limit the transform using isIntOrIntVectorTy() or similar?

Excluding all target-specific types from a given transform looks good to me. If we agree on this, I may refactor some of the code to exclude both MMX and AMX with a common type interface.

[pengfei]

In D98247#2620989, @LuoYuanke wrote:

I don't know anything about the AMX type / functionality, so I'm probably not the best judge.
I agree that we want to avoid type-based hacks (but what does it mean that we even have target-specific types in IR?)...
OTOH there is already precedent for AMX exceptions in instcombine (and MMX before that). I think we managed to make some of the MMX hacks less obviously bad by excluding all target-specific types from a given transform. Is that a possibility here? That is, could we limit the transform using isIntOrIntVectorTy() or similar?

Excluding all target-specific types from a given transform looks good to me. If we agree on this, I may refactor some of the code to exclude both MMX and AMX with a common type interface.

I don't think all target-specific types always have the same behavior. So a functionality repartition looks better to me. For example, under this AMX circumstance, the limitation is we are not preferring any pointer type been bitcasted to AMX* since it much more "expensive". So we may need an interface something like

bool isTriviallyBitcastToPointer() const { return isX86_AMXTy(); }

[LuoYuanke]

Address Pengfei's comments.

[LuoYuanke]

Fix bug. We need check bitcast is lossless.

[lebedev.ri]

In D98247#2615797, @LuoYuanke wrote:

And it cause much work to support it. So we want to prevent generating x86_amx* in our IR.

How much work?
I honestly still don't understand how handling the LHS version of IR in this diff is harder than RHS.

[pengfei]

In D98247#2621399, @lebedev.ri wrote:

In D98247#2615797, @LuoYuanke wrote:

And it cause much work to support it. So we want to prevent generating x86_amx* in our IR.

How much work?
I honestly still don't understand how handling the LHS version of IR in this diff is harder than RHS.

2 aspects. 1) It will increase compiling time. We actually need to handle it in codegen by truning it into load <256 x i32> + bitcast to x86_amx and lower it to a tileload in pass "lower-amx-type". We also need to infer the tile shape by its use. 2) It will increase registers pressure of tile. We only have 8 tile registers. And what makes it worse is that the different shapes cannot be reused after tileconfig.
But if we keep it as non AMX ptr, it's probably lowed to optimizated memcpy. That's why I think it is "cheap".

[LuoYuanke]

@lebedev.ri, after applying Sanjay and Pengfei's suggestion for the patch, current code check if the the bitcast is lossless before transform. I think it conform the previous code infrastructure and looks neat. It is much better that previous patch. May I commit the patch now?

[lebedev.ri]

This is a backend bug.

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

Instruction *InstCombinerImpl::optimizeBitCastFromPhi(CastInst &CI,
   OldPhiNodes.insert(PN);
   while (!PhiWorklist.empty()) {
     auto *OldPN = PhiWorklist.pop_back_val();
     for (Value *IncValue : OldPN->incoming_values()) {
       if (isa<Constant>(IncValue))
         continue;
 
       if (auto *LI = dyn_cast<LoadInst>(IncValue)) {
+        // Don't transform "load <256 x i32>, <256 x i32>*" to
+        // "load x86_amx, x86_amx*", because we don't have a corresponding
+        // instruction to load x86_amx. Doing the transform causes trouble
+        // to lower "load x86_amx" instruction in backend.
+        if (DestTy->isX86_AMXTy())
+          return nullptr;
         // If there is a sequence of one or more load instructions, each loaded
         // value is used as address of later load instruction, bitcast is
         // necessary to change the value type, don't optimize it. For
         // simplicity we give up if the load address comes from another load.
         Value *Addr = LI->getOperand(0);
         if (Addr == &CI || isa<LoadInst>(Addr))
           return nullptr;
         if (LI->hasOneUse() && LI->isSimple())

llvm/test/Transforms/InstCombine/X86/x86-amx.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
 
 define linkonce_odr dso_local void @foo(<256 x i32>* %arrayidx16, <256 x i32>* %arrayidx29, <256 x i32>* %arrayidx35) local_unnamed_addr {
 ; CHECK-LABEL: @foo(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[FOR_COND9:%.*]]
 ; CHECK:       for.cond9:
 ; CHECK-NEXT:    br i1 undef, label [[FOR_BODY14:%.*]], label [[EXIT:%.*]]
 ; CHECK:       for.body14:
-; CHECK-NEXT:    [[TMP0:%.*]] = bitcast <256 x i32>* [[ARRAYIDX16:%.*]] to x86_amx*
-; CHECK-NEXT:    [[T51:%.*]] = load x86_amx, x86_amx* [[TMP0]], align 64
+; CHECK-NEXT:    [[T5:%.*]] = load <256 x i32>, <256 x i32>* [[ARRAYIDX16:%.*]], align 64
 ; CHECK-NEXT:    br label [[FOR_COND18:%.*]]
 ; CHECK:       for.cond18:
-; CHECK-NEXT:    [[TMP1:%.*]] = phi x86_amx [ [[T51]], [[FOR_BODY14]] ], [ [[T11:%.*]], [[FOR_BODY24:%.*]] ]
+; CHECK-NEXT:    [[SUB_C_SROA_0_0:%.*]] = phi <256 x i32> [ [[T5]], [[FOR_BODY14]] ], [ [[T12:%.*]], [[FOR_BODY24:%.*]] ]
 ; CHECK-NEXT:    br i1 undef, label [[FOR_BODY24]], label [[FOR_COND_CLEANUP23:%.*]]
 ; CHECK:       for.cond.cleanup23:
-; CHECK-NEXT:    [[TMP2:%.*]] = bitcast x86_amx [[TMP1]] to <256 x i32>
-; CHECK-NEXT:    store <256 x i32> [[TMP2]], <256 x i32>* [[ARRAYIDX16]], align 64
+; CHECK-NEXT:    store <256 x i32> [[SUB_C_SROA_0_0]], <256 x i32>* [[ARRAYIDX16]], align 64
 ; CHECK-NEXT:    br label [[FOR_COND9]]
 ; CHECK:       for.body24:
 ; CHECK-NEXT:    [[T6:%.*]] = load <256 x i32>, <256 x i32>* [[ARRAYIDX29:%.*]], align 64
 ; CHECK-NEXT:    [[T7:%.*]] = load <256 x i32>, <256 x i32>* [[ARRAYIDX35:%.*]], align 64
+; CHECK-NEXT:    [[T8:%.*]] = bitcast <256 x i32> [[SUB_C_SROA_0_0]] to x86_amx
 ; CHECK-NEXT:    [[T9:%.*]] = bitcast <256 x i32> [[T6]] to x86_amx
 ; CHECK-NEXT:    [[T10:%.*]] = bitcast <256 x i32> [[T7]] to x86_amx
-; CHECK-NEXT:    [[T11]] = call x86_amx @llvm.x86.tdpbssd.internal(i16 1, i16 4, i16 4, x86_amx [[TMP1]], x86_amx [[T9]], x86_amx [[T10]])
+; CHECK-NEXT:    [[T11:%.*]] = call x86_amx @llvm.x86.tdpbssd.internal(i16 1, i16 4, i16 4, x86_amx [[T8]], x86_amx [[T9]], x86_amx [[T10]])
+; CHECK-NEXT:    [[T12]] = bitcast x86_amx [[T11]] to <256 x i32>
 ; CHECK-NEXT:    br label [[FOR_COND18]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %for.cond9
 for.cond9:                                        ; preds = %for.cond, %for.cond.cleanup23
   br i1 undef, label %for.body14, label %exit