[RISCV] Fold narrow binary artifact chains

2 files changed, 118 insertions, 15 deletions

diff --git a/backend/lib/Target/RISCV/RISCVCodeGen.cpp b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
index e790edf..8676a2b 100644
--- a/backend/lib/Target/RISCV/RISCVCodeGen.cpp
+++ b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
@@ -417,7 +417,7 @@ struct NarrowAnyExtBinaryPlan final {
   MachineInstr *Binary = nullptr;
   GenericOpcode Generic = GenericOpcode::Invalid;
   std::array<Register, 2> WideInputs{};
-  std::vector<MachineInstr *> Truncations;
+  std::vector<MachineInstr *> Artifacts;
   std::vector<MachineInstr *> Constants;
   std::uint32_t SourceWidth = 0;
 };
@@ -449,14 +449,35 @@ narrowAnyExtBinaryPlan(const MachineInstr &Extension, const MachineFunction &Fun
     return std::nullopt;
   const GenericOpcode Generic = static_cast<GenericOpcode>(Binary->opcode().value());
   if (Generic != GenericOpcode::G_ADD && Generic != GenericOpcode::G_SUB &&
-      Generic != GenericOpcode::G_AND && Generic != GenericOpcode::G_OR &&
-      Generic != GenericOpcode::G_XOR && Generic != GenericOpcode::G_SHL)
+      Generic != GenericOpcode::G_MUL && Generic != GenericOpcode::G_AND &&
+      Generic != GenericOpcode::G_OR && Generic != GenericOpcode::G_XOR &&
+      Generic != GenericOpcode::G_SHL)
     return std::nullopt;
 
   NarrowAnyExtBinaryPlan Plan;
   Plan.Binary = Binary;
   Plan.Generic = Generic;
   Plan.SourceWidth = SourceType.scalarBitWidth();
