[GlobalISel] Expand combine for (x & mask) -> x when (x & mask) == x

[mbrkusanin]

Summary

We can use KnownBitsAnalysis to cover cases when mask is not trivial. It can
also help with cases when mask is not constant but can still be folded into
one. Since 'and' is commutative we should treat both operands as possible
replacements.

[arsenm]

llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

426

Unnecessary whitespace change

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2860–2861

Constants are canonicalized to the RHS, so you don't need KnownBits and can just directly get the constant.

2873

MaskedValueIsZero helper like the DAG?

[foad]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2860–2861

Or you can implement it generally and symmetrically, so that it works for non-constants as well as constants:

X & Y -> X // if (XBits.Zero | YBits.One).isAllOnesValue()
X & Y -> Y // if (XBits.One | YBits.Zero).isAllOnesValue()

[mbrkusanin]

• Refactored match function
• Updated tests

[mbrkusanin]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2860–2861

@arsenm With KnownBits we can also catch cases like in test_and_non_const. Or the ones in combine-shift-of-shifted-logic.ll generated by D90223

[foad]

LGTM modulo minor comments inline.

llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

433

Maybe just matchRedundantAnd?

llvm/include/llvm/Target/GlobalISel/Combine.td

386–388

I think we should define a standard Register matchdata that we can use here and elsewhere. Doesn't have to be part of this patch though.

[paquette]

llvm/include/llvm/Target/GlobalISel/Combine.td

386–388

+1

[mbrkusanin]

• Renamed
• Updated comments

llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

public:
   bool applyBuildInstructionSteps(MachineInstr &MI,
                                   InstructionStepsMatchInfo &MatchInfo);
 
   /// Match ashr (shl x, C), C -> sext_inreg (C)
   bool matchAshrShlToSextInreg(MachineInstr &MI,
                                std::tuple<Register, int64_t> &MatchInfo);
   bool applyAshShlToSextInreg(MachineInstr &MI,
                               std::tuple<Register, int64_t> &MatchInfo);
+
   /// \return true if \p MI is a G_AND instruction whose RHS is a mask where
   /// LHS & mask == LHS. (E.g., an all-ones value.)
   ///
   /// \param [in] MI - The G_AND instruction.
   /// \param [out] Replacement - A register the G_AND should be replaced with on
   /// success.
-  bool matchAndWithTrivialMask(MachineInstr &MI, Register &Replacement);
+  bool matchAndWithRedundantMask(MachineInstr &MI, Register &Replacement);
 
   /// \return true if \p MI is a G_SEXT_INREG that can be erased.
   bool matchRedundantSExtInReg(MachineInstr &MI);
 
   /// Combine inverting a result of a compare into the opposite cond code.
   bool matchNotCmp(MachineInstr &MI, SmallVectorImpl<Register> &RegsToNegate);
   bool applyNotCmp(MachineInstr &MI, SmallVectorImpl<Register> &RegsToNegate);
 

llvm/include/llvm/Target/GlobalISel/Combine.td

 // Fold ashr (shl x, C), C -> sext_inreg (C)
 def shl_ashr_to_sext_inreg_matchinfo : GIDefMatchData<"std::tuple<Register, int64_t>">;
 def shl_ashr_to_sext_inreg : GICombineRule<
   (defs root:$root, shl_ashr_to_sext_inreg_matchinfo:$info),
   (match (wip_match_opcode G_ASHR): $root,
     [{ return Helper.matchAshrShlToSextInreg(*${root}, ${info}); }]),
   (apply [{ return Helper.applyAshShlToSextInreg(*${root}, ${info});}])
 >;
+
 // Fold (x & mask) -> x when (x & mask) is known to equal x.
-def and_trivial_mask_matchinfo : GIDefMatchData<"Register">;
-def and_trivial_mask: GICombineRule <
-  (defs root:$root, and_trivial_mask_matchinfo:$matchinfo),
+def and_redundant_mask_matchinfo : GIDefMatchData<"Register">;
+def and_redundant_mask: GICombineRule <
+  (defs root:$root, and_redundant_mask_matchinfo:$matchinfo),
   (match (wip_match_opcode G_AND):$root,
-         [{ return Helper.matchAndWithTrivialMask(*${root}, ${matchinfo}); }]),
+         [{ return Helper.matchAndWithRedundantMask(*${root}, ${matchinfo}); }]),
   (apply [{ return Helper.replaceSingleDefInstWithReg(*${root}, ${matchinfo}); }])
 >;
 
 // If the input is already sign extended, just drop the extension.
 // sext_inreg x, K ->
 //   if computeNumSignBits(x) >= (x.getScalarSizeInBits() - K + 1)
 def redundant_sext_inreg: GICombineRule <
   (defs root:$root),
