From ebe408ad8003c946ef871b955ab18e64e82697cb Mon Sep 17 00:00:00 2001
From: Petar Avramovic <Petar.Avramovic@amd.com>
Date: Tue, 27 Apr 2021 10:53:17 +0200
Subject: [PATCH] [MIPatternMatch]: Add ICstRegMatch

Matches G_CONSTANT and returns its def register.

Differential Revision: https://reviews.llvm.org/D99734
---
 llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h  | 16 ++++++++++++++++
 llvm/unittests/CodeGen/GlobalISel/PatternMatchTest.cpp | 11 +++++++++++
 2 files changed, 27 insertions(+)

diff --git a/llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h b/llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h
index ebcb934..e0e6eeb 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h
@@ -68,6 +68,22 @@ struct ConstantMatch {
 
 inline ConstantMatch m_ICst(int64_t &Cst) { return ConstantMatch(Cst); }
 
+struct ICstRegMatch {
+  Register &CR;
+  ICstRegMatch(Register &C) : CR(C) {}
+  bool match(const MachineRegisterInfo &MRI, Register Reg) {
+    if (auto MaybeCst = getConstantVRegValWithLookThrough(
+            Reg, MRI, /*LookThroughInstrs*/ true,
+            /*HandleFConstants*/ false)) {
+      CR = MaybeCst->VReg;
+      return true;
+    }
+    return false;
+  }
+};
+
+inline ICstRegMatch m_ICst(Register &Reg) { return ICstRegMatch(Reg); }
+
 /// Matcher for a specific constant value.
 struct SpecificConstantMatch {
   int64_t RequestedVal;
diff --git a/llvm/unittests/CodeGen/GlobalISel/PatternMatchTest.cpp b/llvm/unittests/CodeGen/GlobalISel/PatternMatchTest.cpp
index a8209fe2..8801d37 100644
--- a/llvm/unittests/CodeGen/GlobalISel/PatternMatchTest.cpp
+++ b/llvm/unittests/CodeGen/GlobalISel/PatternMatchTest.cpp
@@ -40,6 +40,17 @@ TEST_F(AArch64GISelMITest, MatchIntConstant) {
   EXPECT_EQ(Cst, 42);
 }
 
+TEST_F(AArch64GISelMITest, MatchIntConstantRegister) {
+  setUp();
+  if (!TM)
+    return;
+  auto MIBCst = B.buildConstant(LLT::scalar(64), 42);
+  Register Src0;
+  bool match = mi_match(MIBCst.getReg(0), *MRI, m_ICst(Src0));
+  EXPECT_TRUE(match);
+  EXPECT_EQ(Src0, MIBCst.getReg(0));
+}
+
 TEST_F(AArch64GISelMITest, MatchBinaryOp) {
   setUp();
   if (!TM)
-- 
2.7.4