AArch64: Fold immediate into the immediate field of logical instructions

[ahatanak]

Summary

llvm currently turns the following code

void foo1(int a, char *p) {

int t = a & 0xfd;
*p = t;

}

into these instructions:

movz w8, #253
and w8, w0, w8
strb w8, [x1]

This can be done using just two instructions, since we don't care what the upper 24-bits of the "and" instruction are.

and w8, w0, 0xfffffffd
strb w8, [x1]

This patch adds a target hook to TargetLowering::TargetLoweringOpt and overrides it in the AArch64 backend to sign-extend an immediate operand if the upper bits are not demanded and sign-extending enables folding the immediate into the instruction.

This optimization speeds up 253.perlbmk by 5%.

[mcrosier]

Nice gain!! Just a few nits.

include/llvm/Target/TargetLowering.h

2056

s/LTO/TLO ?

lib/Target/AArch64/AArch64ISelLowering.cpp

674

Remove vertical whitespace.

681

Please add an assert message:

assert(cond && "msg");

[ahatanak]

Address Chad's comments.

[ahatanak]

ping

lib/Target/AArch64/AArch64ISelLowering.cpp

685

The code I previously had here didn't make sense, so I cleaned it up. It was checking Size > 0 after Size = std::max(VT.getSizeInBits(), 32u). Also, I changed it to break if VT is a vector, just in case this function is called on a vector node.

[t.p.northover]

Hi Akira,

I've got a couple of comments:

include/llvm/Target/TargetLowering.h

2058

Functions should usually start with a lower-case letter.

lib/Target/AArch64/AArch64ISelLowering.cpp

693–694

OK, so we've got an immediate "ab...iJK..Z" where lower-case digits aren't used. Sign extending converts this to "JJ...JJK...Z". Is there any particular reason to expect that's representable? It seems like a bit of a shot in the dark.

700

Shifting a negative number left is undefined behaviour.

[ahatanak]

Tim, I fixed the undefined behavior and renamed the function to start with lower-case letter.

The reason I only do sign-extension is that it seemed to be the cheapest way to get the most gain without hurting compile time or making the code overly complicated. I looked at the instructions llvm emits, and I saw many places where logical instructions were followed by truncating stores.

I can try changing it to do a more exhaustive search of the bit patterns to see how much further performance can be improved, if that's necessary.

[t.p.northover]

Hi Akira,

The reason I only do sign-extension is that it seemed to be the cheapest way to get the most gain without hurting compile time or making the code overly complicated. I looked at the instructions llvm emits, and I saw many places where logical instructions were followed by truncating stores.

I think the kind of store being done is largely orthogonal to the
calculation most likely to get you a valid immediate. It tells you
what bits can be ignored, but nothing about the best way to fill them.

So I think it's reasonable to start with just assuming that
DemandedBits has some number of low bits set and the rest ignored
based on your observations. I think it's harder to justify coming up
with a single NewImm value by sign extending the existing immediate
and giving up if that fails.

Cheers.

Tim.

[t.p.northover]

Hi again,

I think it's harder to justify coming up
with a single NewImm value by sign extending the existing immediate
and giving up if that fails.

I've been doing some more thinking here, and if we're willing to
assume the truncation is to a power of two type (I am, anyone using
i14 deserves whatever they get) then I think we can cover *all* valid
cases by instead replicating the demanded bits across the 32-bit
register. E.g. try 0xfdfdfdfd instead of 0xfffffffd for the 8-bit
0xfd.

The argument goes that if the input is morally contiguous, then there
are multiple representations involving sign extension to 4, 8, 16 or
32 followed by replication. Otherwise the replication width is less
than the demanded width so we're completely forced and have to
continue the replication that's already started.

Cheers.

Tim.

[ahatanak]

In D5591#15, @t.p.northover wrote:

The argument goes that if the input is morally contiguous, then there
are multiple representations involving sign extension to 4, 8, 16 or
32 followed by replication. Otherwise the replication width is less
than the demanded width so we're completely forced and have to
continue the replication that's already started.

