[GlobalISel] Implement undef combines for funnel shifts

[paquette]

Summary

AFAICT the way SDAG does this is:

fshl (x, y, undef) -> x for scalars
fshr (x, y, undef) -> y for scalars
fshl (x, y, undef) -> undef for vectors
fshr (x, y, undef) -> undef for vectors

https://gcc.godbolt.org/z/xq9PYfhYh

This should fix https://github.com/llvm/llvm-project/issues/57256

[arsenm]

Doesn't remove whatever was incorrectly folding this before

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

763–764

Different behavior for vector and scalar sounds really broken

[aemerson]

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

763–764

Yeah this sounds very sus.

[arsenm]

The different scalar and vector behavior is broken and should be fixed

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

public:
 
   bool matchBuildVectorIdentityFold(MachineInstr &MI, Register &MatchInfo);
   bool matchTruncBuildVectorFold(MachineInstr &MI, Register &MatchInfo);
 
   /// \returns true if it is possible to simplify a select instruction \p MI
   /// to a min/max instruction of some sort.
   bool matchSimplifySelectToMinMax(MachineInstr &MI, BuildFnTy &MatchInfo);
 
+  /// fshl (x, y, undef) -> x for scalars
+  /// fshr (x, y, undef) -> y for scalars
+  /// fshl (x, y, undef) -> undef for vectors
+  /// fshr (x, y, undef) -> undef for vectors
+  void applyUndefFunnelShiftAmount(MachineInstr &MI);
+
 private:
   /// Given a non-indexed load or store instruction \p MI, find an offset that
   /// can be usefully and legally folded into it as a post-indexing operation.
   ///
   /// \returns true if a candidate is found.
   bool findPostIndexCandidate(MachineInstr &MI, Register &Addr, Register &Base,
                               Register &Offset);
 

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

 
 def funnel_shift_to_rotate : GICombineRule<
   (defs root:$root),
   (match (wip_match_opcode G_FSHL, G_FSHR):$root,
     [{ return Helper.matchFunnelShiftToRotate(*${root}); }]),
   (apply [{ Helper.applyFunnelShiftToRotate(*${root}); }])
 >;
 
+def undef_funnel_shift_amt: GICombineRule<
+  (defs root:$root),
+  (match (wip_match_opcode G_FSHL, G_FSHR):$root,
+    [{ return Helper.matchOperandIsUndef(*${root}, 3); }]),
+  (apply [{ Helper.applyUndefFunnelShiftAmount(*${root}); }])
+>;
+
 def rotate_out_of_range : GICombineRule<
   (defs root:$root),
   (match (wip_match_opcode G_ROTR, G_ROTL):$root,
     [{ return Helper.matchRotateOutOfRange(*${root}); }]),
   (apply [{ Helper.applyRotateOutOfRange(*${root}); }])
 >;
 
 def icmp_to_true_false_known_bits : GICombineRule<
def undef_combines : GICombineGroup<[undef_to_fp_zero, undef_to_int_zero,
                                      undef_to_negative_one,
                                      binop_left_undef_to_zero,
                                      binop_right_undef_to_undef,
                                      propagate_undef_any_op,
                                      propagate_undef_all_ops,
                                      propagate_undef_shuffle_mask,
                                      erase_undef_store,
                                      unmerge_undef,
-                                     insert_extract_vec_elt_out_of_bounds]>;
+                                     insert_extract_vec_elt_out_of_bounds,
+                                     undef_funnel_shift_amt]>;
 
 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,
                                         add_sub_reg, buildvector_identity_fold,
                                         trunc_buildvector_fold]>;

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

bool CombinerHelper::matchSimplifySelectToMinMax(MachineInstr &MI,
   if (mi_match(Cond, MRI, m_OneNonDBGUse(m_GTrunc(m_Reg(MaybeTrunc)))))
     Cond = MaybeTrunc;
   Register Dst = MI.getOperand(0).getReg();
   Register TrueVal = MI.getOperand(2).getReg();
   Register FalseVal = MI.getOperand(3).getReg();
   return matchFPSelectToMinMax(Dst, Cond, TrueVal, FalseVal, MatchInfo);
 }
 