def identity_combines : GICombineGroup<[select_same_val, right_identity_zero,
                                         binop_same_val, binop_left_to_zero,
                                         binop_right_to_zero, p2i_to_i2p,
                                         i2p_to_p2i, anyext_trunc_fold,
                                         fneg_fneg_fold, right_identity_one]>;
 
 def const_combines : GICombineGroup<[constant_fp_op, const_ptradd_to_i2p]>;
 
 def known_bits_simplifications : GICombineGroup<[
-  and_trivial_mask, redundant_sext_inreg]>;
+  and_redundant_mask, redundant_sext_inreg]>;
 
 def width_reduction_combines : GICombineGroup<[reduce_shl_of_extend]>;
 
 def select_combines : GICombineGroup<[select_undef_cmp, select_constant_cmp]>;
 
 def trivial_combines : GICombineGroup<[copy_prop, mul_to_shl, add_p2i_to_ptradd,
                                        mul_by_neg_one]>;
 

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

bool CombinerHelper::applyAshShlToSextInreg(
   std::tie(Src, ShiftAmt) = MatchInfo;
   unsigned Size = MRI.getType(Src).getScalarSizeInBits();
   Builder.setInstrAndDebugLoc(MI);
   Builder.buildSExtInReg(MI.getOperand(0).getReg(), Src, Size - ShiftAmt);
   MI.eraseFromParent();
   return true;
 }
 
