[AArch64] Peephole optimization. Merge AND and TST instructions

[kpdev42]

Summary

In some cases Clang does not perform merging of instructions AND and TST (aka ANDS xzr), like

Example:

tst x2, x1
and x3, x2, x1

to:

ands x3, x2, x1

This patch add such merging during instruction selection:
when AND is replaced with ANDS instruction in LowerSELECT_CC,
all users of AND also should be changed for using this ANDS instruction

Short discussion on mailing list
http://llvm.1065342.n5.nabble.com/llvm-dev-ARM-Peephole-optimization-instructions-tst-add-tp133109.html

[kpdev42]

Ping. I would be grateful for your commentaries about this patch. Maybe this code should be changed in some way?

[efriedma]

I don't have time to review this at the moment; adding some other reviewers.

Generally, if you don't get a response in a week, it's appropriate to say something. Sometimes reviews slip through the cracks. See http://llvm.org/docs/DeveloperPolicy.html#code-reviews

[SjoerdMeijer]

I don't mind taking a look.

First question about where this should live. With AArch64InstrInfo::mergeInstructions(MachineFunction& MF) you're iterating of all blocks/instructions, which is typically not something that belongs in AArch64InstrInfo because that works on instructions. The generic peephole class does exactly this, which makes calls to e.g. optimizeCondBranch. As Eli suggested in the email thread on llvm dev, can you implement and do this in AArch64InstrInfo::optimizeCondBranch(MachineInstr &MI)?

[dmgreen]

The other option, which might turn out to be simpler is to try and do this in ISel instead. There is a getNodeIfExists(..), which can be used to check if we have an ISD::AND, and an equivalent AArch64ISD::ANDS. Maybe done in performANDCombine?

This is kind of like doing CSE between two different nodes, but ISel is pretty good at that.

[kpdev42]

In D71701#1844832, @SjoerdMeijer wrote:

I don't mind taking a look.

First question about where this should live. With AArch64InstrInfo::mergeInstructions(MachineFunction& MF) you're iterating of all blocks/instructions, which is typically not something that belongs in AArch64InstrInfo because that works on instructions. The generic peephole class does exactly this, which makes calls to e.g. optimizeCondBranch. As Eli suggested in the email thread on llvm dev, can you implement and do this in AArch64InstrInfo::optimizeCondBranch(MachineInstr &MI)?

It is actually because I was confused that AArch64InstrInfo::optimizeCondBranch(MachineInstr &MI) takes only one instruction as parameter, but I need to iterate over BB to find necessary ANDS instruction. Of course I might get BB with getParent() of MI and iterate over it. But I thought that such interface would be kind of misleading (I mean - if function takes one instruction as a parameter, but works with BB, it should better takes BB as a parameter (or even function)). Seems to be that I am overthinking a bit here :(

I will move all my changes to AArch64InstrInfo::optimizeCondBranch(MachineInstr &MI) asap to proceed with review

In D71701#1845059, @dmgreen wrote:

The other option, which might turn out to be simpler is to try and do this in ISel instead. There is a getNodeIfExists(..), which can be used to check if we have an ISD::AND, and an equivalent AArch64ISD::ANDS. Maybe done in performANDCombine?

This is kind of like doing CSE between two different nodes, but ISel is pretty good at that.

It is an interesting suggestion, thank you. At first glance AArch64TargetLowering looks like more suitable place for this optimization, but I might be wrong.

Should I implement both approaches (in AArch64InstrInfo and in AArch64TargetLowering) to compare?

[dmgreen]

I think the ISel method might be a lot less code to have to deal with if it's done right. ISel already has a lot of the tools to do this kind of transform, even if it's not the normal kind of fold that you would usually do there.

If it is simple then I would recommend doing it there. That keeps all these kinds of optimisations in one place.

[craig.topper]

I assume the TST started as an ISD::SETCC and ISD::AND then the AArch64ISD::ANDS was generated from that. While still leaving the ISD::AND around for the use that wasn't the ISD::SETCC. On X86 when we create X86ISD::AND (which is our flag setting opcode) we replace all other uses of the ISD::AND with the data output of the X86ISD::AND. See EmitTest in X86ISelLowering. Maybe something like that would work here?

[kpdev42]

Update patch: move all changes to AArch64InstrInfo::optimizeCondBranch(MachineInstr &MI) as it was planned.
But I do not know if it will be the latest version, because as I see from the discussion - moving this functionality to ISel is more preferable.
So, do I need to create another patch with ISel implementation (I mean - create another differential revision), or it is better to change this patch?

In D71701#1846226, @craig.topper wrote:

I assume the TST started as an ISD::SETCC and ISD::AND then the AArch64ISD::ANDS was generated from that. While still leaving the ISD::AND around for the use that wasn't the ISD::SETCC. On X86 when we create X86ISD::AND (which is our flag setting opcode) we replace all other uses of the ISD::AND with the data output of the X86ISD::AND. See EmitTest in X86ISelLowering. Maybe something like that would work here?

Thank you! I definitely will take a look at it.

[SjoerdMeijer]

So, do I need to create another patch with ISel implementation (I mean - create another differential revision), or it is better to change this patch?

As is usually the case, there are multiple ways to achieve the same thing.... We don't need 2 implementations, but just pick the right approach. Dave suggested isel, and is open to alternatives. This patch is small and simple, and so I think this is a good place for it (as also remarked by Dave). And @craig.topper suggestion is very interesting, that will probably involve the least amount of code if something similar works for us, so that's probably best to check that first.

[kpdev42]

All changes are moved from InstrInfo to ISelLowering

[kpdev42]

Please take a look at this patch )

