[TLI][X86][AArch64] Generalize isDesirableToCommuteWithShift() hook and enable for X86

[lebedev.ri]

Summary

Resolves PR38938.

[lebedev.ri]

lib/Target/X86/X86ISelDAGToDAG.cpp

2630–2633 ↗

(On Diff #166251)

Actually, this does not have any test coverage it seems.
So i'm not sure what extra tests this needs.

[lebedev.ri]

Comment # 4 on bug 38938 from Craig Topper

We need to get the shl created before the address mode for the inc is
selected
in Simon's test case. So we need to handle this in DAG combine or in
preprocessiseldag.

Ok, makes sense.
This turned out to be rather problematic :S

This is *very* prone to endless loops (we unfold the pattern,
and then SimplifyDemandedBits folds it back).
I suspect the change in DAGCombiner::SimplifyDemandedBits()
is too ugly, but i don't have any better ideas.
Anything else appears to deadlock.

These aren't all the test changes, but most of them.

[lebedev.ri]

Ok, so the fix for the *given* problem is so much simpler.

AArch64 already overrode isDesirableToCommuteWithShift() with this very pattern.
But we clearly want to do the same for x86.
So just generalize it.

FIXME: i'm not sure if isBitFieldExtractPattern() should be a virtual function?

[lebedev.ri]

Rebased ontop of more tests, NFC.
Will look into x86 dagcombine next..

[lebedev.ri]

In D52426#1250354, @RKSimon wrote:

@lebedev.ri Are you able to refactor this now that D52570 has landed?

Rebased now that D52570 has landed.

[lebedev.ri]

In D52426#1250878, @RKSimon wrote:

Can you remove the D52293 dependency and just move the current code to DAG so we can fix PR38938?

We don't want this hook for X86?

[lebedev.ri]

Does not appear to be worthwhile; effectively superseeded by D53126.

[RKSimon]

Abandon this?

[lebedev.ri]

In D52293#1262639, @RKSimon wrote:

In D52293#1262261, @lebedev.ri wrote:

Does not appear to be worthwhile; effectively superseeded by D53126.

Abandon this?

Oops, meant to and forgot :)

include/llvm/CodeGen/TargetLowering.h

public:
   ///   SDValue.Val == N   - N was replaced, is dead, and is already handled.
   ///   otherwise          - N should be replaced by the returned Operand.
   ///
   /// In addition, methods provided by DAGCombinerInfo may be used to perform
   /// more complex transformations.
   ///
   virtual SDValue PerformDAGCombine(SDNode *N, DAGCombinerInfo &DCI) const;
 
+  // Returns true if the target has a pattern
+  // to extract N bits starting with M'th bit.
+  // On x86 that could be BMI1's BEXTR, on AArch64 it is UBFX.
+  virtual bool haveBitFieldExtractPattern(EVT VT) const { return false; }
+
+  // True if N is  (X >> C1) & C2  where C2 is a mask with all-ones in low bits.
+  virtual bool isBitFieldExtractPattern(const SDNode *N) const {
+    if (!haveBitFieldExtractPattern(N->getValueType(0)))
+      return false;
+    if (N->getOpcode() != ISD::AND)
+      return false;
+    const auto *AndCstMask = dyn_cast<ConstantSDNode>(N->getOperand(1));
+    if (!AndCstMask)
+      return false;
+    SDValue N0 = N->getOperand(0);
+    if (N0.getOpcode() != ISD::SRL || !isa<ConstantSDNode>(N0.getOperand(1)))
+      return false;
+    // Ok, if the mask is all-ones in low bits, then it is 'bit field extract'.
+    return AndCstMask->getAPIntValue().isMask();
+  }
+
   /// Return true if it is profitable to move this shift by a constant amount
   /// though its operand, adjusting any immediate operands as necessary to
   /// preserve semantics. This transformation may not be desirable if it