Sign-extension followed by replication enables converting constants like 0xfdfd (demanded = 16-bits) or 0x3dfd (demanded = 14bits) to bimm, where just sign-extending fails, but I think we should also try to handle cases like 0x19 (demanded = 5bits). This isn't a bimm as it stands and sign-extending doesn't make it a bimm either. In this case, we have to copy bits 1-3 to bits 5-7.

I think I can come up with a patch that does this.

[ahatanak]

Rewrote the algorithm for searching for a bitmask immediate based on Tim's feedback.

[ahatanak]

I had to make changes to DAGCombiner::visitAND because it was transforming the DAG in way that made it impossible to do any optimization in optimizeConstant (the code here was originally committed in r97616). This change doesn't seem to have any noticeable impact on performance, but I'm still investigating.

[ahatanak]

Sorry for not updating the patch for a long time. Here is my new patch.

For the most part, the new patch takes the same approach as the previous patch to find an immediate operand, but there were a couple of changes made.

• Function optimizeLogicalImm emits AArch64 machine nodes to prevent the target-independent dagcombine from undoing the optimization. With this change, there is no need to make changes in DAGCombiner::visitOR as I did in my previous patch.
• In optimizeLogicalImm, rotation is used to avoid using branches and simplify the logic.
• A target-specific function object for optimizing nodes with immediates is passed to the constructor of TargetLoweringOpt, which gets called later in TargetLoweringOpt::ShrinkDemandedConstant.

[ahatanak]

Rebase and make a couple of changes:

• Add comments.
• In optimizeLogicalImm, return false instead of true when the immediate is already a bimm32 or bimm64 so that it doesn't inhibit the optimization done later in AArch64ISelDAGToDAG.cpp that emits BFXIL.
• Remove redundant instructions in bitreverse.ll.

[ab]

Some drive-by comments; I haven't looked into the optimizeLogicalImm logic yet.

include/llvm/Target/TargetLowering.h

2026

Add more detail on when this is called and what the parameters are?

2036–2039

I don't think this is the best location for this: I'd rather have TargetLoweringOpt be "the result of an optimization", and TargetLowering be "how to do optimizations".

What do you think of going back to the TLI virtual hook, and fixing the other TLO methods to do the same: https://reviews.llvm.org/differential/diff/73495/

lib/Target/AArch64/AArch64ISelLowering.cpp

672–673

Check this only once when initializing the std::function? (or leave it here if you remove the std::function, I suppose)

test/CodeGen/AArch64/optimize-imm.ll

2

Triple can be simplified to: -mtriple=aarch64--

[ahatanak]

include/llvm/Target/TargetLowering.h

2036–2039

I'm not sure what TargetLoweringOpt is supposed to do, but using TLI virtual hooks looks like a better approach.

[ahatanak]

Address review comments.

[ahatanak]

Fix comment.

[ahatanak]

Rebase. Move TargetLoweringOpt::SimplifyDemandedBits into TargetLowering.

[ahatanak]

Rebase.

[mcrosier]

lib/Target/AArch64/AArch64ISelLowering.cpp

676

You should use AArch64_AM::isLogicalImmediate() here.

771

If you put this after the below switch, you'll guarantee Op has operand 1.

777

Can't we early exit if we demand all of the bits? E.g.,

if (Demanded.countPopulation() == Size)
  return false;

792

Default case goes at top of switch, per coding guidelines.

[mcrosier]

FWIW, I think this is in pretty good shape. Also, I ran correctness tests on everything I've got (e.g., llvm-ts, SPEC200X, internal tests) and saw no correctness failures.

[ahatanak]

Address Chad's comments.

[mcrosier]

Thanks for working on this, Akira. I believe you've addressed all of my concerns as well as Tim's and Arnaud's. LGTM, assuming you've done the due diligence to ensure this doesn't dramatically increase compile-time.

[ahatanak]

I compiled the files in MultiSource/Applications of the test suite and didn't see any measurable increase in compile time. I'll commit this patch today.

Thank you for the review.

[mcrosier]

Thanks for working on this patch and for being *extremely* patient with the 2.5 year review, Akira!

include/llvm/Target/TargetLowering.h