[SjoerdMeijer]

Nice one, that's very minimal now.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1682

Some real nits: there's probably very little benefit of declaring RHSOp and LHSOp here. For readability, it's probably better to just feed the LHS.getOperand() directly to getNode()

llvm/test/CodeGen/AArch64/peephole-and-tst.ll

1

You don't need the temp file, probably this is enough:

; RUN: llc < %s -mtriple=aarch64-- | FileCheck %s

2

Probably good to add a test case where the AND has multiple users.

16

It's not wrong, but you don't need the second check-label here, that can be a normal check. And it's probably easier to just let script utils/update_llc_test_checks.py generate all the checks for you.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1687

Can this just be DAG.ReplaceAllUsesWith(LHS, ANDSNode) ?

Otherwise it looks like it might try to treat &ANDSNode as an array.

[kpdev42]

Add a test case with two uses of AND instruction
Plus make some simplification of code in AArch64ISelLowering

[kpdev42]

All commentaries are resolved, patch is ready for review )

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1682

Thank you. I removed these temporary (and a little bit redundant :) ) variables

1687

Thank you. I changed call, this one looks much cleaner.

[SjoerdMeijer]

llvm/test/CodeGen/AArch64/peephole-and-tst.ll

33

Selection DAG isel has basic block scope, so I don't think this will be recognised as another use, and thus this would not be affected by a replaceAllUsesWith call when the transformation kicks in. So, best to move the 2nd use to the same block.

[kpdev42]

Update test with two uses of AND

[kpdev42]

llvm/test/CodeGen/AArch64/peephole-and-tst.ll

33

Ouch, mea culpa, fixed

[SjoerdMeijer]

Thanks, LGTM

[kpdev42]

Thank you for the review. (All of you)
Unfortunately I don't have commit access, would you please submit this change on my behalf? )

P.S: As I understand I can request commit access according to: https://llvm.org/docs/DeveloperPolicy.html#obtaining-commit-access ?

[SjoerdMeijer]

Yep, those are the instructions to request commit access, which might be good to request if you plan to do more llvm work. I guess that will probably take a few days. And if you're not in a hurry with getting this committed, you could wait for your commit access and then commit this yourself. But I would of course be happy to commit this on your behalf, which I will then do tomorrow morning, just let me know what you prefer.

[craig.topper]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

2574

nit. the && should be on the line above

[kpdev42]

Change formatting in AArch64ISelDAGToDAG.cpp according to codestyle

[kpdev42]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

2574

