[AArch64][Isel] Avoid implicit zext for SIGN_EXTEND_INREG (TRUNCATE)

[guopeilin]

Summary

In some cases, AArch64 use SBFMXri + EXTRACT_SUBREG to implement "sign_extend_inreg (truncate)",
However COPY is usually erased after Register Coalescing, so if we want to zext this result,
we need to avoid using implicit zero-extend.

[t.p.northover]

llvm/lib/Target/AArch64/AArch64ISelLowering.h

436

Could we fix isBitfieldExtractOpFromSExtInReg to produce SBFMWri instead?

The other cases here seem to be nodes that we'd have to pessimize in the general case to ensure they zero the high bits on the off-chance they're needed, but for sign-extension it's just a quirk of the implementation.

llvm/test/CodeGen/AArch64/aarch64-avoid-implicit-zext-for-SIGN_EXTEND_INREG.ll

2 ↗

(On Diff #305402)

Could you try to simplify this a bit more? There are lots of stray annotations, both calls appear to be unnecessary, and as far as I can tell %bf.load could just as easily be a function argument.

[guopeilin]

llvm/lib/Target/AArch64/AArch64ISelLowering.h

436

I guess the main job of isBitfieldExtractOpFromSExtInReg is trying to change Truncate + SBFMWri into SBFMXri + COPY, So I think if we fix it to produce SBFMWri, this function may become meaningless.
And, I think maybe it's ok to use a pattern like SBFMXri + COPY. The problem is that COPY is likely to be erased after Register Coalescing. And I found a patch that modifies the function shouleCoalesce() (https://reviews.llvm.org/D85956) to avoid coalescing COPY in some cases. However, this solution will cause performance regression and cannot fix this case.
So I think the key point is how to keep the operation that extract the lower bits after Register Coalescing, and I found this patch (https://reviews.llvm.org/D87771). So when we zero extend sign_exten_inreg, we need to look forward, if it comes from a truncate, then we should not use implicit zero extend.

llvm/test/CodeGen/AArch64/aarch64-avoid-implicit-zext-for-SIGN_EXTEND_INREG.ll

2 ↗

(On Diff #305402)

Yes, these calls can be eliminated. However, with these calls, we can directly see the difference of runtime results between -O0 and -O3 after generating a binary program from this .ll file.
So I guess maybe it's ok to keep these calls, although it may look redundant.

[t.p.northover]

llvm/lib/Target/AArch64/AArch64ISelLowering.h

436

I don't think that'''s what isBitfieldExtractOpFromSExtInReg is trying to do. Any COPY it results in would be far too implicit, and as you note completely unrelable. We just can't rely on it being preserved for correctness and shouldn't try.

Instead, I see a function that tries to select the correct SBFM variant, but is tripped up by the truncate detection code changing VT under it. N itself produces an i32 so SBFMWri is correct, but that's forgotten in the truncate case.

llvm/test/CodeGen/AArch64/aarch64-avoid-implicit-zext-for-SIGN_EXTEND_INREG.ll

2 ↗

(On Diff #305402)

.ll tests aren't designed to be executed and we don't optimize them for that.

[t.p.northover]

llvm/lib/Target/AArch64/AArch64ISelLowering.h

436

Oh hang on, I think I see what you mean. Let me think a bit more.

[t.p.northover]

llvm/lib/Target/AArch64/AArch64ISelLowering.h

436

OK, I'm convinced, but please add a comment explaining why this is a special case, and probably mentioning isBitfieldExtractOpFromSExtInReg explicitly as the source.

[guopeilin]

llvm/lib/Target/AArch64/AArch64ISelLowering.h

436

Ok, I will add this comment, and a simplifier case will be added later, please wait a bit more.

[guopeilin]

llvm/lib/Target/AArch64/AArch64ISelLowering.h

436

a isBitfieldExtractOpFromSExtInReg related comment and a simpler case are added

[RKSimon]

aarch64-avoid-implicit-zext-for-SIGN_EXTEND_INREG.ll - please can you avoid using capital letters? Every so often it causes problems on windows dev machines.... (TBH its been better since the git transition).

[guopeilin]

In D91512#2417984, @RKSimon wrote:

aarch64-avoid-implicit-zext-for-SIGN_EXTEND_INREG.ll - please can you avoid using capital letters? Every so often it causes problems on windows dev machines.... (TBH its been better since the git transition).

OK, the filename has been changed.

[guopeilin]

llvm/lib/Target/AArch64/AArch64ISelLowering.h

436

If everything is ok, could you please approve this patch? Thanks a lot!

[guopeilin]

Pinging reviewers...

[craig.topper]

llvm/lib/Target/AArch64/AArch64ISelLowering.h

437

Should COPY here be EXTRACT_SUBREG? I think that’s what it is when isel creates it.

[guopeilin]

llvm/lib/Target/AArch64/AArch64ISelLowering.h

437

Thanks for reviewing, the comment is now modified.

[jaykang10]

I wonder we can extend this check to other basic blocks.

For example, if the trunc is in other basic block, there could be CopyFromReg instead of trunc in current basic block. Can we detect the trunc in this case please?

[efriedma]

In D91512#2999435, @jaykang10 wrote:

I wonder we can extend this check to other basic blocks.

For example, if the trunc is in other basic block, there could be CopyFromReg instead of trunc in current basic block. Can we detect the trunc in this case please?

Ideally, I think we should kill off isDef32, and run a pass after isel to clean up redundant extension operations. That would have fewer bugs, and we'd get cross-block optimization for free.

[jaykang10]

In D91512#3000179, @efriedma wrote:

In D91512#2999435, @jaykang10 wrote:

I wonder we can extend this check to other basic blocks.

For example, if the trunc is in other basic block, there could be CopyFromReg instead of trunc in current basic block. Can we detect the trunc in this case please?

Ideally, I think we should kill off isDef32, and run a pass after isel to clean up redundant extension operations. That would have fewer bugs, and we'd get cross-block optimization for free.

I agree with you. I feel we need custom peephole optimization pass for AArch64. If there is no working patch for it, I could start to implement it.

llvm/lib/Target/AArch64/AArch64ISelLowering.h

 // Any instruction that defines a 32-bit result zeros out the high half of the
 // register. Truncate can be lowered to EXTRACT_SUBREG. CopyFromReg may
 // be copying from a truncate. But any other 32-bit operation will zero-extend
 // up to 64 bits. AssertSext/AssertZext aren't saying anything about the upper
 // 32 bits, they're probably just qualifying a CopyFromReg.
 // FIXME: X86 also checks for CMOV here. Do we need something similar?
 static inline bool isDef32(const SDNode &N) {
   unsigned Opc = N.getOpcode();
+
+  // For some cases, isBitfieldExtractOpFromSExtInReg() will try to use
+  // SBFMXri + COPY to implement sign_extend_inreg(truncate) during ISel,
+  // However COPY is likely to be erased after register coalescing, So if
+  // we want to zext this value, we need avoid using implicit way.
+  if (Opc == ISD::SIGN_EXTEND_INREG) {
+    SDValue Opr = N.getOperand(0);
+    return Opr->getOpcode() != ISD::TRUNCATE;
+  }
+
   return Opc != ISD::TRUNCATE && Opc != TargetOpcode::EXTRACT_SUBREG &&
          Opc != ISD::CopyFromReg && Opc != ISD::AssertSext &&
          Opc != ISD::AssertZext;
 }
 
 } // end anonymous namespace
 
 class AArch64Subtarget;

llvm/test/CodeGen/AArch64/aarch64-avoid-implicit-zext-for-sign_extend_inreg.ll

+; RUN: llc -O3 -mtriple=aarch64-unknown-linux-gnu < %s | FileCheck %s
+@c = common dso_local global i64 0, align 8
+
+define dso_local i32 @main(i64 %bf.load) {
+entry:
+  %0 = lshr i64 %bf.load, 30
+  %sext3 = shl nuw nsw i64 %0, 16
+  %sext = trunc i64 %sext3 to i32
+  %conv.i = ashr exact i32 %sext, 16
+  %conv.i.i = zext i32 %conv.i to i64
+; CHECK:        // %bb.0:
+; CHECK:        sbfx    x{{[0-9]+}}, x{{[0-9]+}}, #{{[0-9]+}}, #{{[0-9]+}}
+; CHECK-NEXT:   mov     w{{[0-9]+}}, w{{[0-9]+}}
+  store i64 %conv.i.i, i64* @c, align 8
+  ret i32 0
+}