[GISel] Add new GISel combiners for G_SELECT

https://reviews.llvm.org/D83833

Patch adds two new GICombinerRules for G_SELECT. The rules include:
combining selects with undef comparisons into their first selectee value,
and to combine away selects with constant comparisons. Patch additionally
adds a new combiner test for the AArch64 target to test these new G_SELECT
combiner rules and the existing select_same_val combiner rule.

Patch by  mkitzan

include/llvm/CodeGen/GlobalISel/CombinerHelper.h

@@ -272,6 +272,13 @@ public:
   /// Return true if a G_STORE instruction \p MI is storing an undef value.
   bool matchUndefStore(MachineInstr &MI);
 
+  /// Return true if a G_SELECT instruction \p MI has an undef comparison.
+  bool matchUndefSelectCmp(MachineInstr &MI);
+
+  /// Return true if a G_SELECT instruction \p MI has a constant comparison. If
+  /// true, \p OpIdx will store the operand index of the known selected value.
+  bool matchConstantSelectCmp(MachineInstr &MI, unsigned &OpIdx);
+
   /// Replace an instruction with a G_FCONSTANT with value \p C.
   bool replaceInstWithFConstant(MachineInstr &MI, double C);
 

include/llvm/Target/GlobalISel/Combine.td

@@ -225,6 +225,24 @@ def select_same_val: GICombineRule<
   (apply [{ return Helper.replaceSingleDefInstWithOperand(*${root}, 2); }])
 >;
 
+// Fold (undef ? x : y) -> y
+def select_undef_cmp: GICombineRule<
+  (defs root:$root),
+  (match (wip_match_opcode G_SELECT):$root,
+    [{ return Helper.matchUndefSelectCmp(*${root}); }]),
+  (apply [{ return Helper.replaceSingleDefInstWithOperand(*${root}, 2); }])
+>;
+
+// Fold (true ? x : y) -> x
+// Fold (false ? x : y) -> y
+def select_constant_cmp_matchdata : GIDefMatchData<"unsigned">;
+def select_constant_cmp: GICombineRule<
+  (defs root:$root, select_constant_cmp_matchdata:$matchinfo),
+  (match (wip_match_opcode G_SELECT):$root,
+    [{ return Helper.matchConstantSelectCmp(*${root}, ${matchinfo}); }]),
+  (apply [{ return Helper.replaceSingleDefInstWithOperand(*${root}, ${matchinfo}); }])
+>;
+
 // Fold x op 0 -> x
 def right_identity_zero: GICombineRule<
   (defs root:$root),
@@ -341,10 +359,12 @@ def identity_combines : GICombineGroup<[select_same_val, right_identity_zero,
 
 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]>;
 def all_combines : GICombineGroup<[trivial_combines, ptr_add_immed_chain,
     combines_for_extload, combine_indexed_load_store, undef_combines,
     identity_combines, simplify_add_to_sub,
     hoist_logic_op_with_same_opcode_hands,
     shl_ashr_to_sext_inreg, sext_inreg_of_load,
-    width_reduction_combines]>;
+    width_reduction_combines, select_combines]>;

lib/CodeGen/GlobalISel/CombinerHelper.cpp

@@ -1770,6 +1770,22 @@ bool CombinerHelper::matchUndefStore(MachineInstr &MI) {
                       MRI);
 }
 