Thank you ) Fixed

[kpdev42]

In D71701#1894226, @SjoerdMeijer wrote:

Yep, those are the instructions to request commit access, which might be good to request if you plan to do more llvm work. I guess that will probably take a few days. And if you're not in a hurry with getting this committed, you could wait for your commit access and then commit this yourself. But I would of course be happy to commit this on your behalf, which I will then do tomorrow morning, just let me know what you prefer.

I will definitely apply for commit access (there is a lot of work ahead), but for this patch I would really appreciate if you commit it. :)

[SjoerdMeijer]

Sure, no problem, this is committed in: https://github.com/llvm/llvm-project/commit/13db7490fa67e22605dec4ab824121230b0fd928

[kpdev42]

Thank you. Should I close this revision?

[SjoerdMeijer]

That should happen automatically when you put the tag Differential Revision: ... in the commit message.
It doesn't look like a made spell mistake in it, so am a bit puzzled why this didn't happen.... but anyway, it's closed now.

[kpdev42]

Ok, thanks a lot )

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

else if (ShiftAmt->getOpcode() == ISD::SUB &&
       NewShiftAmt = SDValue(Neg, 0);
     } else
       return false;
   } else {
     // If the shift amount is masked with an AND, check that the mask covers the
     // bits that are implicitly ANDed off by the above opcodes and if so, skip
     // the AND.
     uint64_t MaskImm;
-    if (!isOpcWithIntImmediate(ShiftAmt.getNode(), ISD::AND, MaskImm))
+    if (!isOpcWithIntImmediate(ShiftAmt.getNode(), ISD::AND, MaskImm)
+        && !isOpcWithIntImmediate(ShiftAmt.getNode(), AArch64ISD::ANDS, MaskImm))
       return false;
 
     if (countTrailingOnes(MaskImm) < Bits)
       return false;
 
     NewShiftAmt = ShiftAmt->getOperand(0);
   }
 

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (isCMN(RHS, CC)) {
     // Can we combine a (CMP op1, (sub 0, op2) into a CMN instruction ?
     Opcode = AArch64ISD::ADDS;
     RHS = RHS.getOperand(1);
   } else if (isCMN(LHS, CC)) {
     // As we are looking for EQ/NE compares, the operands can be commuted ; can
     // we combine a (CMP (sub 0, op1), op2) into a CMN instruction ?
     Opcode = AArch64ISD::ADDS;
     LHS = LHS.getOperand(1);
-  } else if (LHS.getOpcode() == ISD::AND && isNullConstant(RHS) &&
-             !isUnsignedIntSetCC(CC)) {
+  } else if (isNullConstant(RHS) && !isUnsignedIntSetCC(CC)) {
+    if (LHS.getOpcode() == ISD::AND) {
       // Similarly, (CMP (and X, Y), 0) can be implemented with a TST
       // (a.k.a. ANDS) except that the flags are only guaranteed to work for one
       // of the signed comparisons.
-    Opcode = AArch64ISD::ANDS;
-    RHS = LHS.getOperand(1);
-    LHS = LHS.getOperand(0);
+      const SDValue RHSOp = LHS.getOperand(1);
+      const SDValue LHSOp = LHS.getOperand(0);
+      const SDValue ANDSNode = DAG.getNode(AArch64ISD::ANDS, dl,
+                                           DAG.getVTList(VT, MVT_CC),
+                                           LHSOp, RHSOp);
+      // Replace all users of (and X, Y) with newly generated (ands X, Y)
+      DAG.ReplaceAllUsesWith(LHS.getNode(), &ANDSNode);
+      return ANDSNode.getValue(1);
+    } else if (LHS.getOpcode() == AArch64ISD::ANDS) {
+      // Use result of ANDS comparison
+      return LHS.getValue(1);
+    }
   }
 
   return DAG.getNode(Opcode, dl, DAG.getVTList(VT, MVT_CC), LHS, RHS)
       .getValue(1);
 }
 
 /// \defgroup AArch64CCMP CMP;CCMP matching
 ///

