[CodeGen] Match LLVM branch relaxation semantics

4 files changed, 87 insertions, 15 deletions

diff --git a/backend/lib/CodeGen/MachineBranchRelaxation.cpp b/backend/lib/CodeGen/MachineBranchRelaxation.cpp
index 9f14fc4..4df4cc8 100644
--- a/backend/lib/CodeGen/MachineBranchRelaxation.cpp
+++ b/backend/lib/CodeGen/MachineBranchRelaxation.cpp
@@ -2,7 +2,6 @@
 
 #include "rcc/CodeGen/MachineBranchRelaxation.h"
 
-#include "rcc/CodeGen/LiveVariables.h"
 #include "rcc/CodeGen/MachineBasicBlock.h"
 #include "rcc/CodeGen/MachineFunction.h"
 #include "rcc/CodeGen/MachineFunctionCloner.h"
@@ -29,16 +28,12 @@ struct MachineLayout final {
 [[nodiscard]] const MachinePassContract &relaxationContract() {
   static const MachinePassContract Contract = [] {
     MachinePassContract Value{
-        .Name = "machine-branch-relaxation",
-        .Requires = {MachineFunctionProperty::Selected, MachineFunctionProperty::NoPHIs,
-                     MachineFunctionProperty::TiedOpsRewritten, MachineFunctionProperty::NoVRegs,
-                     MachineFunctionProperty::TracksLiveness,
-                     MachineFunctionProperty::FrameFinalized,
-                     MachineFunctionProperty::NoFrameIndices},
+        .Name = "branch-relaxation",
+        .Requires = {MachineFunctionProperty::TracksLiveness},
         .Sets = {},
         .Clears = {},
         .RequiredAnalyses = {},
-        .VerificationStage = PipelineStage::PostRAFinalized,
+        .VerificationStage = std::nullopt,
     };
     Value.ThreadSafety = PassThreadSafety::ThreadSafe;
     Value.DeclaredChanges = {MachinePassChange::Instructions, MachinePassChange::Operands};
@@ -107,7 +102,7 @@ MachineBranchRelaxationPass::run(MachineFunction &Function, const TargetDescript
                                  MachineFunctionAnalysisManager &Analyses) {
   if (!Function.properties().containsAll(contract().Requires)) {
     return makeUnexpected(ErrorCode::VerificationError,
-                          "branch relaxation requires finalized frame and tracked liveness");
+                          "branch relaxation requires tracked liveness");
   }
   auto Trial = cloneMachineFunction(
       Function, Target,
@@ -185,9 +180,6 @@ MachineBranchRelaxationPass::runImplementation(MachineFunction &Function,
     if (Changed)
       continue;
     if (Result.Changed) {
-      auto Rebuilt = rebuildMachineLiveness(Function, Target);
-      if (!Rebuilt)
-        return std::unexpected(Rebuilt.error());
       Result.Changes = {MachinePassChange::Instructions, MachinePassChange::Operands};
       Result.Preserved = PreservedAnalyses::none();
     } else {
diff --git a/backend/lib/TableGen/Parser.cpp b/backend/lib/TableGen/Parser.cpp
index a14385c..66d10e0 100644
--- a/backend/lib/TableGen/Parser.cpp
+++ b/backend/lib/TableGen/Parser.cpp
@@ -1027,8 +1027,8 @@ private:
                      {"id", "symbol", "name", "assembly_mnemonic", "encoding", "operands",
                       "num_explicit_defs", "variadic", "implicit_defs", "implicit_uses",
                       "allow_extra_implicit_operands", "flags", "scheduling_class",
-                      "encoding_category", "assembly_format", "requires_memory_operands", "owner",
-                      "deadline"},
+                      "encoding_category", "size", "assembly_format", "requires_memory_operands",
+                      "owner", "deadline"},
                      "instruction")) {
         return false;
       }
@@ -1080,6 +1080,12 @@ private:
           return false;
         Instruction.EncodingCategorySymbol = *Value;
       }
+      if (const yaml::Node *Size = optionalNode(*Element, "size")) {
+        const auto Value = unsignedInteger32(*Size, "instruction size");
+        if (!Value)
+          return false;
+        Instruction.SizeInBytes = *Value;
+      }
       if (const yaml::Node *AssemblyFormat = optionalNode(*Element, "assembly_format")) {
         const auto Value = scalar(*AssemblyFormat, "instruction assembly_format");
         if (!Value)
diff --git a/backend/lib/Target/RISCV/RISCV.rccdesc b/backend/lib/Target/RISCV/RISCV.rccdesc
index c4a87e9..910ea3f 100644
--- a/backend/lib/Target/RISCV/RISCV.rccdesc
+++ b/backend/lib/Target/RISCV/RISCV.rccdesc
@@ -3962,6 +3962,7 @@ instructions:
     implicit_defs: [{ register: X1 }]
     allow_extra_implicit_operands: true
     flags: [pseudo, produces_code, allowed_pre_legalizer, call, may_load, may_store, has_side_effects, not_duplicable, unmodeled_side_effects]
+    size: 8
     owner: post_ra_expansion
     deadline: post_ra_finalized
   -
@@ -3974,6 +3975,7 @@ instructions:
     num_explicit_defs: 0
     allow_extra_implicit_operands: true
     flags: [pseudo, produces_code, allowed_pre_legalizer, terminator, branch, call, barrier, may_load, may_store, has_side_effects, not_duplicable, unmodeled_side_effects]
+    size: 8
     owner: post_ra_expansion
     deadline: post_ra_finalized
   -
@@ -3984,6 +3986,7 @@ instructions:
     num_explicit_defs: 0
     allow_extra_implicit_operands: true
     flags: [pseudo, produces_code, allowed_pre_legalizer, terminator, return, barrier]
+    size: 4
     owner: post_ra_expansion
     deadline: post_ra_finalized
   -
@@ -4044,5 +4047,6 @@ instructions:
       - { kinds: [basic_block], name: TARGET }
     num_explicit_defs: 0
     flags: [pseudo, produces_code, allowed_pre_legalizer, terminator, branch, barrier]
+    size: 4
     owner: post_ra_expansion
     deadline: post_ra_finalized
diff --git a/backend/lib/Target/RISCV/RISCVInstrInfo.cpp b/backend/lib/Target/RISCV/RISCVInstrInfo.cpp
index 7706cf3..886bf2b 100644
--- a/backend/lib/Target/RISCV/RISCVInstrInfo.cpp
+++ b/backend/lib/Target/RISCV/RISCVInstrInfo.cpp
@@ -2,6 +2,7 @@
 
 #include "RISCVInstrInfo.h"
 
+#include "rcc/CodeGen/GenericOpcodes.h"
 #include "rcc/CodeGen/InstructionDescriptor.h"
 #include "rcc/CodeGen/MachineBasicBlock.h"
 #include "rcc/CodeGen/MachineFunction.h"
@@ -16,6 +17,7 @@
 
 #include <algorithm>
 #include <array>
+#include <charconv>
 #include <cstddef>
 #include <cstdint>
 #include <limits>
@@ -30,6 +32,57 @@ namespace {
 constexpr std::uint32_t PCRelativeHighFlag = 2;
 constexpr std::uint32_t PCRelativeLowFlag = 4;
 
+[[nodiscard]] bool isAsmSpace(char Value) noexcept {
+  return Value == ' ' || Value == '\t' || Value == '\r' || Value == '\f' || Value == '\v';
+}
+
+[[nodiscard]] std::string_view trimAsmSpace(std::string_view Text) noexcept {
+  while (!Text.empty() && isAsmSpace(Text.front()))
+    Text.remove_prefix(1);
+  while (!Text.empty() && isAsmSpace(Text.back()))
+    Text.remove_suffix(1);
+  return Text;
+}
+
+[[nodiscard]] Expected<std::uint32_t> inlineAsmSize(const MachineInstr &Instruction) {
+  const MachineOperand *Operand = Instruction.operandAt(0);
+  const auto *Symbol = Operand == nullptr ? nullptr : Operand->getIf<ExternalSymbolOperand>();
+  if (Symbol == nullptr)
+    return makeUnexpected(ErrorCode::VerificationError, "RISC-V INLINEASM has no assembly text");
+
+  std::uint64_t Size = 0;
+  std::string_view Remaining = Symbol->Name;
+  while (!Remaining.empty()) {
+    const std::size_t Separator = Remaining.find_first_of("\n;");
+    std::string_view Statement = trimAsmSpace(Remaining.substr(0, Separator));
+    Remaining =
+        Separator == std::string_view::npos ? std::string_view{} : Remaining.substr(Separator + 1U);
+    if (Statement.empty() || Statement.front() == '#')
+      continue;
+
+    std::uint64_t StatementSize = 4;
+    constexpr std::string_view SpaceDirective = ".space";
+    if (Statement.starts_with(SpaceDirective) && (Statement.size() == SpaceDirective.size() ||
+                                                  isAsmSpace(Statement[SpaceDirective.size()]))) {
+      std::string_view Argument = trimAsmSpace(Statement.substr(SpaceDirective.size()));
+      std::int64_t Count = 0;
+      const auto Parsed =
+          std::from_chars(Argument.data(), Argument.data() + Argument.size(), Count);
+      if (Parsed.ec == std::errc{}) {
+        std::string_view Trailing(
+            Parsed.ptr, static_cast<std::size_t>(Argument.data() + Argument.size() - Parsed.ptr));
+        Trailing = trimAsmSpace(Trailing);
+        if (Trailing.empty() || Trailing.front() == '#')
+          StatementSize = Count < 0 ? 0 : static_cast<std::uint64_t>(Count);
+      }
+    }
+    if (StatementSize > std::numeric_limits<std::uint32_t>::max() - Size)
+      return makeUnexpected(ErrorCode::VerificationError, "RISC-V INLINEASM size overflowed");
+    Size += StatementSize;
+  }
+  return static_cast<std::uint32_t>(Size);
+}
+
 [[nodiscard]] bool isConditionalBranch(Opcode Value) noexcept {
   return Value == instructionOpcode(InstructionOpcode::BEQ) ||
          Value == instructionOpcode(InstructionOpcode::BNE) ||
@@ -333,6 +386,18 @@ selectLongBranchScratch(const MachineInstr &Instruction, const MachineBasicBlock
   auto TrueTarget = directTarget(Instruction);
   if (!TrueTarget)
     return std::unexpected(TrueTarget.error());
+  MachineInstr *Following = Instruction.next();
+  if (Following != nullptr && isUnconditionalBranch(*Following)) {
+    auto FalseTarget = directTarget(*Following);
+    if (!FalseTarget)
+      return std::unexpected(FalseTarget.error());
+    if (*FalseTarget == *TrueTarget) {
+      auto Removed = removeInstruction(Instruction);
+      if (!Removed)
+        return std::unexpected(Removed.error());
+      return true;
+    }
+  }
   MachineBasicBlock *FalseTarget = Block->next();
   if (FalseTarget == nullptr || FalseTarget == *TrueTarget ||
       std::ranges::none_of(Block->successors(), [FalseTarget](const MachineSuccessor &Successor) {
@@ -353,7 +418,6 @@ selectLongBranchScratch(const MachineInstr &Instruction, const MachineBasicBlock
   if (!SetTarget)
     return std::unexpected(SetTarget.error());
 
-  MachineInstr *Following = Instruction.next();
   if (Following != nullptr && isUnconditionalBranch(*Following)) {
     auto Retargeted = setDirectTarget(*Following, **TrueTarget);
     if (!Retargeted)
@@ -702,6 +766,12 @@ public:
                             "RISC-V instruction-size query encountered an unknown opcode");
     if (!Descriptor->Flags.has(InstructionFlag::ProducesCode))
       return 0;
+    if (Instruction.opcode() == opcode(GenericOpcode::INLINEASM) ||
+        Instruction.opcode() == opcode(GenericOpcode::INLINEASM_BR)) {
+      return inlineAsmSize(Instruction);
+    }
+    if (Descriptor->SizeInBytes != 0)
+      return Descriptor->SizeInBytes;
     const auto Category =
         std::ranges::find_if(encodingCategories(), [Descriptor](const EncodingCategoryInfo &Info) {
           return Info.ID == Descriptor->EncodingCategory;