std::pair<SDValue, SDValue> makeLibCall(SelectionDAG &DAG, RTLIB::Libcall LC,
                                           EVT RetVT, const SDValue *Ops,
                                           unsigned NumOps, bool isSigned,
                                           SDLoc dl, bool doesNotReturn = false,
                                           bool isReturnValueUsed = true) const;
 
   //===--------------------------------------------------------------------===//
   // TargetLowering Optimization Methods
   //
 
   /// A convenience struct that encapsulates a DAG, and two SDValues for
   /// returning information from TargetLowering to its clients that want to
   /// combine.
   struct TargetLoweringOpt {
     SelectionDAG &DAG;
     bool LegalTys;
     bool LegalOps;
     SDValue Old;
     SDValue New;
 
     explicit TargetLoweringOpt(SelectionDAG &InDAG,
                                bool LT, bool LO) :
       DAG(InDAG), LegalTys(LT), LegalOps(LO) {}
 
     bool LegalTypes() const { return LegalTys; }
     bool LegalOperations() const { return LegalOps; }
 
     bool CombineTo(SDValue O, SDValue N) {
       Old = O;
       New = N;
       return true;
     }
 
     /// Check to see if the specified operand of the specified instruction is a
     /// constant integer.  If so, check to see if there are any bits set in the
     /// constant that are not demanded.  If so, shrink the constant and return
     /// true.
     bool ShrinkDemandedConstant(SDValue Op, const APInt &Demanded);
 
     /// Convert x+y to (VT)((SmallVT)x+(SmallVT)y) if the casts are free.  This
     /// uses isZExtFree and ZERO_EXTEND for the widening cast, but it could be
     /// generalized for targets with other types of implicit widening casts.
     bool ShrinkDemandedOp(SDValue Op, unsigned BitWidth, const APInt &Demanded,
                           SDLoc dl);
   };
 
   /// Look at Op.  At this point, we know that only the DemandedMask bits of the
   /// result of Op are ever used downstream.  If we can use this information to
   /// simplify Op, create a new simplified DAG node and return true, returning
   /// the original and new nodes in Old and New.  Otherwise, analyze the
   /// expression and return a mask of KnownOne and KnownZero bits for the
   /// expression (used to simplify the caller).  The KnownZero/One bits may only
   /// be accurate for those bits in the DemandedMask.
   bool SimplifyDemandedBits(SDValue Op, const APInt &DemandedMask,
                             APInt &KnownZero, APInt &KnownOne,
                             TargetLoweringOpt &TLO, unsigned Depth = 0) const;
 