llvm/test/CodeGen/AArch64/funnel-shift.ll

 declare i64 @llvm.fshr.i64(i64, i64, i64)
 declare <4 x i32> @llvm.fshr.v4i32(<4 x i32>, <4 x i32>, <4 x i32>)
 
 ; General case - all operands can be variables.
 
 define i32 @fshl_i32(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: fshl_i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    and w9, w2, #0x1f
+; CHECK-NEXT:    ands w9, w2, #0x1f
 ; CHECK-NEXT:    neg w9, w9
 ; CHECK-NEXT:    lsl w8, w0, w2
 ; CHECK-NEXT:    lsr w9, w1, w9
 ; CHECK-NEXT:    orr w8, w8, w9
-; CHECK-NEXT:    tst w2, #0x1f
 ; CHECK-NEXT:    csel w0, w0, w8, eq
 ; CHECK-NEXT:    ret
   %f = call i32 @llvm.fshl.i32(i32 %x, i32 %y, i32 %z)
   ret i32 %f
 }
 
 ; Verify that weird types are minimally supported.
 declare i37 @llvm.fshl.i37(i37, i37, i37)

 
 ; Repeat everything for funnel shift right.
 
 ; General case - all operands can be variables.
 
 define i32 @fshr_i32(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: fshr_i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    and w9, w2, #0x1f
+; CHECK-NEXT:    ands w9, w2, #0x1f
 ; CHECK-NEXT:    neg w9, w9
 ; CHECK-NEXT:    lsr w8, w1, w2
 ; CHECK-NEXT:    lsl w9, w0, w9
 ; CHECK-NEXT:    orr w8, w9, w8
-; CHECK-NEXT:    tst w2, #0x1f
 ; CHECK-NEXT:    csel w0, w1, w8, eq
 ; CHECK-NEXT:    ret
   %f = call i32 @llvm.fshr.i32(i32 %x, i32 %y, i32 %z)
   ret i32 %f
 }
 
 ; Verify that weird types are minimally supported.
 declare i37 @llvm.fshr.i37(i37, i37, i37)

llvm/test/CodeGen/AArch64/peephole-and-tst.ll

+; RUN: llc -mtriple=aarch64-linux-gnu %s -o %t; FileCheck %s --input-file=%t
+
+%struct.anon = type { i32*, i32* }
+
+@ptr_wrapper = common dso_local local_unnamed_addr global %struct.anon* null, align 8
+
+define dso_local i32 @test_func_i32(i32 %in, i32 %bit, i32 %mask) local_unnamed_addr {
+entry:
+  %0 = load %struct.anon*, %struct.anon** @ptr_wrapper, align 8
+  %result = getelementptr inbounds %struct.anon, %struct.anon* %0, i64 0, i32 1
+  %tobool2 = icmp ne i32 %mask, 0
+  br label %do.body
+
+do.body:                                          ; preds = %4, %entry
+; CHECK-LABEL: test_func_i32:
+; CHECK-LABEL: .LBB0_2:
+; CHECK: ands [[DSTREG:w[0-9]+]], [[SRCREG1:w[0-9]+]], [[SRCREG2:w[0-9]+]]
+  %bit.addr.0 = phi i32 [ %bit, %entry ], [ %shl, %4 ]
+  %retval1.0 = phi i32 [ 0, %entry ], [ %retval1.1, %4 ]
+  %and = and i32 %bit.addr.0, %in
+  %tobool = icmp eq i32 %and, 0
+  %not.tobool = xor i1 %tobool, true
+  %inc = zext i1 %not.tobool to i32
+  %retval1.1 = add nuw nsw i32 %retval1.0, %inc
+  %1 = xor i1 %tobool, true
+  %2 = or i1 %tobool2, %1
+  br i1 %2, label %3, label %4
+
+3:                                                ; preds = %do.body
+  store i32* null, i32** %result, align 8
+  br label %4
+
+4:                                                ; preds = %do.body, %3
+  %shl = shl i32 %bit.addr.0, 1
+  %tobool6 = icmp eq i32 %shl, 0
+  br i1 %tobool6, label %do.end, label %do.body
+
+do.end:                                           ; preds = %4
+  ret i32 %retval1.1
+}
+
+
+define dso_local i32 @test_func_i64(i64 %in, i64 %bit, i64 %mask) local_unnamed_addr #0 {
+entry:
+  %0 = load %struct.anon*, %struct.anon** @ptr_wrapper, align 8
+  %result = getelementptr inbounds %struct.anon, %struct.anon* %0, i64 0, i32 1
+  %tobool2 = icmp ne i64 %mask, 0
+  br label %do.body
+
+do.body:                                          ; preds = %4, %entry
+; CHECK-LABEL: test_func_i64:
+; CHECK-LABEL: .LBB1_2:
+; CHECK: ands [[DSTREG:x[0-9]+]], [[SRCREG1:x[0-9]+]], [[SRCREG2:x[0-9]+]]
+  %bit.addr.0 = phi i64 [ %bit, %entry ], [ %shl, %4 ]
+  %retval1.0 = phi i32 [ 0, %entry ], [ %retval1.1, %4 ]
+  %and = and i64 %bit.addr.0, %in
+  %tobool = icmp eq i64 %and, 0
+  %not.tobool = xor i1 %tobool, true
+  %inc = zext i1 %not.tobool to i32
+  %retval1.1 = add nuw nsw i32 %retval1.0, %inc
+  %1 = xor i1 %tobool, true
+  %2 = or i1 %tobool2, %1
+  br i1 %2, label %3, label %4
+
+3:                                                ; preds = %do.body
+  store i32* null, i32** %result, align 8
+  br label %4
+
+4:                                                ; preds = %do.body, %3
+  %shl = shl i64 %bit.addr.0, 1
+  %tobool6 = icmp eq i64 %shl, 0
+  br i1 %tobool6, label %do.end, label %do.body
+
+do.end:                                           ; preds = %4
+  ret i32 %retval1.1
+}