+bool CombinerHelper::matchUndefSelectCmp(MachineInstr &MI) {
+  assert(MI.getOpcode() == TargetOpcode::G_SELECT);
+  return getOpcodeDef(TargetOpcode::G_IMPLICIT_DEF, MI.getOperand(1).getReg(),
+                      MRI);
+}
+
+bool CombinerHelper::matchConstantSelectCmp(MachineInstr &MI, unsigned &OpIdx) {
+  assert(MI.getOpcode() == TargetOpcode::G_SELECT);
+  if (auto MaybeCstCmp =
+          getConstantVRegValWithLookThrough(MI.getOperand(1).getReg(), MRI)) {
+    OpIdx = MaybeCstCmp->Value ? 2 : 3;
+    return true;
+  }
+  return false;
+}
+
 bool CombinerHelper::eraseInst(MachineInstr &MI) {
   MI.eraseFromParent();
   return true;

test/CodeGen/AArch64/GlobalISel/combine-select.mir

@@ -0,0 +1,62 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -run-pass=aarch64-prelegalizer-combiner -verify-machineinstrs -mtriple aarch64-unknown-unknown %s -o - | FileCheck %s
+# RUN: llc -debugify-and-strip-all-safe -run-pass=aarch64-prelegalizer-combiner -verify-machineinstrs -mtriple aarch64-unknown-unknown %s -o - | FileCheck %s
+---
+# select (c, x, x) -> x
+name:            test_combine_select_same_res
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: test_combine_select_same_res
+    ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
+    ; CHECK: $x0 = COPY [[COPY]](s64)
+    %0:_(s64) = COPY $x0
+    %1:_(s1) = G_TRUNC %0
+    %2:_(s64) = G_SELECT %1, %0, %0
+    $x0 = COPY %2(s64)
+...
+---
+# select (undef, x, y) -> y
+name:            test_combine_select_undef_res0_res1
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: test_combine_select_undef_res0_res1
+    ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
+    ; CHECK: $x0 = COPY [[COPY]](s64)
+    %0:_(s64) = COPY $x0
+    %1:_(s64) = COPY $x1
+    %2:_(s1) = G_IMPLICIT_DEF
+    %3:_(s64) = G_SELECT %2, %0, %1
+    $x0 = COPY %3(s64)
+...
+---
+# select (false, x, y) -> y
+name:            test_combine_select_false_res0_res1
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: test_combine_select_false_res0_res1
+    ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x1
+    ; CHECK: $x0 = COPY [[COPY]](s64)
+    %0:_(s64) = COPY $x0
+    %1:_(s64) = COPY $x1
+    %2:_(s1) = G_CONSTANT i1 false
+    %3:_(s64) = G_SELECT %2, %0, %1
+    $x0 = COPY %3(s64)
+...
+---
+# select (true, x, y) -> x
+name:            test_combine_select_true_res0_res1
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: test_combine_select_true_res0_res1
+    ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
+    ; CHECK: $x0 = COPY [[COPY]](s64)
+    %0:_(s64) = COPY $x0
+    %1:_(s64) = COPY $x1
+    %2:_(s1) = G_CONSTANT i1 true
+    %3:_(s64) = G_SELECT %2, %0, %1
+    $x0 = COPY %3(s64)
+...

test/CodeGen/AMDGPU/GlobalISel/postlegalizercombiner-select.mir

@@ -6,19 +6,23 @@ name:            select_from_different_results_of_unmerge_values
 tracksRegLiveness: true
 body:             |
   bb.0:
+    liveins: $vgpr0
 
     ; GCN-LABEL: name: select_from_different_results_of_unmerge_values
+    ; GCN: liveins: $vgpr0
     ; GCN: [[DEF:%[0-9]+]]:_(<2 x s32>) = G_IMPLICIT_DEF
-    ; GCN: [[DEF1:%[0-9]+]]:_(s1) = G_IMPLICIT_DEF
+    ; GCN: [[COPY:%[0-9]+]]:_(s32) = COPY $vgpr0
+    ; GCN: [[TRUNC:%[0-9]+]]:_(s1) = G_TRUNC [[COPY]](s32)
     ; GCN: [[UV:%[0-9]+]]:_(s32), [[UV1:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[DEF]](<2 x s32>)
-    ; GCN: [[SELECT:%[0-9]+]]:_(s32) = G_SELECT [[DEF1]](s1), [[UV]], [[UV1]]
+    ; GCN: [[SELECT:%[0-9]+]]:_(s32) = G_SELECT [[TRUNC]](s1), [[UV]], [[UV1]]
     ; GCN: $vgpr0 = COPY [[SELECT]](s32)
     ; GCN: SI_RETURN_TO_EPILOG $vgpr0
-    %2:_(<2 x s32>) = G_IMPLICIT_DEF
-    %4:_(s1) = G_IMPLICIT_DEF
-    %0:_(s32), %1:_(s32) = G_UNMERGE_VALUES %2:_(<2 x s32>)
-    %3:_(s32) = G_SELECT %4:_(s1), %0:_, %1:_
-    $vgpr0 = COPY %3
+    %0:_(<2 x s32>) = G_IMPLICIT_DEF
+    %1:_(s32) = COPY $vgpr0
+    %2:_(s1) = G_TRUNC %1:_(s32)
+    %3:_(s32), %4:_(s32) = G_UNMERGE_VALUES %0:_(<2 x s32>)
+    %5:_(s32) = G_SELECT %2:_(s1), %3:_, %4:_
+    $vgpr0 = COPY %5
     SI_RETURN_TO_EPILOG $vgpr0
 
 ...
@@ -28,17 +32,20 @@ name:            select_from_same_results_of_unmerge_values
 tracksRegLiveness: true
 body:             |
   bb.0:
+    liveins: $vgpr0
 
     ; GCN-LABEL: name: select_from_same_results_of_unmerge_values
+    ; GCN: liveins: $vgpr0
     ; GCN: [[DEF:%[0-9]+]]:_(<2 x s32>) = G_IMPLICIT_DEF
     ; GCN: [[UV:%[0-9]+]]:_(s32), [[UV1:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[DEF]](<2 x s32>)
     ; GCN: $vgpr0 = COPY [[UV]](s32)
     ; GCN: SI_RETURN_TO_EPILOG $vgpr0
-    %2:_(<2 x s32>) = G_IMPLICIT_DEF
-    %4:_(s1) = G_IMPLICIT_DEF
-    %0:_(s32), %1:_(s32) = G_UNMERGE_VALUES %2:_(<2 x s32>)
-    %3:_(s32) = G_SELECT %4:_(s1), %0:_, %0:_
-    $vgpr0 = COPY %3
+    %0:_(<2 x s32>) = G_IMPLICIT_DEF
+    %1:_(s32) = COPY $vgpr0
+    %2:_(s1) = G_TRUNC %1:_(s32)
+    %3:_(s32), %4:_(s32) = G_UNMERGE_VALUES %0:_(<2 x s32>)
+    %5:_(s32) = G_SELECT %2:_(s1), %3:_, %3:_
+    $vgpr0 = COPY %5
     SI_RETURN_TO_EPILOG $vgpr0
 
 ...