+  /// Do a target-specific optimization of Op's constant immediate operand using
+  /// the demanded bits information. For example, targets can sign-extend the
+  /// immediate in order to have it folded into the immediate field of the
+  /// instruction. Return true if the immediate node doesn't need further
+  /// optimization and set the New and Old fields of TLO if a new immediate
+  /// node was created.
+  virtual bool optimizeConstant(SDValue Op, const APInt &Demanded,
+                                TargetLoweringOpt &TLO) const;
+
   /// Determine which of the bits specified in Mask are known to be either zero
   /// or one and return them in the KnownZero/KnownOne bitsets.
   virtual void computeKnownBitsForTargetNode(const SDValue Op,
                                              APInt &KnownZero,
                                              APInt &KnownOne,
                                              const SelectionDAG &DAG,
                                              unsigned Depth = 0) const;
 

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitOR(SDNode *N) {
   if (BSwap.getNode())
     return BSwap;
 
   // reassociate or
   SDValue ROR = ReassociateOps(ISD::OR, SDLoc(N), N0, N1);
   if (ROR.getNode())
     return ROR;
   // Canonicalize (or (and X, c1), c2) -> (and (or X, c2), c1|c2)
-  // iff (c1 & c2) == 0.
+  // iff (c1 & c2) == c1.
   if (N1C && N0.getOpcode() == ISD::AND && N0.getNode()->hasOneUse() &&
              isa<ConstantSDNode>(N0.getOperand(1))) {
     ConstantSDNode *C1 = cast<ConstantSDNode>(N0.getOperand(1));
-    if ((C1->getAPIntValue() & N1C->getAPIntValue()) != 0) {
+    if ((C1->getAPIntValue() & N1C->getAPIntValue()) == C1->getAPIntValue()) {
       SDValue COR = DAG.FoldConstantArithmetic(ISD::OR, VT, N1C, C1);
       if (!COR.getNode())
         return SDValue();
       return DAG.getNode(ISD::AND, SDLoc(N), VT,
                          DAG.getNode(ISD::OR, SDLoc(N0), VT,
                                      N0.getOperand(0), N1), COR);
     }
   }

lib/CodeGen/SelectionDAG/TargetLowering.cpp

 /// ShrinkDemandedConstant - Check to see if the specified operand of the
 /// specified instruction is a constant integer.  If so, check to see if there
 /// are any bits set in the constant that are not demanded.  If so, shrink the
 /// constant and return true.
 bool TargetLowering::TargetLoweringOpt::ShrinkDemandedConstant(SDValue Op,
                                                         const APInt &Demanded) {
   SDLoc dl(Op);
 
+  if (Op.getOpcode() == ISD::XOR) {
+    ConstantSDNode *C = dyn_cast<ConstantSDNode>(Op.getOperand(1));
+
+    // If an XOR already has all the bits set, nothing to change but don't
+    // shrink either.
+    if (!C || (C->getAPIntValue() | (~Demanded)).isAllOnesValue())
+      return false;
+  }
+
+  assert(!Old.getNode() && !New.getNode());
+
+  // Return if the constant operand doesn't need further optimization.
+  if (DAG.getTargetLoweringInfo().optimizeConstant(Op, Demanded, *this))
+    return New.getNode();
+
   // FIXME: ISD::SELECT, ISD::SELECT_CC
   switch (Op.getOpcode()) {
   default: break;
   case ISD::XOR:
   case ISD::AND:
   case ISD::OR: {
     ConstantSDNode *C = dyn_cast<ConstantSDNode>(Op.getOperand(1));
     if (!C) return false;
 
-    if (Op.getOpcode() == ISD::XOR &&
-        (C->getAPIntValue() | (~Demanded)).isAllOnesValue())
-      return false;
-
     // if we can expand it to have all bits set, do it
     if (C->getAPIntValue().intersects(~Demanded)) {
       EVT VT = Op.getValueType();
       SDValue New = DAG.getNode(Op.getOpcode(), dl, VT, Op.getOperand(0),
                                 DAG.getConstant(Demanded &
                                                 C->getAPIntValue(),
                                                 VT));
       return CombineTo(Op, New);
bool TargetLowering::SimplifyDemandedBits(SDValue Op,
   // If we know the value of all of the demanded bits, return this as a
   // constant.
   if ((NewMask & (KnownZero|KnownOne)) == NewMask)
     return TLO.CombineTo(Op, TLO.DAG.getConstant(KnownOne, Op.getValueType()));
 
   return false;
 }
 
+bool TargetLowering::optimizeConstant(SDValue Op, const APInt &Demanded,
+                                      TargetLoweringOpt &TLO) const {
+  return false;
+}
+
 /// computeKnownBitsForTargetNode - Determine which of the bits specified
 /// in Mask are known to be either zero or one and return them in the
 /// KnownZero/KnownOne bitsets.
 void TargetLowering::computeKnownBitsForTargetNode(const SDValue Op,
                                                    APInt &KnownZero,
                                                    APInt &KnownOne,
                                                    const SelectionDAG &DAG,
                                                    unsigned Depth) const {

lib/Target/AArch64/AArch64ISelLowering.h

class AArch64TargetLowering : public TargetLowering {
   bool RequireStrictAlign;
 
 public:
   explicit AArch64TargetLowering(const TargetMachine &TM);
 
   /// Selects the correct CCAssignFn for a given CallingConvention value.
   CCAssignFn *CCAssignFnForCall(CallingConv::ID CC, bool IsVarArg) const;
 
+  bool optimizeConstant(SDValue Op, const APInt &Demanded,
+                        TargetLoweringOpt &TLO) const override;
+
   /// computeKnownBitsForTargetNode - Determine which of the bits specified in
   /// Mask are known to be either zero or one and return them in the
   /// KnownZero/KnownOne bitsets.
   void computeKnownBitsForTargetNode(const SDValue Op, APInt &KnownZero,
                                      APInt &KnownOne, const SelectionDAG &DAG,
                                      unsigned Depth = 0) const override;
 
   MVT getScalarShiftAmountTy(EVT LHSTy) const override;

lib/Target/AArch64/AArch64ISelLowering.cpp

 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetOptions.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "aarch64-lower"
 
 STATISTIC(NumTailCalls, "Number of tail calls");
 STATISTIC(NumShiftInserts, "Number of vector shift inserts");
+STATISTIC(NumOptimizedImms, "Number of times immediates were optimized");
 
 namespace {
 enum AlignMode {
   StrictAlign,
   NoStrictAlign
 };
 }
 

 }
 
 EVT AArch64TargetLowering::getSetCCResultType(LLVMContext &, EVT VT) const {
   if (!VT.isVector())
     return MVT::i32;
   return VT.changeVectorElementTypeToInteger();
 }
 
+static bool optimizeConstant(unsigned Size, uint64_t Imm, uint64_t &NewImm,
+                             const APInt &Demanded) {
+  uint64_t Enc;
+  uint64_t Mask = ((uint64_t)(-1LL) >> (64 - Size));
+
+  // Return if the immediate is already a bimm32 or bimm64.
+  if (AArch64_AM::processLogicalImmediate(Imm & Mask, Size, Enc)) {
+    NewImm = Imm;
+    return true;
+  }
+
+  unsigned EltSize = Size;
+  uint64_t DemandedBits = Demanded.getZExtValue();
+
+  // Clear bits that are not demanded.
+  Imm &= DemandedBits;
+
+  while (true) {
+    // The goal here is to set the non-demanded bits in a way that minimizes
+    // the number of switching between 0 and 1. In order to achieve this goal,
+    // we set the non-demanded bits to the value of the preceding demanded bits.
+    // For example, if we have an immediate 0bx10xx0x1 ('x' indicates a
+    // non-demanded bit), we copy bit0 (1) to the least significant 'x',
+    // bit2 (0) to 'xx', and bit6 (1) to the most significant 'x'.
+    // The final result is 0b11000011.
+    uint64_t NonDemandedBits = ~DemandedBits;
+    uint64_t Ones =
+        (NonDemandedBits + ((~Imm & DemandedBits) << 1)) & NonDemandedBits;
+
+    // We cannot copy the preceding demanded bit if the rightmost bits are not
+    // demanded. In that case, we replicate the least significant demanded bit
+    // that follows the trailing non-demanded bits.
+    if (Ones & 1) {
+      bool FirstDemandedBitIsZero = ~Imm & ((Ones + 1 | Ones) ^ Ones) & Mask;
+
+      // If the first demanded bit is zero, clear the trailing ones in Ones.
+      if (FirstDemandedBitIsZero)
+        Ones = (Ones + 1) & Ones;
+    }
+
+    uint64_t TmpImm = (Imm | Ones) & Mask;
+
+    // If TmpImm or its bitwise NOT is a shifted mask, it is a bitmask immediate
+    // or all-ones or all-zeros, in which case we can stop searching. Otherwise,
+    // we halve the element size and continue the search.
+    if (isShiftedMask_64(TmpImm) || isShiftedMask_64(~(TmpImm | ~Mask))) {
+      Imm = TmpImm;
+      break;
+    }
+
+    // We cannot shrink the element size any further if it is 2-bits.
+    if (EltSize == 2)
+      return false;
+
+    EltSize /= 2;
+    Mask >>= EltSize;
+    uint64_t Hi = Imm >> EltSize, DemandedBitsHi = DemandedBits >> EltSize;
+
+    // Return if there is mismatch in any of the demanded bits of Imm and Hi.
+    if (((Imm ^ Hi) & (DemandedBits & DemandedBitsHi) & Mask) != 0)
+      return false;
+
+    // Merge the upper and lower halves of Imm and DemandedBits.
+    Imm |= Hi;
+    DemandedBits |= DemandedBitsHi;
+  }
+
+  // Replicate the element across the register width.
+  while (EltSize < Size) {
+    Imm |= Imm << EltSize;
+    EltSize *= 2;
+  }
+
+  NewImm = Imm;
+  return true;
+}
+
+bool AArch64TargetLowering::optimizeConstant(SDValue Op, const APInt &Demanded,
+                                             TargetLoweringOpt &TLO) const {
+  // Delay this optimization to as late as possible.
+  if (!TLO.LegalOps)
+    return false;
+
+  switch (Op.getOpcode()) {
+  default: break;
+  case ISD::AND:
+  case ISD::OR:
+  case ISD::XOR:
+    ConstantSDNode *C = dyn_cast<ConstantSDNode>(Op.getOperand(1));
+    if (!C)
+      break;
+
+    EVT VT = Op.getValueType();
+    if (VT.isVector())
+      break;
+
+    unsigned Size = VT.getSizeInBits();
+    assert((Size == 32 || Size == 64) &&
+           "i32 or i64 is expected after legalization.");
+    uint64_t Imm = C->getZExtValue(), NewImm;
+
+    // Break if we cannot turn Imm into a bitmask immediate.
+    if (!::optimizeConstant(Size, Imm, NewImm, Demanded))
+      break;
+
+    // Return if Imm was already a bitmask immediate.
+    if (Imm == NewImm)
+      return true;
+
+    ++NumOptimizedImms;
+    assert(((NewImm ^ Imm) & Demanded.getZExtValue()) == 0 &&
+           "demanded bits should never be altered");
+
+    // Create the new constant immediate node.
+    SDValue New = TLO.DAG.getNode(Op.getOpcode(), SDLoc(Op), VT,
+                                  Op.getOperand(0),
+                                  TLO.DAG.getConstant(NewImm, VT));
+    return TLO.CombineTo(Op, New);
+  }
+
+  return false;
+}
+
 /// computeKnownBitsForTargetNode - Determine which of the bits specified in
 /// Mask are known to be either zero or one and return them in the
 /// KnownZero/KnownOne bitsets.
 void AArch64TargetLowering::computeKnownBitsForTargetNode(
     const SDValue Op, APInt &KnownZero, APInt &KnownOne,
     const SelectionDAG &DAG, unsigned Depth) const {
   switch (Op.getOpcode()) {
   default:
     break;
   case AArch64ISD::CSEL: {

test/CodeGen/AArch64/optimize-imm.ll

+; RUN: llc -o - %s -march=aarch64 | FileCheck %s
+
+; CHECK-LABEL: _and1:
+; CHECK: and {{w[0-9]+}}, w0, #0xfffffffd
+
+define void @and1(i32 %a, i8* nocapture %p) {
+entry:
+  %and = and i32 %a, 253
+  %conv = trunc i32 %and to i8
+  store i8 %conv, i8* %p, align 1
+  ret void
+}
+
+; (a & 0x3dfd) | 0xffffc000
+;
+; CHECK-LABEL: _and2:
+; CHECK: and {{w[0-9]+}}, w0, #0xfdfdfdfd
+
+define i32 @and2(i32 %a) {
+entry:
+  %and = and i32 %a, 15869
+  %or = or i32 %and, -16384
+  ret i32 %or
+}
+
+; (a & 0x19) | 0xffffffc0
+;
+; CHECK-LABEL: _and3:
+; CHECK: and {{w[0-9]+}}, w0, #0x99999999
+
+define i32 @and3(i32 %a) {
+entry:
+  %and = and i32 %a, 25
+  %or = or i32 %and, -64
+  ret i32 %or
+}
+
+; (a & 0xc5600) | 0xfff1f1ff
+;
+; CHECK-LABEL: _and4:
+; CHECK: and {{w[0-9]+}}, w0, #0xfffc07ff
+
+define i32 @and4(i32 %a) {
+entry:
+  %and = and i32 %a, 787968
+  %or = or i32 %and, -921089
+  ret i32 %or
+}
+
+; Make sure we don't shrink or optimize an XOR's immediate operand if the
+; immediate is -1. Instruction selection turns (and ((xor $mask, -1), $v0)) into
+; a BIC.
+
+; CHECK-LABEL: _xor1:
+; CHECK: orr [[R0:w[0-9]+]], wzr, #0x38
+; CHECK: bic {{w[0-9]+}}, [[R0]], w0, lsl #3
+
+define i32 @xor1(i32 %a) {
+entry:
+  %shl = shl i32 %a, 3
+  %xor = and i32 %shl, 56
+  %and = xor i32 %xor, 56
+  ret i32 %and
+}