-  /// disrupts a particularly auspicious target-specific tree (e.g. bitfield
-  /// extraction in AArch64). By default, it returns true.
+  /// disrupts a particularly suspicious target-specific tree (e.g. bitfield
+  /// extraction in x86/AArch64).
   ///
   /// @param N the shift node
   /// @param Level the current DAGCombine legalization level.
   virtual bool isDesirableToCommuteWithShift(const SDNode *N,
                                              CombineLevel Level) const {
-    return true;
+    // Desirable if it will not disturb the 'bit field extract' pattern.
+    return !isBitFieldExtractPattern(N->getOperand(0).getNode());
   }
 
   // Return true if it is profitable to combine a BUILD_VECTOR with a stride-pattern
   // to a shuffle and a truncate.
   // Example of such a combine:
   // v4i32 build_vector((extract_elt V, 1),
   //                    (extract_elt V, 3),
   //                    (extract_elt V, 5),

lib/Target/AArch64/AArch64ISelLowering.h

public:
 
   /// Return true if an FMA operation is faster than a pair of fmul and fadd
   /// instructions. fmuladd intrinsics will be expanded to FMAs when this method
   /// returns true, otherwise fmuladd is expanded to fmul + fadd.
   bool isFMAFasterThanFMulAndFAdd(EVT VT) const override;
 
   const MCPhysReg *getScratchRegisters(CallingConv::ID CC) const override;
 
-  /// Returns false if N is a bit extraction pattern of (X >> C) & Mask.
-  bool isDesirableToCommuteWithShift(const SDNode *N,
-                                     CombineLevel Level) const override;
+  bool haveBitFieldExtractPattern(EVT VT) const override {
+    // UBFX is valid for i32 and i64 types.
+    return VT == MVT::i32 || VT == MVT::i64;
+  }
 
   /// Returns true if it is beneficial to convert a load of a constant
   /// to just the constant itself.
   bool shouldConvertConstantLoadToIntImm(const APInt &Imm,
                                          Type *Ty) const override;
 
   /// Return true if EXTRACT_SUBVECTOR is cheap for this result type
   /// with this index.

lib/Target/AArch64/AArch64ISelLowering.cpp

AArch64TargetLowering::getScratchRegisters(CallingConv::ID) const {
   // site. Hence we include LR in the scratch registers, which are in turn added
   // as implicit-defs for stackmaps and patchpoints.
   static const MCPhysReg ScratchRegs[] = {
     AArch64::X16, AArch64::X17, AArch64::LR, 0
   };
   return ScratchRegs;
 }
 