-bool CombinerHelper::matchAndWithTrivialMask(MachineInstr &MI,
+bool CombinerHelper::matchAndWithRedundantMask(MachineInstr &MI,
                                                Register &Replacement) {
   // Given
   //
-  // %mask:_(sN) = G_CONSTANT iN 000...0111...1
+  // %mask:_(sN) = G_CONSTANT iN %c
   // %x:_(sN) = G_SOMETHING
   // %y:_(sN) = G_AND %x, %mask
   //
   // Eliminate the G_AND when it is known that x & mask == x.
   //
   // Patterns like this can appear as a result of legalization. E.g.
   //
   // %cmp:_(s32) = G_ICMP intpred(pred), %x(s32), %y
   // %one:_(s32) = G_CONSTANT i32 1
   // %and:_(s32) = G_AND %cmp, %one
   //
   // In this case, G_ICMP only produces a single bit, so x & 1 == x.
   assert(MI.getOpcode() == TargetOpcode::G_AND);
   if (!KB)
     return false;
 
-  // Replacement = %x, AndDst = %y. Check that we can replace AndDst with the
-  // LHS of the G_AND.
-  Replacement = MI.getOperand(1).getReg();
   Register AndDst = MI.getOperand(0).getReg();
   LLT DstTy = MRI.getType(AndDst);
 
   // FIXME: This should be removed once GISelKnownBits supports vectors.
   if (DstTy.isVector())
     return false;
-  if (!canReplaceReg(AndDst, Replacement, MRI))
-    return false;
 
-  // Check that we have a constant on the RHS of the G_AND, which is of the form
-  // 000...0111...1.
-  int64_t Cst;
-  if (!mi_match(MI.getOperand(2).getReg(), MRI, m_ICst(Cst)))
-    return false;
-  APInt Mask(DstTy.getSizeInBits(), Cst);
-  if (!Mask.isMask())
-    return false;
+  Register LHS = MI.getOperand(1).getReg();
+  Register RHS = MI.getOperand(2).getReg();
 
-  // Now, let's check that x & Mask == x. If this is true, then x & ~Mask == 0.
-  return KB->maskedValueIsZero(Replacement, ~Mask);
+  // While we need to know all bits of the Mask, for x we only need the bits
+  // that go against zeros in Mask.
+  KnownBits LHSBits = KB->getKnownBits(LHS);
+  KnownBits RHSBits = KB->getKnownBits(RHS);
+
+  // Check that x & Mask == x.
+  // x & 1 == x, always
+  // x & 0 == x, only if x is also 0
+  // Meaning Mask has no effect if all it's zeros match zeros in x.
+  //
+  // Check that we can replace AndDst with the LHS of the G_AND and that all
+  // zero bits in RHS are also zero in LHS.
+  if (RHSBits.isConstant() && canReplaceReg(AndDst, LHS, MRI) &&
+      RHSBits.Zero.isSubsetOf(LHSBits.Zero)) {
+    Replacement = LHS;
+    return true;
+  }
+
+  // And is commutative so we need to check both cases.
+  if (LHSBits.isConstant() && canReplaceReg(AndDst, RHS, MRI) &&
+      LHSBits.Zero.isSubsetOf(RHSBits.Zero)) {
+    Replacement = RHS;
+    return true;
+  }
+
+  return false;
 }
 
 bool CombinerHelper::matchRedundantSExtInReg(MachineInstr &MI) {
   // If the input is already sign extended, just drop the extension.
   Register Src = MI.getOperand(1).getReg();
   unsigned ExtBits = MI.getOperand(2).getImm();
   unsigned TypeSize = MRI.getType(Src).getScalarSizeInBits();
   return KB->computeNumSignBits(Src) >= (TypeSize - ExtBits + 1);

llvm/lib/Target/AArch64/AArch64Combine.td

 }
 
 // Post-legalization combines which are primarily optimizations.
 def AArch64PostLegalizerCombinerHelper
     : GICombinerHelper<"AArch64GenPostLegalizerCombinerHelper",
                        [copy_prop, erase_undef_store, combines_for_extload,
                         sext_trunc_sextload,
                         hoist_logic_op_with_same_opcode_hands,
-                        and_trivial_mask, xor_of_and_with_same_reg]> {
+                        and_redundant_mask, xor_of_and_with_same_reg]> {
   let DisableRuleOption = "aarch64postlegalizercombiner-disable-rule";
 }

llvm/test/CodeGen/AMDGPU/GlobalISel/combine-and-redundant.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=amdgcn -run-pass=amdgpu-prelegalizer-combiner -verify-machineinstrs %s -o - | FileCheck %s
+
+---
+name:            test_const_const_lhs
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: test_const_const_lhs
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 15
+    ; CHECK: $sgpr0 = COPY [[C]](s32)
+    ; CHECK: SI_RETURN_TO_EPILOG implicit $sgpr0
+    %0:_(s32) = G_CONSTANT i32 15
+    %1:_(s32) = G_CONSTANT i32 255
+    %2:_(s32) = G_AND %0(s32), %1(s32)
+    $sgpr0 = COPY %2(s32)
+    SI_RETURN_TO_EPILOG implicit $sgpr0
+...
+
+---
+name:            test_const_const_rhs
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: test_const_const_rhs
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 15
+    ; CHECK: $vgpr0 = COPY [[C]](s32)
+    ; CHECK: SI_RETURN_TO_EPILOG implicit $vgpr0
+    %0:_(s32) = G_CONSTANT i32 15
+    %1:_(s32) = G_CONSTANT i32 255
+    %2:_(s32) = G_AND %1(s32), %0(s32)
+    $vgpr0 = COPY %2(s32)
+    SI_RETURN_TO_EPILOG implicit $vgpr0
+...
+
+---
+name:            test_and_consts
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: test_and_consts
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 15
+    ; CHECK: $sgpr0 = COPY [[C]](s32)
+    ; CHECK: SI_RETURN_TO_EPILOG implicit $sgpr0
+    %0:_(s32) = G_CONSTANT i32 15
+    %1:_(s32) = G_CONSTANT i32 255
+    %2:_(s32) = G_AND %0(s32), %1(s32)
+    $sgpr0 = COPY %2(s32)
+    SI_RETURN_TO_EPILOG implicit $sgpr0
+...
+
+---
+name:            test_and_and
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+
+    ; CHECK-LABEL: name: test_and_and
+    ; CHECK: liveins: $vgpr0
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $vgpr0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 15
+    ; CHECK: [[AND:%[0-9]+]]:_(s32) = G_AND [[COPY]], [[C]]
+    ; CHECK: $vgpr0 = COPY [[AND]](s32)
+    ; CHECK: SI_RETURN_TO_EPILOG implicit $vgpr0
+    %0:_(s32) = COPY $vgpr0
+    %1:_(s32) = G_CONSTANT i32 15
+    %2:_(s32) = G_CONSTANT i32 255
+    %3:_(s32) = G_AND %0, %1(s32)
+    %4:_(s32) = G_AND %3, %2
+    $vgpr0 = COPY %4(s32)
+    SI_RETURN_TO_EPILOG implicit $vgpr0
+...
+
+---
+name:            test_shl_and
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $sgpr0
+
+    ; CHECK-LABEL: name: test_shl_and
+    ; CHECK: liveins: $sgpr0
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $sgpr0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 5
+    ; CHECK: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[COPY]], [[C]](s32)
+    ; CHECK: $sgpr0 = COPY [[SHL]](s32)
+    ; CHECK: SI_RETURN_TO_EPILOG implicit $sgpr0
+    %0:_(s32) = COPY $sgpr0
+    %1:_(s32) = G_CONSTANT i32 5
+    %2:_(s32) = G_CONSTANT i32 4294967264
+    %3:_(s32) = G_SHL %0, %1(s32)
+    %4:_(s32) = G_AND %3, %2
+    $sgpr0 = COPY %4(s32)
+    SI_RETURN_TO_EPILOG implicit $sgpr0
+...
+
+---
+name:            test_lshr_and
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+
+    ; CHECK-LABEL: name: test_lshr_and
+    ; CHECK: liveins: $vgpr0
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $vgpr0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 5
+    ; CHECK: [[LSHR:%[0-9]+]]:_(s32) = G_LSHR [[COPY]], [[C]](s32)
+    ; CHECK: $vgpr0 = COPY [[LSHR]](s32)
+    ; CHECK: SI_RETURN_TO_EPILOG implicit $vgpr0
+    %0:_(s32) = COPY $vgpr0
+    %1:_(s32) = G_CONSTANT i32 5
+    %2:_(s32) = G_CONSTANT i32 134217727
+    %3:_(s32) = G_LSHR %0, %1(s32)
+    %4:_(s32) = G_AND %3, %2
+    $vgpr0 = COPY %4(s32)
+    SI_RETURN_TO_EPILOG implicit $vgpr0
+...

