GlobalISel: Add MIPatternMatch for G_ICMP/G_FCMP

arsenm · arsenm · commit c77bbea9a671 · 2020-01-23T13:30:47.000-08:00
diff --git a/llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h b/llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h
@@ -16,6 +16,7 @@
 #include "llvm/ADT/APInt.h"
 #include "llvm/CodeGen/GlobalISel/Utils.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/IR/InstrTypes.h"
 
 namespace llvm {
 namespace MIPatternMatch {
@@ -163,6 +164,8 @@ template <typename Class> struct bind_ty {
 inline bind_ty<Register> m_Reg(Register &R) { return R; }
 inline bind_ty<MachineInstr *> m_MInstr(MachineInstr *&MI) { return MI; }
 inline bind_ty<LLT> m_Type(LLT &Ty) { return Ty; }
+inline bind_ty<CmpInst::Predicate> m_Pred(CmpInst::Predicate &P) { return P; }
+inline operand_type_match m_Pred() { return operand_type_match(); }
 
 // Helper for matching G_FCONSTANT
 inline bind_ty<const ConstantFP *> m_GFCst(const ConstantFP *&C) { return C; }
@@ -320,6 +323,45 @@ inline UnaryOp_match<SrcTy, TargetOpcode::COPY> m_Copy(SrcTy &&Src) {
   return UnaryOp_match<SrcTy, TargetOpcode::COPY>(std::forward<SrcTy>(Src));
 }
 
+// General helper for generic MI compares, i.e. G_ICMP and G_FCMP
+// TODO: Allow checking a specific predicate.
+template <typename Pred_P, typename LHS_P, typename RHS_P, unsigned Opcode>
+struct CompareOp_match {
+  Pred_P P;
+  LHS_P L;
+  RHS_P R;
+
+  CompareOp_match(const Pred_P &Pred, const LHS_P &LHS, const RHS_P &RHS)
+      : P(Pred), L(LHS), R(RHS) {}
+
+  template <typename OpTy>
+  bool match(const MachineRegisterInfo &MRI, OpTy &&Op) {
+    MachineInstr *TmpMI;
+    if (!mi_match(Op, MRI, m_MInstr(TmpMI)) || TmpMI->getOpcode() != Opcode)
+      return false;
+
+    auto TmpPred =
+        static_cast<CmpInst::Predicate>(TmpMI->getOperand(1).getPredicate());
+    if (!P.match(MRI, TmpPred))
+      return false;
+
+    return L.match(MRI, TmpMI->getOperand(2).getReg()) &&
+           R.match(MRI, TmpMI->getOperand(3).getReg());
+  }
+};
+
+template <typename Pred, typename LHS, typename RHS>
+inline CompareOp_match<Pred, LHS, RHS, TargetOpcode::G_ICMP>
+m_GICmp(const Pred &P, const LHS &L, const RHS &R) {
+  return CompareOp_match<Pred, LHS, RHS, TargetOpcode::G_ICMP>(P, L, R);
+}
+
+template <typename Pred, typename LHS, typename RHS>
+inline CompareOp_match<Pred, LHS, RHS, TargetOpcode::G_FCMP>
+m_GFCmp(const Pred &P, const LHS &L, const RHS &R) {
+  return CompareOp_match<Pred, LHS, RHS, TargetOpcode::G_FCMP>(P, L, R);
+}
+
 // Helper for checking if a Reg is of specific type.
 struct CheckType {
   LLT Ty;
diff --git a/llvm/unittests/CodeGen/GlobalISel/PatternMatchTest.cpp b/llvm/unittests/CodeGen/GlobalISel/PatternMatchTest.cpp
@@ -129,6 +129,56 @@ TEST_F(GISelMITest, MatchBinaryOp) {
   EXPECT_EQ(Src1, Copies[1]);
 }
 
+TEST_F(GISelMITest, MatchICmp) {
+  setUp();
+  if (!TM)
+    return;
+
+  const LLT s1 = LLT::scalar(1);
+  auto CmpEq = B.buildICmp(CmpInst::ICMP_EQ, s1, Copies[0], Copies[1]);
+
+  // Check match any predicate.
+  bool match =
+      mi_match(CmpEq.getReg(0), *MRI, m_GICmp(m_Pred(), m_Reg(), m_Reg()));
+  EXPECT_TRUE(match);
+
+  // Check we get the predicate and registers.
+  CmpInst::Predicate Pred;
+  Register Reg0;
+  Register Reg1;
+  match = mi_match(CmpEq.getReg(0), *MRI,
+                   m_GICmp(m_Pred(Pred), m_Reg(Reg0), m_Reg(Reg1)));
+  EXPECT_TRUE(match);
+  EXPECT_EQ(CmpInst::ICMP_EQ, Pred);
+  EXPECT_EQ(Copies[0], Reg0);
+  EXPECT_EQ(Copies[1], Reg1);
+}
+
+TEST_F(GISelMITest, MatchFCmp) {
+  setUp();
+  if (!TM)
+    return;
+
+  const LLT s1 = LLT::scalar(1);
+  auto CmpEq = B.buildFCmp(CmpInst::FCMP_OEQ, s1, Copies[0], Copies[1]);
+
+  // Check match any predicate.
+  bool match =
+      mi_match(CmpEq.getReg(0), *MRI, m_GFCmp(m_Pred(), m_Reg(), m_Reg()));
+  EXPECT_TRUE(match);
+
+  // Check we get the predicate and registers.
+  CmpInst::Predicate Pred;
+  Register Reg0;
+  Register Reg1;
+  match = mi_match(CmpEq.getReg(0), *MRI,
+                   m_GFCmp(m_Pred(Pred), m_Reg(Reg0), m_Reg(Reg1)));
+  EXPECT_TRUE(match);
+  EXPECT_EQ(CmpInst::FCMP_OEQ, Pred);
+  EXPECT_EQ(Copies[0], Reg0);
+  EXPECT_EQ(Copies[1], Reg1);
+}
+
 TEST_F(GISelMITest, MatchFPUnaryOp) {
   setUp();
   if (!TM)