-bool
-AArch64TargetLowering::isDesirableToCommuteWithShift(const SDNode *N,
-                                                     CombineLevel Level) const {
-  N = N->getOperand(0).getNode();
-  EVT VT = N->getValueType(0);
-    // If N is unsigned bit extraction: ((x >> C) & mask), then do not combine
-    // it with shift to let it be lowered to UBFX.
-  if (N->getOpcode() == ISD::AND && (VT == MVT::i32 || VT == MVT::i64) &&
-      isa<ConstantSDNode>(N->getOperand(1))) {
-    uint64_t TruncMask = N->getConstantOperandVal(1);
-    if (isMask_64(TruncMask) &&
-      N->getOperand(0).getOpcode() == ISD::SRL &&
-      isa<ConstantSDNode>(N->getOperand(0)->getOperand(1)))
-      return false;
-  }
-  return true;
-}
-
 bool AArch64TargetLowering::shouldConvertConstantLoadToIntImm(const APInt &Imm,
                                                               Type *Ty) const {
   assert(Ty->isIntegerTy());
 
   unsigned BitSize = Ty->getPrimitiveSizeInBits();
   if (BitSize == 0)
     return false;
 

lib/Target/X86/X86ISelLowering.h

public:
 
     /// Return true if the specified scalar FP type is computed in an SSE
     /// register, not on the X87 floating point stack.
     bool isScalarFPTypeInSSEReg(EVT VT) const {
       return (VT == MVT::f64 && X86ScalarSSEf64) || // f64 is when SSE2
              (VT == MVT::f32 && X86ScalarSSEf32);   // f32 is when SSE1
     }
 
+    bool haveBitFieldExtractPattern(EVT VT) const override;
+
     /// Returns true if it is beneficial to convert a load of a constant
     /// to just the constant itself.
     bool shouldConvertConstantLoadToIntImm(const APInt &Imm,
                                            Type *Ty) const override;
 
     bool convertSelectOfConstantsToMath(EVT VT) const override;
 
     bool decomposeMulByConstant(EVT VT, SDValue C) const override;

lib/Target/X86/X86ISelLowering.cpp

bool X86TargetLowering::shouldReduceLoadWidth(SDNode *Load,
   // relocation target a movq or addq instruction: don't let the load shrink.
   SDValue BasePtr = cast<LoadSDNode>(Load)->getBasePtr();
   if (BasePtr.getOpcode() == X86ISD::WrapperRIP)
     if (const auto *GA = dyn_cast<GlobalAddressSDNode>(BasePtr.getOperand(0)))
       return GA->getTargetFlags() != X86II::MO_GOTTPOFF;
   return true;
 }
 
+bool X86TargetLowering::haveBitFieldExtractPattern(EVT VT) const {
+  // BEXTR is valid for i32 and i64 types.
+  return Subtarget.hasBMI() && (VT == MVT::i32 || VT == MVT::i64);
+}
+
 /// Returns true if it is beneficial to convert a load of a constant
 /// to just the constant itself.
 bool X86TargetLowering::shouldConvertConstantLoadToIntImm(const APInt &Imm,
                                                           Type *Ty) const {
   assert(Ty->isIntegerTy());
 
   unsigned BitSize = Ty->getPrimitiveSizeInBits();
   if (BitSize == 0 || BitSize > 64)

test/CodeGen/X86/extract-bits.ll

; X64-BMI1BMI2-NEXT:    retq
   ret i64 %masked
 }
 
 ; ---------------------------------------------------------------------------- ;
 ; Constant
 ; ---------------------------------------------------------------------------- ;
 
 define void @pr38938(i32* %a0, i64* %a1) {
-; X86-LABEL: pr38938:
-; X86:       # %bb.0:
-; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; X86-NEXT:    movl {{[0-9]+}}(%esp), %ecx
-; X86-NEXT:    movl (%ecx), %ecx
-; X86-NEXT:    shrl $19, %ecx
-; X86-NEXT:    andl $4092, %ecx # imm = 0xFFC
-; X86-NEXT:    incl (%eax,%ecx)
-; X86-NEXT:    retl
-;
-; X64-LABEL: pr38938:
-; X64:       # %bb.0:
-; X64-NEXT:    movq (%rsi), %rax
-; X64-NEXT:    shrq $19, %rax
-; X64-NEXT:    andl $4092, %eax # imm = 0xFFC
-; X64-NEXT:    incl (%rdi,%rax)
-; X64-NEXT:    retq
+; X86-NOBMI-LABEL: pr38938:
+; X86-NOBMI:       # %bb.0:
+; X86-NOBMI-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; X86-NOBMI-NEXT:    movl {{[0-9]+}}(%esp), %ecx
+; X86-NOBMI-NEXT:    movl (%ecx), %ecx
+; X86-NOBMI-NEXT:    shrl $19, %ecx
+; X86-NOBMI-NEXT:    andl $4092, %ecx # imm = 0xFFC
+; X86-NOBMI-NEXT:    incl (%eax,%ecx)
+; X86-NOBMI-NEXT:    retl
+;
+; X86-BMI1NOTBM-LABEL: pr38938:
+; X86-BMI1NOTBM:       # %bb.0:
+; X86-BMI1NOTBM-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; X86-BMI1NOTBM-NEXT:    movl {{[0-9]+}}(%esp), %ecx
+; X86-BMI1NOTBM-NEXT:    movl $2581, %edx # imm = 0xA15
+; X86-BMI1NOTBM-NEXT:    bextrl %edx, (%ecx), %ecx
+; X86-BMI1NOTBM-NEXT:    incl (%eax,%ecx,4)
+; X86-BMI1NOTBM-NEXT:    retl
+;
+; X86-BMI1TBM-LABEL: pr38938:
+; X86-BMI1TBM:       # %bb.0:
+; X86-BMI1TBM-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; X86-BMI1TBM-NEXT:    movl {{[0-9]+}}(%esp), %ecx
+; X86-BMI1TBM-NEXT:    bextrl $2581, (%ecx), %ecx # imm = 0xA15
+; X86-BMI1TBM-NEXT:    incl (%eax,%ecx,4)
+; X86-BMI1TBM-NEXT:    retl
+;
+; X86-BMI1NOTBMBMI2-LABEL: pr38938:
+; X86-BMI1NOTBMBMI2:       # %bb.0:
+; X86-BMI1NOTBMBMI2-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; X86-BMI1NOTBMBMI2-NEXT:    movl {{[0-9]+}}(%esp), %ecx
+; X86-BMI1NOTBMBMI2-NEXT:    movl $2581, %edx # imm = 0xA15
+; X86-BMI1NOTBMBMI2-NEXT:    bextrl %edx, (%ecx), %ecx
+; X86-BMI1NOTBMBMI2-NEXT:    incl (%eax,%ecx,4)
+; X86-BMI1NOTBMBMI2-NEXT:    retl
+;
+; X64-NOBMI-LABEL: pr38938:
+; X64-NOBMI:       # %bb.0:
+; X64-NOBMI-NEXT:    movq (%rsi), %rax
+; X64-NOBMI-NEXT:    shrq $19, %rax
+; X64-NOBMI-NEXT:    andl $4092, %eax # imm = 0xFFC
+; X64-NOBMI-NEXT:    incl (%rdi,%rax)
+; X64-NOBMI-NEXT:    retq
+;
+; X64-BMI1NOTBM-LABEL: pr38938:
+; X64-BMI1NOTBM:       # %bb.0:
+; X64-BMI1NOTBM-NEXT:    movl $2581, %eax # imm = 0xA15
+; X64-BMI1NOTBM-NEXT:    bextrl %eax, (%rsi), %eax
+; X64-BMI1NOTBM-NEXT:    incl (%rdi,%rax,4)
+; X64-BMI1NOTBM-NEXT:    retq
+;
+; X64-BMI1TBM-LABEL: pr38938:
+; X64-BMI1TBM:       # %bb.0:
+; X64-BMI1TBM-NEXT:    bextrl $2581, (%rsi), %eax # imm = 0xA15
+; X64-BMI1TBM-NEXT:    incl (%rdi,%rax,4)
+; X64-BMI1TBM-NEXT:    retq
+;
+; X64-BMI1NOTBMBMI2-LABEL: pr38938:
+; X64-BMI1NOTBMBMI2:       # %bb.0:
+; X64-BMI1NOTBMBMI2-NEXT:    movl $2581, %eax # imm = 0xA15
+; X64-BMI1NOTBMBMI2-NEXT:    bextrl %eax, (%rsi), %eax
+; X64-BMI1NOTBMBMI2-NEXT:    incl (%rdi,%rax,4)
+; X64-BMI1NOTBMBMI2-NEXT:    retq
   %tmp = load i64, i64* %a1, align 8
   %tmp1 = lshr i64 %tmp, 21
   %tmp2 = and i64 %tmp1, 1023
   %tmp3 = getelementptr inbounds i32, i32* %a0, i64 %tmp2
   %tmp4 = load i32, i32* %tmp3, align 4
   %tmp5 = add nsw i32 %tmp4, 1
   store i32 %tmp5, i32* %tmp3, align 4
   ret void
 }