[RISCV] Preserve live constants when folding extensions

2 files changed, 52 insertions, 2 deletions

diff --git a/backend/lib/Target/RISCV/RISCVCodeGen.cpp b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
index 9cf7b53..bc437c0 100644
--- a/backend/lib/Target/RISCV/RISCVCodeGen.cpp
+++ b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
@@ -127,7 +127,8 @@ extendIntegerImmediate(const IntegerImmediate &Value, std::uint32_t DestinationW
     return makeUnexpected(ErrorCode::VerificationError,
                           "RISC-V constant extension requires virtual-register operands");
   }
-  const MachineOperand *SourceDefinition = Function.registerInfo().getUniqueVRegDef(*Source);
+  const Register SourceRegister = *Source;
+  const MachineOperand *SourceDefinition = Function.registerInfo().getUniqueVRegDef(SourceRegister);
   const MachineInstr *SourceDef = SourceDefinition == nullptr ? nullptr : SourceDefinition->parent();
   if (SourceDef == nullptr || SourceDef->opcode() != opcode(GenericOpcode::G_CONSTANT))
     return makeUnexpected(ErrorCode::Unsupported,
@@ -161,7 +162,7 @@ extendIntegerImmediate(const IntegerImmediate &Value, std::uint32_t DestinationW
   if (!Added)
     return std::unexpected(Added.error());
 
-  if (Function.registerInfo().useEmpty(*Source)) {
+  if (Function.registerInfo().useEmpty(SourceRegister)) {
     MachineInstr *Def = const_cast<MachineInstr *>(SourceDef);
     MachineBasicBlock *Parent = Def->parent();
     if (Parent == nullptr)
diff --git a/backend/unittests/Target/RISCVCodeGenTest.cpp b/backend/unittests/Target/RISCVCodeGenTest.cpp
index 9c44b28..12e59ba 100644
--- a/backend/unittests/Target/RISCVCodeGenTest.cpp
+++ b/backend/unittests/Target/RISCVCodeGenTest.cpp
@@ -982,6 +982,54 @@ void testAtomicCompareExchangeLegalization(TestContext &Context) {
       "compare-exchange-with-success must lower to atomic, equality test, and old-value copy");
 }
 
+void testConstantExtensionPreservesLiveSource(TestContext &Context) {
+  auto Target = rcc::riscv::createTargetDescription();
+  auto Function = rcc::MachineFunction::create("riscv_live_constant_extension",
+                                               rcc::MachineFunctionStableID::fromValue(1300));
+  auto I32 = rcc::LLT::makeAnyScalar(32);
+  auto I64 = rcc::LLT::makeAnyScalar(64);
+  if (!Target || !Function || !I32 || !I64)
+    return;
+  auto Block = (*Function)->createBasicBlock();
+  if (!Block || !(*Function)->appendBlock(**Block))
+    return;
+  auto Source = (*Function)->registerInfo().createVirtualRegister(*I32);
+  auto Extended = (*Function)->registerInfo().createVirtualRegister(*I64);
+  auto PreservedUse = (*Function)->registerInfo().createVirtualRegister(*I32);
+  auto Immediate = rcc::IntegerImmediate::fromUnsigned(32, 7);
+  if (!Source || !Extended || !PreservedUse || !Immediate)
+    return;
+
+  rcc::MachineIRBuilder Builder(**Function, *Target);
+  if (!Builder.setInsertPoint(**Block))
+    return;
+  auto Constant = Builder.buildInstr(rcc::opcode(rcc::GenericOpcode::G_CONSTANT));
+  if (!Constant || !Constant->addDef(*Source) ||
+      !Constant->addIntegerImmediate(std::move(*Immediate)))
+    return;
+  auto BuiltConstant = Constant->build();
+  auto Extension = Builder.buildInstr(rcc::opcode(rcc::GenericOpcode::G_ANYEXT));
+  if (!BuiltConstant || !Extension || !Extension->addDef(*Extended) || !Extension->addUse(*Source))
+    return;
+  auto BuiltExtension = Extension->build();
+  auto Copy = Builder.buildInstr(rcc::opcode(rcc::GenericOpcode::COPY));
+  if (!BuiltExtension || !Copy || !Copy->addDef(*PreservedUse) || !Copy->addUse(*Source) ||
+      !Copy->build())
+    return;
+
+  auto Action = rcc::riscv::legalizerInfo().classify(**BuiltExtension, **Function, *Target);
+  auto Lowered =
+      Action ? rcc::riscv::legalizerInfo().lower(**BuiltExtension, **Function, *Target, *Action)
+             : rcc::Expected<void>(std::unexpected(Action.error()));
+  Context.expect(Action && *Action == rcc::LegalizationAction::Custom && Lowered &&
+                     (*BuiltExtension)->opcode() == rcc::opcode(rcc::GenericOpcode::G_CONSTANT) &&
+                     (*BuiltConstant)->parent() != nullptr &&
+                     (*Function)->registerInfo().getUniqueVRegDef(*Source) ==
+                         (*BuiltConstant)->operandAt(0) &&
+                     (*Function)->registerInfo().verifyDefUseChains(),
+                 "constant-extension folding must preserve a source definition with live uses");
+}
+
 void testAtomicSelection(TestContext &Context) {
   auto Target = rcc::riscv::createTargetDescription();
   if (!Target)
@@ -1260,6 +1308,7 @@ int main() {
   testCompositeScalarSelection(Context);
   testOverflowSelection(Context);
   testAtomicCompareExchangeLegalization(Context);
+  testConstantExtensionPreservesLiveSource(Context);
   testAtomicSelection(Context);
   testUnsupportedTypes(Context);
   testMemoryOffsetSelection(Context);