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

(On Diff #245627)

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

2

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

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

3

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

17

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 ↗

(On Diff #245627)

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 ↗

(On Diff #245627)

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

1687 ↗

(On Diff #245627)

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 ↗

(On Diff #246670)

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 ↗

(On Diff #246670)

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/AArch64InstrInfo.h

void instantiateCondBranch(MachineBasicBlock &MBB, const DebugLoc &DL,
                              MachineBasicBlock *TBB,
                              ArrayRef<MachineOperand> Cond) const;
   bool substituteCmpToZero(MachineInstr &CmpInstr, unsigned SrcReg,
                            const MachineRegisterInfo *MRI) const;
 
   /// Returns an unused general-purpose register which can be used for
   /// constructing an outlined call if one exists. Returns 0 otherwise.
   unsigned findRegisterToSaveLRTo(const outliner::Candidate &C) const;
+
+  ///
+  /// Merge AND and TST instructions to ANDS when:
+  /// 1. AND and TST has the same operands
+  /// 2. Flag register (which is defined by TST) is not used between AND and ANDS
+  /// NB: TST is alias for ANDS xzr
+  ///
+  /// Example:
+  ///   \code
+  ///   tst x2, x1
+  ///   and x3, x2, x1
+  ///   \endcode
+  /// to
+  ///   \code
+  ///   ands x3, x2, x1
+  ///   \endcode
+  ///
+  /// \param  MBB BasicBlock in which instructions will be checked (and merged if suitable)
+  /// \return True when the instructions were merged
+  ///
+  bool mergeAndAnds(MachineBasicBlock& MBB) const;
 };
 
 /// emitFrameOffset - Emit instructions as needed to set DestReg to SrcReg
 /// plus Offset.  This is intended to be used from within the prolog/epilog
 /// insertion (PEI) pass, where a virtual scratch register may be allocated
 /// if necessary, to be replaced by the scavenger at the end of PEI.
 void emitFrameOffset(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
                      const DebugLoc &DL, unsigned DestReg, unsigned SrcReg,

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

 #include "AArch64InstrInfo.h"
 #include "AArch64MachineFunctionInfo.h"
 #include "AArch64Subtarget.h"
 #include "MCTargetDesc/AArch64AddressingModes.h"
 #include "Utils/AArch64BaseInfo.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"

 #include <iterator>
 #include <utility>
 
 using namespace llvm;
 
 #define GET_INSTRINFO_CTOR_DTOR
 #include "AArch64GenInstrInfo.inc"
 
+#define DEBUG_TYPE "aarch64-instr-info"
+STATISTIC(NumOfMergesTstAnd, "Number of merges of AND and TST instructions");
+
 static cl::opt<unsigned> TBZDisplacementBits(
     "aarch64-tbz-offset-bits", cl::Hidden, cl::init(14),
     cl::desc("Restrict range of TB[N]Z instructions (DEBUG)"));
 
 static cl::opt<unsigned> CBZDisplacementBits(
     "aarch64-cbz-offset-bits", cl::Hidden, cl::init(19),
     cl::desc("Restrict range of CB[N]Z instructions (DEBUG)"));
 

 ///    \code
 ///   tbnz w8, #10, L1
 ///    \endcode
 ///
 /// \param  MI Conditional Branch
 /// \return True when the simple conditional branch is generated
 ///
 bool AArch64InstrInfo::optimizeCondBranch(MachineInstr &MI) const {
+
+  if (mergeAndAnds(*MI.getParent()))
+    return true;
+
   bool IsNegativeBranch = false;
   bool IsTestAndBranch = false;
   unsigned TargetBBInMI = 0;
   switch (MI.getOpcode()) {
   default:
     llvm_unreachable("Unknown branch instruction?");
   case AArch64::Bcc:
     return false;
if (IsNegativeBranch)
       CC = AArch64CC::getInvertedCondCode(CC);
     BuildMI(RefToMBB, MI, DL, get(AArch64::Bcc)).addImm(CC).addMBB(TBB);
     MI.eraseFromParent();
     return true;
   }
   }
 }
 
+// Helper functions for mergeAndAnds
+static MachineBasicBlock::iterator FindInstrWithReg(const MachineBasicBlock::iterator IterBegin,
+                                                    const MachineBasicBlock::iterator IterEnd,
+                                                    const Register Reg,
+                                                    const TargetRegisterInfo& TRI) {
+  return std::find_if(IterBegin, IterEnd,
+                      [Reg, &TRI](const auto& MI){ return MI.readsRegister(Reg, &TRI); });
+}
+
+static bool HasSameOperands(const MachineInstr& FirstInstr,
+                            const MachineInstr& SecondInstr) {
+  const auto FirstOp1 = FirstInstr.getOperand(1).getReg();
+  const auto FirstOp2 = SecondInstr.getOperand(1).getReg();
+  const auto SecondOp1 = FirstInstr.getOperand(2).getReg();
+  const auto SecondOp2 = SecondInstr.getOperand(2).getReg();
+
+  return (FirstOp1 == FirstOp2 && SecondOp1 == SecondOp2)
+          ||(FirstOp1 == SecondOp2 && SecondOp1 == FirstOp2);
+}
+
+bool AArch64InstrInfo::mergeAndAnds(MachineBasicBlock& MBB) const {
+  const auto EndIter = MBB.end();
+  auto MIAndIter = EndIter;
+
+  for (auto CurrentInstrIter = MBB.begin();
+       CurrentInstrIter != EndIter; ++CurrentInstrIter) {
+
+    const auto Opcode = CurrentInstrIter->getOpcode();
+    if (Opcode == AArch64::ANDXrr || Opcode == AArch64::ANDWrr) {
+      assert(CurrentInstrIter->getNumOperands() == 3 && "Unexpected number of operands in ANDXrr");
+      MIAndIter = CurrentInstrIter;
+    }
+    else if ((Opcode == AArch64::ANDSXrr || Opcode == AArch64::ANDSWrr)
+              && (MIAndIter != EndIter)) {
+      assert(CurrentInstrIter->getNumOperands() == 4  && "Unexpected number of operands in ANDSXrr");
+
+      auto SiutableForMerge = [](const TargetRegisterInfo& TRI,
+                                  MachineBasicBlock::iterator& MIAndIter,
+                                  const MachineBasicBlock::iterator& CurrentInstrIter) {
+        MachineInstr* MIAnd = &*MIAndIter;
+        MachineInstr* MITst = &*CurrentInstrIter;
+        const auto MIterEnd = MIAnd->getParent()->end();
+
+        if (!HasSameOperands(*MIAnd, *MITst)) {
+          // Operands are differs, not a candidate for merge
+          return false;
+        }
+
+        const auto& RetVal = MITst->getOperand(0);
+        auto TmpIter = CurrentInstrIter;
+        assert(RetVal.isReg() && "Return value should be stored in register");
+        const auto RetValUseIter = FindInstrWithReg(++TmpIter, MIterEnd, RetVal.getReg(), TRI);
+
+        if (RetValUseIter != MIterEnd) {
+          // Return value of ANDS is used
+          // This case requires additional analysis
+          return false;
+        }
+
+        const auto FlagReg = MITst->getOperand(3);
+        assert(FlagReg.isReg() && "Third operand in ANDS instruction can be register only");
+        const auto FlagUseIter = FindInstrWithReg(MIAndIter, CurrentInstrIter, FlagReg.getReg(), TRI);
+
+        if (FlagUseIter != CurrentInstrIter) {
+          // Flag is used between AND and TST
+          // This case requires additional analysis
+          // Drop this instruction
+          MIAndIter = MIterEnd;
+          return false;
+        }
+
+        return true;
+      };
+
+      if (!SiutableForMerge(getRegisterInfo(), MIAndIter, CurrentInstrIter)) {
+        continue;
+      }
+
+      // All conditions are satisfied,
+      // instructions can be merged
+      MachineInstr* MIAnd = &*MIAndIter;
+      MachineInstr* MITst = &*CurrentInstrIter;
+      DebugLoc DL = MIAnd->getDebugLoc();
+      BuildMI(MBB, MIAndIter, DL, MITst->getDesc(), MIAnd->getOperand(0).getReg())
+        .addReg(MIAnd->getOperand(1).getReg())
+        .addReg(MIAnd->getOperand(2).getReg());
+      MIAnd->eraseFromParent();
+      MITst->eraseFromParent();
+      ++NumOfMergesTstAnd;
+      return true;
+    }
+  }
+
+  return false;
+}
+
 std::pair<unsigned, unsigned>
 AArch64InstrInfo::decomposeMachineOperandsTargetFlags(unsigned TF) const {
   const unsigned Mask = AArch64II::MO_FRAGMENT;
   return std::make_pair(TF & Mask, TF & ~Mask);
 }
 
 ArrayRef<std::pair<unsigned, const char *>>
 AArch64InstrInfo::getSerializableDirectMachineOperandTargetFlags() const {

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
+}