+void CombinerHelper::applyUndefFunnelShiftAmount(MachineInstr &MI) {
+  if (MRI.getType(MI.getOperand(0).getReg()).isVector()) {
+    replaceInstWithUndef(MI);
+    return;
+  }
+  replaceSingleDefInstWithOperand(
+      MI, MI.getOpcode() == TargetOpcode::G_FSHL ? 1 : 2);
+}
+
 bool CombinerHelper::tryCombine(MachineInstr &MI) {
   if (tryCombineCopy(MI))
     return true;
   if (tryCombineExtendingLoads(MI))
     return true;
   if (tryCombineIndexedLoadStore(MI))
     return true;
   return false;
 }

llvm/test/CodeGen/AArch64/GlobalISel/combine-undef-funnel-shift-amt.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64-unknown-unknown -run-pass=aarch64-prelegalizer-combiner -verify-machineinstrs -o - %s | FileCheck %s
+
+...
+---
+name:            scalar_fshl
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; fshl (op1, op2, undef) -> op1 for scalars
+    ; CHECK-LABEL: name: scalar_fshl
+    ; CHECK: %op1:_(s32) = G_CONSTANT i32 1
+    ; CHECK-NEXT: $w0 = COPY %op1(s32)
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+    %op1:_(s32) = G_CONSTANT i32 1
+    %op2:_(s32) = G_CONSTANT i32 10000
+    %undef:_(s32) = G_IMPLICIT_DEF
+    %shift:_(s32) = G_FSHL %op1, %op2, %undef(s32)
+    $w0 = COPY %shift(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            scalar_fshr
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; fshr (op1, op2, undef) -> op2 for scalars
+    ; CHECK-LABEL: name: scalar_fshr
+    ; CHECK: %op2:_(s32) = G_CONSTANT i32 2
+    ; CHECK-NEXT: $w0 = COPY %op2(s32)
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+    %op1:_(s32) = G_CONSTANT i32 1
+    %op2:_(s32) = G_CONSTANT i32 2
+    %undef:_(s32) = G_IMPLICIT_DEF
+    %shift:_(s32) = G_FSHR %op1, %op2, %undef(s32)
+    $w0 = COPY %shift(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            vector_fshl
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $q0
+    ; fshl (op1, op2, undef) -> undef for vectors
+    ; CHECK-LABEL: name: vector_fshl
+    ; CHECK: liveins: $q0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %undef:_(<4 x s32>) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: $q0 = COPY %undef(<4 x s32>)
+    ; CHECK-NEXT: RET_ReallyLR implicit $q0
+    %op1:_(<4 x s32>) = COPY $q0
+    %op2:_(<4 x s32>) = COPY $q0
+    %undef:_(<4 x s32>) = G_IMPLICIT_DEF
+    %shift:_(<4 x s32>) = G_FSHL %op1, %op2, %undef(<4 x s32>)
+    $q0 = COPY %shift(<4 x s32>)
+    RET_ReallyLR implicit $q0
+...
+---
+name:            vector_fshr
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $q0
+    ; fshr (op1, op2, undef) -> undef for vectors
+    ; CHECK-LABEL: name: vector_fshr
+    ; CHECK: liveins: $q0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %undef:_(<4 x s32>) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: $q0 = COPY %undef(<4 x s32>)
+    ; CHECK-NEXT: RET_ReallyLR implicit $q0
+    %op1:_(<4 x s32>) = COPY $q0
+    %op2:_(<4 x s32>) = COPY $q0
+    %undef:_(<4 x s32>) = G_IMPLICIT_DEF
+    %shift:_(<4 x s32>) = G_FSHR %op1, %op2, %undef(<4 x s32>)
+    $q0 = COPY %shift(<4 x s32>)
+    RET_ReallyLR implicit $q0
+...