llvm/test/CodeGen/AMDGPU/GlobalISel/combine-shift-of-shifted-logic.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -global-isel -march=amdgcn -verify-machineinstrs < %s | FileCheck %s
 
 define amdgpu_cs i32 @test_shl_and_1(i32 inreg %arg1) {
 ; CHECK-LABEL: test_shl_and_1:
 ; CHECK:       ; %bb.0: ; %.entry
 ; CHECK-NEXT:    s_lshl_b32 s0, s0, 4
-; CHECK-NEXT:    s_and_b32 s0, s0, -16
 ; CHECK-NEXT:    ; return to shader part epilog
 .entry:
   %z1 = shl i32 %arg1, 2
   %z2 = and i32 %z1, 1073741820
   %z3 = shl i32 %z2, 2
   ret i32 %z3
 }
 
 define amdgpu_cs i32 @test_shl_and_2(i32 inreg %arg1) {
 ; CHECK-LABEL: test_shl_and_2:
 ; CHECK:       ; %bb.0: ; %.entry
 ; CHECK-NEXT:    s_lshl_b32 s0, s0, 8
-; CHECK-NEXT:    s_and_b32 s0, s0, 0xffffff00
 ; CHECK-NEXT:    ; return to shader part epilog
 .entry:
   %z1 = shl i32 %arg1, 5
   %z2 = and i32 %z1, 536870880
   %z3 = shl i32 %z2, 3
   ret i32 %z3
 }
 
.entry:
   %z3 = shl i32 %z2, 2
   ret i32 %z3
 }
 
 define amdgpu_cs i32 @test_lshr_and_1(i32 inreg %arg1) {
 ; CHECK-LABEL: test_lshr_and_1:
 ; CHECK:       ; %bb.0: ; %.entry
 ; CHECK-NEXT:    s_lshr_b32 s0, s0, 4
-; CHECK-NEXT:    s_and_b32 s0, s0, 0xfffffff
 ; CHECK-NEXT:    ; return to shader part epilog
 .entry:
   %z1 = lshr i32 %arg1, 2
   %z2 = and i32 %z1, 1073741820
   %z3 = lshr i32 %z2, 2
   ret i32 %z3
 }
 
.entry:
   %z3 = lshr i32 %z2, 3
   ret i32 %z3
 }
 
 define amdgpu_cs i32 @test_lshr_and_3(i32 inreg %arg1) {
 ; CHECK-LABEL: test_lshr_and_3:
 ; CHECK:       ; %bb.0: ; %.entry
 ; CHECK-NEXT:    s_lshr_b32 s0, s0, 5
-; CHECK-NEXT:    s_and_b32 s0, s0, 0x7ffffff
 ; CHECK-NEXT:    ; return to shader part epilog
 .entry:
   %z1 = lshr i32 %arg1, 3
   %z2 = and i32 %z1, 536870908
   %z3 = lshr i32 %z2, 2
   ret i32 %z3
 }