+  const auto AddArtifact = [&Plan](MachineInstr *Artifact) {
+    if (std::ranges::find(Plan.Artifacts, Artifact) == Plan.Artifacts.end())
+      Plan.Artifacts.push_back(Artifact);
+  };
+  const auto TruncatedWideSource = [&](MachineInstr *Truncation,
+                                       MachineInstr *OnlyUse) -> std::optional<Register> {
+    if (Truncation == nullptr || Truncation->opcode() != opcode(GenericOpcode::G_TRUNC) ||
+        Truncation->operandCount() != 2)
+      return std::nullopt;
+    const MachineOperand *NarrowOperand = Truncation->operandAt(0);
+    const Register *Narrow = NarrowOperand == nullptr ? nullptr : NarrowOperand->getIf<Register>();
+    const MachineOperand *WideOperand = Truncation->operandAt(1);
+    const Register *Wide = WideOperand == nullptr ? nullptr : WideOperand->getIf<Register>();
+    const std::optional<LLT> WideType =
+        WideOperand == nullptr ? std::nullopt : registerType(Function, *WideOperand);
+    if (Narrow == nullptr || Wide == nullptr || !Wide->isVirtual() || !WideType ||
+        *WideType != DestinationType || singleRegisterUse(Function, *Narrow) != OnlyUse)
+      return std::nullopt;
+    return *Wide;
+  };
   for (std::size_t Index = 0; Index != Plan.WideInputs.size(); ++Index) {
     const MachineOperand *Operand = Binary->operandAt(Index + 1U);
     const Register *Reg = Operand == nullptr ? nullptr : Operand->getIf<Register>();
@@ -474,19 +495,30 @@ narrowAnyExtBinaryPlan(const MachineInstr &Extension, const MachineFunction &Fun
 
     const MachineOperand *Definition = Function.registerInfo().getUniqueVRegDef(*Reg);
     MachineInstr *Def = Definition == nullptr ? nullptr : const_cast<MachineInstr *>(Definition->parent());
-    if (Def != nullptr && Def->opcode() == opcode(GenericOpcode::G_TRUNC) &&
-        Def->operandCount() == 2) {
-      const MachineOperand *WideOperand = Def->operandAt(1);
-      const Register *Wide = WideOperand == nullptr ? nullptr : WideOperand->getIf<Register>();
-      if (Wide == nullptr || !Wide->isVirtual())
-        return std::nullopt;
-      const auto WideInfo = Function.registerInfo().virtualRegisterInfo(*Wide);
-      if (!WideInfo || !WideInfo->get().Type || *WideInfo->get().Type != DestinationType)
-        return std::nullopt;
+    if (auto Wide = TruncatedWideSource(Def, Binary)) {
       Plan.WideInputs[Index] = *Wide;
-      Plan.Truncations.push_back(Def);
+      AddArtifact(Def);
       continue;
     }
+    if (Def != nullptr &&
+        (Def->opcode() == opcode(GenericOpcode::COPY) ||
+         Def->opcode() == opcode(GenericOpcode::G_ANYEXT)) &&
+        Def->operandCount() == 2 && singleRegisterUse(Function, *Reg) == Binary) {
+      const MachineOperand *NarrowOperand = Def->operandAt(1);
+      const Register *Narrow =
+          NarrowOperand == nullptr ? nullptr : NarrowOperand->getIf<Register>();
+      const MachineOperand *NarrowDefinition =
+          Narrow == nullptr ? nullptr : Function.registerInfo().getUniqueVRegDef(*Narrow);
+      MachineInstr *Truncation = NarrowDefinition == nullptr
+                                     ? nullptr
+                                     : const_cast<MachineInstr *>(NarrowDefinition->parent());
+      if (auto Wide = TruncatedWideSource(Truncation, Def)) {
+        Plan.WideInputs[Index] = *Wide;
+        AddArtifact(Def);
+        AddArtifact(Truncation);
+        continue;
+      }
+    }
     if (Def != nullptr && Def->opcode() == opcode(GenericOpcode::G_CONSTANT) &&
         singleRegisterUse(Function, *Reg) == Binary) {
       Plan.WideInputs[Index] = *Reg;
@@ -887,8 +919,10 @@ lowerAtomicCompareExchangeWithSuccess(MachineInstr &Instruction, MachineFunction
     if (!Replaced)
       return std::unexpected(Replaced.error());
   }
-  for (MachineInstr *Truncation : Plan->Truncations) {
-    auto Removed = Truncation->parent()->removeInstruction(*Truncation);
+  for (MachineInstr *Artifact : Plan->Artifacts) {
+    if (Artifact->parent() == nullptr)
+      continue;
+    auto Removed = Artifact->parent()->removeInstruction(*Artifact);
     if (!Removed)
       return std::unexpected(Removed.error());
   }
diff --git a/backend/unittests/Target/RISCVCodeGenTest.cpp b/backend/unittests/Target/RISCVCodeGenTest.cpp
index 40f84f2..ac79328 100644
--- a/backend/unittests/Target/RISCVCodeGenTest.cpp
+++ b/backend/unittests/Target/RISCVCodeGenTest.cpp
@@ -1126,6 +1126,74 @@ void testIntegerCompareReusesConstantExtension(TestContext &Context) {
                  "single-use constant extension folding must preserve the source layout position");
 }
 
+void testNarrowAnyExtBinaryArtifacts(TestContext &Context) {
+  auto Target = rcc::riscv::createTargetDescription();
+  auto Function = rcc::MachineFunction::create("riscv_narrow_binary_artifacts",
+                                               rcc::MachineFunctionStableID::fromValue(1301));
+  auto I1 = rcc::LLT::makeAnyScalar(1);
+  auto I32 = rcc::LLT::makeAnyScalar(32);
+  auto I64 = rcc::LLT::makeAnyScalar(64);
+  if (!Target || !Function || !I1 || !I32 || !I64)
+    return;
+  auto Block = (*Function)->createBasicBlock();
+  if (!Block || !(*Function)->appendBlock(**Block))
+    return;
+  auto WideLeft = (*Function)->registerInfo().createVirtualRegister(*I64);
+  auto WideRight = (*Function)->registerInfo().createVirtualRegister(*I64);
+  auto NarrowLeft = (*Function)->registerInfo().createVirtualRegister(*I32);
+  auto CopiedLeft = (*Function)->registerInfo().createVirtualRegister(*I32);
+  auto NarrowRight = (*Function)->registerInfo().createVirtualRegister(*I1);
+  auto ExtendedRight = (*Function)->registerInfo().createVirtualRegister(*I32);
+  auto Product = (*Function)->registerInfo().createVirtualRegister(*I32);
+  auto Result = (*Function)->registerInfo().createVirtualRegister(*I64);
+  if (!WideLeft || !WideRight || !NarrowLeft || !CopiedLeft || !NarrowRight || !ExtendedRight ||
+      !Product || !Result)
+    return;
+
+  rcc::MachineIRBuilder Builder(**Function, *Target);
+  if (!Builder.setInsertPoint(**Block))
+    return;
+  const auto BuildUnary = [&](rcc::GenericOpcode Generic, rcc::Register Destination,
+                              rcc::Register Source) -> rcc::MachineInstr * {
+    auto Pending = Builder.buildInstr(rcc::opcode(Generic));
+    if (!Pending || !Pending->addDef(Destination) || !Pending->addUse(Source))
+      return nullptr;
+    auto Built = Pending->build();
+    return Built ? *Built : nullptr;
+  };
+  rcc::MachineInstr *LeftTrunc = BuildUnary(rcc::GenericOpcode::G_TRUNC, *NarrowLeft, *WideLeft);
+  rcc::MachineInstr *LeftCopy = BuildUnary(rcc::GenericOpcode::COPY, *CopiedLeft, *NarrowLeft);
+  rcc::MachineInstr *RightTrunc = BuildUnary(rcc::GenericOpcode::G_TRUNC, *NarrowRight, *WideRight);
+  rcc::MachineInstr *RightExtension =
+      BuildUnary(rcc::GenericOpcode::G_ANYEXT, *ExtendedRight, *NarrowRight);
+  auto Multiply = Builder.buildInstr(rcc::opcode(rcc::GenericOpcode::G_MUL));
+  if (LeftTrunc == nullptr || LeftCopy == nullptr || RightTrunc == nullptr ||
+      RightExtension == nullptr || !Multiply || !Multiply->addDef(*Product) ||
+      !Multiply->addUse(*CopiedLeft) || !Multiply->addUse(*ExtendedRight))
+    return;
+  auto BuiltMultiply = Multiply->build();
+  rcc::MachineInstr *OuterExtension = BuildUnary(rcc::GenericOpcode::G_ANYEXT, *Result, *Product);
+  if (!BuiltMultiply || OuterExtension == nullptr)
+    return;
+
+  auto Action = rcc::riscv::legalizerInfo().classify(*OuterExtension, **Function, *Target);
+  auto Lowered =
+      Action ? rcc::riscv::legalizerInfo().lower(*OuterExtension, **Function, *Target, *Action)
+             : rcc::Expected<void>(std::unexpected(Action.error()));
+  const rcc::Register *LegalLeft = OuterExtension->operandAt(1)->getIf<rcc::Register>();
+  const rcc::Register *LegalRight = OuterExtension->operandAt(2)->getIf<rcc::Register>();
+  Context.expect(
+      Action && *Action == rcc::LegalizationAction::Custom && Lowered &&
+          OuterExtension->opcode() == rcc::opcode(rcc::GenericOpcode::G_MUL) &&
+          LegalLeft != nullptr && *LegalLeft == *WideLeft && LegalRight != nullptr &&
+          *LegalRight == *WideRight && LeftTrunc->parent() == nullptr &&
+          LeftCopy->parent() == nullptr && RightTrunc->parent() == nullptr &&
+          RightExtension->parent() == nullptr && (*BuiltMultiply)->parent() == nullptr &&
+          (*Function)->registerInfo().verifyDefUseChains() &&
+          Target->validateInstruction(*OuterExtension, (*Function)->registerInfo()),
+      "narrow any-extension binary folding must remove copy, truncation, and extension artifacts");
+}
+
 void testAtomicSelection(TestContext &Context) {
   auto Target = rcc::riscv::createTargetDescription();
   if (!Target)
@@ -1406,6 +1474,7 @@ int main() {
   testAtomicCompareExchangeLegalization(Context);
   testConstantExtensionPreservesLiveSource(Context);
   testIntegerCompareReusesConstantExtension(Context);
+  testNarrowAnyExtBinaryArtifacts(Context);
   testAtomicSelection(Context);
   testUnsupportedTypes(Context);
   testMemoryOffsetSelection(Context);