1 files changed, 36 insertions, 6 deletions

diff --git a/llvm/lib/CodeGen/GlobalISel/Utils.cpp b/llvm/lib/CodeGen/GlobalISel/Utils.cpp
index 080600d3cc98..eaf829f562b2 100644
--- a/llvm/lib/CodeGen/GlobalISel/Utils.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/Utils.cpp
@@ -205,8 +205,15 @@ bool llvm::canReplaceReg(Register DstReg, Register SrcReg,
     return false;
   // Replace if either DstReg has no constraints or the register
   // constraints match.
-  return !MRI.getRegClassOrRegBank(DstReg) ||
-         MRI.getRegClassOrRegBank(DstReg) == MRI.getRegClassOrRegBank(SrcReg);
+  const auto &DstRBC = MRI.getRegClassOrRegBank(DstReg);
+  if (!DstRBC || DstRBC == MRI.getRegClassOrRegBank(SrcReg))
+    return true;
+
+  // Otherwise match if the Src is already a regclass that is covered by the Dst
+  // RegBank.
+  return DstRBC.is<const RegisterBank *>() && MRI.getRegClassOrNull(SrcReg) &&
+         DstRBC.get<const RegisterBank *>()->covers(
+             *MRI.getRegClassOrNull(SrcReg));
 }
 
 bool llvm::isTriviallyDead(const MachineInstr &MI,
@@ -773,6 +780,29 @@ std::optional<APInt> llvm::ConstantFoldExtOp(unsigned Opcode,
   return std::nullopt;
 }
 
+std::optional<APInt> llvm::ConstantFoldCastOp(unsigned Opcode, LLT DstTy,
+                                              const Register Op0,
+                                              const MachineRegisterInfo &MRI) {
+  std::optional<APInt> Val = getIConstantVRegVal(Op0, MRI);
+  if (!Val)
+    return Val;
+
+  const unsigned DstSize = DstTy.getScalarSizeInBits();
+
+  switch (Opcode) {
+  case TargetOpcode::G_SEXT:
+    return Val->sext(DstSize);
+  case TargetOpcode::G_ZEXT:
+  case TargetOpcode::G_ANYEXT:
+    // TODO: DAG considers target preference when constant folding any_extend.
+    return Val->zext(DstSize);
+  default:
+    break;
+  }
+
+  llvm_unreachable("unexpected cast opcode to constant fold");
+}
+
 std::optional<APFloat>
 llvm::ConstantFoldIntToFloat(unsigned Opcode, LLT DstTy, Register Src,
                              const MachineRegisterInfo &MRI) {
@@ -1086,9 +1116,9 @@ std::optional<APInt>
 llvm::getIConstantSplatVal(const Register Reg, const MachineRegisterInfo &MRI) {
   if (auto SplatValAndReg =
           getAnyConstantSplat(Reg, MRI, /* AllowUndef */ false)) {
-    std::optional<ValueAndVReg> ValAndVReg =
-        getIConstantVRegValWithLookThrough(SplatValAndReg->VReg, MRI);
-    return ValAndVReg->Value;
+    if (std::optional<ValueAndVReg> ValAndVReg =
+        getIConstantVRegValWithLookThrough(SplatValAndReg->VReg, MRI))
+      return ValAndVReg->Value;
   }
 
   return std::nullopt;
@@ -1143,7 +1173,7 @@ llvm::getVectorSplat(const MachineInstr &MI, const MachineRegisterInfo &MRI) {
   if (auto Splat = getIConstantSplatSExtVal(MI, MRI))
     return RegOrConstant(*Splat);
   auto Reg = MI.getOperand(1).getReg();
-  if (any_of(make_range(MI.operands_begin() + 2, MI.operands_end()),
+  if (any_of(drop_begin(MI.operands(), 2),
              [&Reg](const MachineOperand &Op) { return Op.getReg() != Reg; }))
     return std::nullopt;
   return RegOrConstant(Reg);