[GISel] Fold bitreverse(shl/srl(bitreverse(x),y)) -> srl/shl(x,y)

RKSimon · RKSimon · commit bf7343aaffeb · 2024-05-07T18:05:04.000+01:00
Sibling patch to llvm#89897
diff --git a/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h b/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h
@@ -320,6 +320,9 @@ class CombinerHelper {
   void applyCombineShlOfExtend(MachineInstr &MI,
                                const RegisterImmPair &MatchData);
 
+  /// Fold bitreverse(shift (bitreverse x), y)) -> (shift x, y)
+  bool matchBitreverseShift(MachineInstr &MI, BuildFnTy &MatchInfo);
+
   /// Fold away a merge of an unmerge of the corresponding values.
   bool matchCombineMergeUnmerge(MachineInstr &MI, Register &MatchInfo);
 
diff --git a/llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h b/llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h
@@ -598,6 +598,12 @@ m_GBitcast(const SrcTy &Src) {
   return UnaryOp_match<SrcTy, TargetOpcode::G_BITCAST>(Src);
 }
 
+template <typename SrcTy>
+inline UnaryOp_match<SrcTy, TargetOpcode::G_BITREVERSE>
+m_GBitreverse(const SrcTy &Src) {
+  return UnaryOp_match<SrcTy, TargetOpcode::G_BITREVERSE>(Src);
+}
+
 template <typename SrcTy>
 inline UnaryOp_match<SrcTy, TargetOpcode::G_PTRTOINT>
 m_GPtrToInt(const SrcTy &Src) {
diff --git a/llvm/include/llvm/Target/GlobalISel/Combine.td b/llvm/include/llvm/Target/GlobalISel/Combine.td
@@ -323,6 +323,13 @@ def reduce_shl_of_extend : GICombineRule<
          [{ return Helper.matchCombineShlOfExtend(*${mi}, ${matchinfo}); }]),
   (apply [{ Helper.applyCombineShlOfExtend(*${mi}, ${matchinfo}); }])>;
 
+// Combine bitreverse(shift (bitreverse x), y)) -> (shift x, y)
+def bitreverse_shift : GICombineRule<
+  (defs root:$d, build_fn_matchinfo:$matchinfo),
+  (match (wip_match_opcode G_BITREVERSE):$d,
+         [{ return Helper.matchBitreverseShift(*${d}, ${matchinfo}); }]),
+  (apply [{ Helper.applyBuildFn(*${d}, ${matchinfo}); }])>;
+
 // Combine (shl (add x, c1), c2) -> (add (shl x, c2), c1 << c2)
 // Combine (shl (or x, c1), c2) -> (or (shl x, c2), c1 << c2)
 def commute_shift : GICombineRule<
@@ -1658,7 +1665,7 @@ def all_combines : GICombineGroup<[trivial_combines, vector_ops_combines,
     unmerge_zext_to_zext, merge_unmerge, trunc_ext_fold, trunc_shift,
     const_combines, xor_of_and_with_same_reg, ptr_add_with_zero,
     shift_immed_chain, shift_of_shifted_logic_chain, load_or_combine,
-    div_rem_to_divrem, funnel_shift_combines, commute_shift,
+    div_rem_to_divrem, funnel_shift_combines, bitreverse_shift, commute_shift,
     form_bitfield_extract, constant_fold_binops, constant_fold_fma,
     constant_fold_cast_op, fabs_fneg_fold,
     intdiv_combines, mulh_combines, redundant_neg_operands,
diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
@@ -2016,6 +2016,32 @@ void CombinerHelper::applyCombineShlOfExtend(MachineInstr &MI,
   MI.eraseFromParent();
 }
 
+bool CombinerHelper::matchBitreverseShift(MachineInstr &MI,
+                                          BuildFnTy &MatchInfo) {
+  assert(MI.getOpcode() == TargetOpcode::G_BITREVERSE && "Expected BITREVERSE");
+  Register Dst = MI.getOperand(0).getReg();
+  Register Src = MI.getOperand(1).getReg();
+  Register Val, Amt;
+
+  // fold (bitreverse (shl (bitreverse x), y)) -> (lshr x, y)
+  if (mi_match(Src, MRI, m_GShl(m_GBitreverse(m_Reg(Val)), m_Reg(Amt))) &&
+      isLegalOrBeforeLegalizer(
+          {TargetOpcode::G_LSHR, {MRI.getType(Val), MRI.getType(Amt)}})) {
+    MatchInfo = [=](MachineIRBuilder &B) { B.buildLShr(Dst, Val, Amt); };
+    return true;
+  }
+
+  // fold (bitreverse (lshr (bitreverse x), y)) -> (shl x, y)
+  if (mi_match(Src, MRI, m_GLShr(m_GBitreverse(m_Reg(Val)), m_Reg(Amt))) &&
+      isLegalOrBeforeLegalizer(
+          {TargetOpcode::G_SHL, {MRI.getType(Val), MRI.getType(Amt)}})) {
+    MatchInfo = [=](MachineIRBuilder &B) { B.buildShl(Dst, Val, Amt); };
+    return true;
+  }
+
+  return false;
+}
+
 bool CombinerHelper::matchCombineMergeUnmerge(MachineInstr &MI,
                                               Register &MatchInfo) {
   GMerge &Merge = cast<GMerge>(MI);
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/combine-bitreverse-shift.ll b/llvm/test/CodeGen/AArch64/GlobalISel/combine-bitreverse-shift.ll
@@ -1,6 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc < %s -mtriple=aarch64-unknown-unknown | FileCheck %s --check-prefixes=SDAG
-; RUN: llc < %s -mtriple=aarch64-unknown-unknown -global-isel | FileCheck %s --check-prefixes=GISEL
+; RUN: llc < %s -mtriple=aarch64-unknown-unknown | FileCheck %s
+; RUN: llc < %s -mtriple=aarch64-unknown-unknown -global-isel | FileCheck %s
 
 ; These tests can be optimised
 ;       fold (bitreverse(srl (bitreverse c), x)) -> (shl c, x)
@@ -12,151 +12,87 @@ declare i32 @llvm.bitreverse.i32(i32)
 declare i64 @llvm.bitreverse.i64(i64)
 
 define i8 @test_bitreverse_srli_bitreverse_i8(i8 %a) nounwind {
-; SDAG-LABEL: test_bitreverse_srli_bitreverse_i8:
-; SDAG:       // %bb.0:
-; SDAG-NEXT:    lsl w0, w0, #3
-; SDAG-NEXT:    ret
-;
-; GISEL-LABEL: test_bitreverse_srli_bitreverse_i8:
-; GISEL:       // %bb.0:
-; GISEL-NEXT:    rbit w8, w0
-; GISEL-NEXT:    lsr w8, w8, #24
-; GISEL-NEXT:    lsr w8, w8, #3
-; GISEL-NEXT:    rbit w8, w8
-; GISEL-NEXT:    lsr w0, w8, #24
-; GISEL-NEXT:    ret
+; CHECK-LABEL: test_bitreverse_srli_bitreverse_i8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    lsl w0, w0, #3
+; CHECK-NEXT:    ret
   %1 = call i8 @llvm.bitreverse.i8(i8 %a)
   %2 = lshr i8 %1, 3
   %3 = call i8 @llvm.bitreverse.i8(i8 %2)
   ret i8 %3
 }
 
 define i16 @test_bitreverse_srli_bitreverse_i16(i16 %a) nounwind {
-; SDAG-LABEL: test_bitreverse_srli_bitreverse_i16:
-; SDAG:       // %bb.0:
-; SDAG-NEXT:    lsl w0, w0, #7
-; SDAG-NEXT:    ret
-;
-; GISEL-LABEL: test_bitreverse_srli_bitreverse_i16:
-; GISEL:       // %bb.0:
-; GISEL-NEXT:    rbit w8, w0
-; GISEL-NEXT:    lsr w8, w8, #16
-; GISEL-NEXT:    lsr w8, w8, #7
-; GISEL-NEXT:    rbit w8, w8
-; GISEL-NEXT:    lsr w0, w8, #16
-; GISEL-NEXT:    ret
+; CHECK-LABEL: test_bitreverse_srli_bitreverse_i16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    lsl w0, w0, #7
+; CHECK-NEXT:    ret
   %1 = call i16 @llvm.bitreverse.i16(i16 %a)
   %2 = lshr i16 %1, 7
   %3 = call i16 @llvm.bitreverse.i16(i16 %2)
   ret i16 %3
 }
 
 define i32 @test_bitreverse_srli_bitreverse_i32(i32 %a) nounwind {
-; SDAG-LABEL: test_bitreverse_srli_bitreverse_i32:
-; SDAG:       // %bb.0:
-; SDAG-NEXT:    lsl w0, w0, #15
-; SDAG-NEXT:    ret
-;
-; GISEL-LABEL: test_bitreverse_srli_bitreverse_i32:
-; GISEL:       // %bb.0:
-; GISEL-NEXT:    rbit w8, w0
-; GISEL-NEXT:    lsr w8, w8, #15
-; GISEL-NEXT:    rbit w0, w8
-; GISEL-NEXT:    ret
+; CHECK-LABEL: test_bitreverse_srli_bitreverse_i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    lsl w0, w0, #15
+; CHECK-NEXT:    ret
   %1 = call i32 @llvm.bitreverse.i32(i32 %a)
   %2 = lshr i32 %1, 15
   %3 = call i32 @llvm.bitreverse.i32(i32 %2)
   ret i32 %3
 }
 
 define i64 @test_bitreverse_srli_bitreverse_i64(i64 %a) nounwind {
-; SDAG-LABEL: test_bitreverse_srli_bitreverse_i64:
-; SDAG:       // %bb.0:
-; SDAG-NEXT:    lsl x0, x0, #33
-; SDAG-NEXT:    ret
-;
-; GISEL-LABEL: test_bitreverse_srli_bitreverse_i64:
-; GISEL:       // %bb.0:
-; GISEL-NEXT:    rbit x8, x0
-; GISEL-NEXT:    lsr x8, x8, #33
-; GISEL-NEXT:    rbit x0, x8
-; GISEL-NEXT:    ret
+; CHECK-LABEL: test_bitreverse_srli_bitreverse_i64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    lsl x0, x0, #33
+; CHECK-NEXT:    ret
   %1 = call i64 @llvm.bitreverse.i64(i64 %a)
   %2 = lshr i64 %1, 33
   %3 = call i64 @llvm.bitreverse.i64(i64 %2)
   ret i64 %3
 }
 
 define i8 @test_bitreverse_shli_bitreverse_i8(i8 %a) nounwind {
-; SDAG-LABEL: test_bitreverse_shli_bitreverse_i8:
-; SDAG:       // %bb.0:
-; SDAG-NEXT:    ubfx w0, w0, #3, #5
-; SDAG-NEXT:    ret
-;
-; GISEL-LABEL: test_bitreverse_shli_bitreverse_i8:
-; GISEL:       // %bb.0:
-; GISEL-NEXT:    rbit w8, w0
-; GISEL-NEXT:    lsr w8, w8, #24
-; GISEL-NEXT:    lsl w8, w8, #3
-; GISEL-NEXT:    rbit w8, w8
-; GISEL-NEXT:    lsr w0, w8, #24
-; GISEL-NEXT:    ret
+; CHECK-LABEL: test_bitreverse_shli_bitreverse_i8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ubfx w0, w0, #3, #5
+; CHECK-NEXT:    ret
   %1 = call i8 @llvm.bitreverse.i8(i8 %a)
   %2 = shl i8 %1, 3
   %3 = call i8 @llvm.bitreverse.i8(i8 %2)
   ret i8 %3
 }
 
 define i16 @test_bitreverse_shli_bitreverse_i16(i16 %a) nounwind {
-; SDAG-LABEL: test_bitreverse_shli_bitreverse_i16:
-; SDAG:       // %bb.0:
-; SDAG-NEXT:    ubfx w0, w0, #7, #9
-; SDAG-NEXT:    ret
-;
-; GISEL-LABEL: test_bitreverse_shli_bitreverse_i16:
-; GISEL:       // %bb.0:
-; GISEL-NEXT:    rbit w8, w0
-; GISEL-NEXT:    lsr w8, w8, #16
-; GISEL-NEXT:    lsl w8, w8, #7
-; GISEL-NEXT:    rbit w8, w8
-; GISEL-NEXT:    lsr w0, w8, #16
-; GISEL-NEXT:    ret
+; CHECK-LABEL: test_bitreverse_shli_bitreverse_i16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ubfx w0, w0, #7, #9
+; CHECK-NEXT:    ret
   %1 = call i16 @llvm.bitreverse.i16(i16 %a)
   %2 = shl i16 %1, 7
   %3 = call i16 @llvm.bitreverse.i16(i16 %2)
   ret i16 %3
 }
 
 define i32 @test_bitreverse_shli_bitreverse_i32(i32 %a) nounwind {
-; SDAG-LABEL: test_bitreverse_shli_bitreverse_i32:
-; SDAG:       // %bb.0:
-; SDAG-NEXT:    lsr w0, w0, #15
-; SDAG-NEXT:    ret
-;
-; GISEL-LABEL: test_bitreverse_shli_bitreverse_i32:
-; GISEL:       // %bb.0:
-; GISEL-NEXT:    rbit w8, w0
-; GISEL-NEXT:    lsl w8, w8, #15
-; GISEL-NEXT:    rbit w0, w8
-; GISEL-NEXT:    ret
+; CHECK-LABEL: test_bitreverse_shli_bitreverse_i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    lsr w0, w0, #15
+; CHECK-NEXT:    ret
   %1 = call i32 @llvm.bitreverse.i32(i32 %a)
   %2 = shl i32 %1, 15
   %3 = call i32 @llvm.bitreverse.i32(i32 %2)
   ret i32 %3
 }
 
 define i64 @test_bitreverse_shli_bitreverse_i64(i64 %a) nounwind {
-; SDAG-LABEL: test_bitreverse_shli_bitreverse_i64:
-; SDAG:       // %bb.0:
-; SDAG-NEXT:    lsr x0, x0, #33
-; SDAG-NEXT:    ret
-;
-; GISEL-LABEL: test_bitreverse_shli_bitreverse_i64:
-; GISEL:       // %bb.0:
-; GISEL-NEXT:    rbit x8, x0
-; GISEL-NEXT:    lsl x8, x8, #33
-; GISEL-NEXT:    rbit x0, x8
-; GISEL-NEXT:    ret
+; CHECK-LABEL: test_bitreverse_shli_bitreverse_i64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    lsr x0, x0, #33
+; CHECK-NEXT:    ret
   %1 = call i64 @llvm.bitreverse.i64(i64 %a)
   %2 = shl i64 %1, 33
   %3 = call i64 @llvm.bitreverse.i64(i64 %2)