llvm/test/CodeGen/AArch64/shift-by-signext.ll

 }
 
 ; This is not valid for funnel shifts
 declare i32 @llvm.fshl.i32(i32 %a, i32 %b, i32 %c)
 declare i32 @llvm.fshr.i32(i32 %a, i32 %b, i32 %c)
 define i32 @n6_fshl(i32 %x, i32 %y, i8 %shamt) nounwind {
 ; CHECK-LABEL: n6_fshl:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    and w9, w2, #0x1f
+; CHECK-NEXT:    ands w9, w2, #0x1f
 ; CHECK-NEXT:    neg w9, w9
 ; CHECK-NEXT:    lsl w8, w0, w2
 ; CHECK-NEXT:    lsr w9, w1, w9
 ; CHECK-NEXT:    orr w8, w8, w9
-; CHECK-NEXT:    tst w2, #0x1f
 ; CHECK-NEXT:    csel w0, w0, w8, eq
 ; CHECK-NEXT:    ret
   %shamt_wide = sext i8 %shamt to i32
   %r = call i32 @llvm.fshl.i32(i32 %x, i32 %y, i32 %shamt_wide)
   ret i32 %r
 }
 define i32 @n7_fshr(i32 %x, i32 %y, i8 %shamt) nounwind {
 ; CHECK-LABEL: n7_fshr:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    and w9, w2, #0x1f
+; CHECK-NEXT:    ands w9, w2, #0x1f
 ; CHECK-NEXT:    neg w9, w9
 ; CHECK-NEXT:    lsr w8, w1, w2
 ; CHECK-NEXT:    lsl w9, w0, w9
 ; CHECK-NEXT:    orr w8, w9, w8
-; CHECK-NEXT:    tst w2, #0x1f
 ; CHECK-NEXT:    csel w0, w1, w8, eq
 ; CHECK-NEXT:    ret
   %shamt_wide = sext i8 %shamt to i32
   %r = call i32 @llvm.fshr.i32(i32 %x, i32 %y, i32 %shamt_wide)
   ret i32 %r
 }
 
 define i32 @n8_extrause(i32 %x, i8 %shamt, i32* %shamt_wide_store) nounwind {