[RISCV][GlobalISel] Generate integer extension selections

5 files changed, 380 insertions, 187 deletions

diff --git a/backend/lib/CodeGen/GlobalISel/GIMatchTableExecutor.cpp b/backend/lib/CodeGen/GlobalISel/GIMatchTableExecutor.cpp
index 97b6cfe..739f12b 100644
--- a/backend/lib/CodeGen/GlobalISel/GIMatchTableExecutor.cpp
+++ b/backend/lib/CodeGen/GlobalISel/GIMatchTableExecutor.cpp
@@ -146,21 +146,15 @@ operandOrDefConstantValue(const MachineInstr &Instruction, const MachineFunction
   return operandDefConstantValue(Instruction, Function, OperandIndex);
 }
 
-[[nodiscard]] Expected<bool> matchesOperandDefConstantRanges(const GIMatchTable &Table,
-                                                             const MachineInstr &Instruction,
-                                                             const MachineFunction &Function,
-                                                             const GIMatchTableInstruction &Step) {
-  const std::size_t First = Step.Value;
-  const std::size_t Count = Step.Auxiliary;
+[[nodiscard]] Expected<bool> matchesConstantRanges(const GIMatchTable &Table, std::size_t First,
+                                                   std::size_t Count,
+                                                   std::optional<std::int64_t> Value) {
   if (Count == 0 || First > Table.ConstantRanges.size() ||
       Count > Table.ConstantRanges.size() - First) {
     return makeUnexpected(ErrorCode::VerificationError,
                           "generated match table has an invalid constant range");
   }
-  auto Value = operandDefConstantValue(Instruction, Function, Step.OperandIndex);
   if (!Value)
-    return std::unexpected(Value.error());
-  if (!*Value)
     return false;
   for (std::size_t Index = First; Index < First + Count; ++Index) {
     const GIMatchTableConstantRange &Range = Table.ConstantRanges[Index];
@@ -168,12 +162,31 @@ operandOrDefConstantValue(const MachineInstr &Instruction, const MachineFunction
       return makeUnexpected(ErrorCode::VerificationError,
                             "generated match table has an inverted constant range");
     }
-    if (**Value >= Range.Minimum && **Value <= Range.Maximum)
+    if (*Value >= Range.Minimum && *Value <= Range.Maximum)
       return true;
   }
   return false;
 }
 
+[[nodiscard]] Expected<bool> matchesOperandImmediateRanges(const GIMatchTable &Table,
+                                                           const MachineInstr &Instruction,
+                                                           const GIMatchTableInstruction &Step) {
+  const MachineOperand *Operand = Instruction.operandAt(Step.OperandIndex);
+  const std::optional<std::int64_t> Value =
+      Operand == nullptr ? std::optional<std::int64_t>{} : immediateValue(*Operand);
+  return matchesConstantRanges(Table, Step.Value, Step.Auxiliary, Value);
+}
+
+[[nodiscard]] Expected<bool> matchesOperandDefConstantRanges(const GIMatchTable &Table,
+                                                             const MachineInstr &Instruction,
+                                                             const MachineFunction &Function,
+                                                             const GIMatchTableInstruction &Step) {
+  auto Value = operandDefConstantValue(Instruction, Function, Step.OperandIndex);
+  if (!Value)
+    return std::unexpected(Value.error());
+  return matchesConstantRanges(Table, Step.Value, Step.Auxiliary, *Value);
+}
+
 [[nodiscard]] const MachineInstr *operandDefinition(const MachineInstr &Instruction,
                                                     const MachineFunction &Function,
                                                     std::size_t OperandIndex) {
@@ -201,30 +214,14 @@ matchesOperandDefOperandConstantRanges(const GIMatchTable &Table, const MachineI
     return makeUnexpected(ErrorCode::VerificationError,
                           "generated match table has an invalid def-use constant range count");
   }
-  const std::size_t First = Step.Value;
   const std::size_t Count = static_cast<std::size_t>(Step.Immediate);
-  if (First > Table.ConstantRanges.size() || Count > Table.ConstantRanges.size() - First) {
-    return makeUnexpected(ErrorCode::VerificationError,
-                          "generated match table has an invalid def-use constant range");
-  }
   const MachineInstr *Definition = operandDefinition(Instruction, Function, Step.OperandIndex);
   if (Definition == nullptr)
     return false;
   auto Value = operandOrDefConstantValue(*Definition, Function, Step.Auxiliary);
   if (!Value)
     return std::unexpected(Value.error());
-  if (!*Value)
-    return false;
-  for (std::size_t Index = First; Index < First + Count; ++Index) {
-    const GIMatchTableConstantRange &Range = Table.ConstantRanges[Index];
-    if (Range.Minimum > Range.Maximum) {
-      return makeUnexpected(ErrorCode::VerificationError,
-                            "generated match table has an inverted def-use constant range");
-    }
-    if (**Value >= Range.Minimum && **Value <= Range.Maximum)
-      return true;
-  }
-  return false;
+  return matchesConstantRanges(Table, Step.Value, Count, *Value);
 }
 
 [[nodiscard]] Expected<RegisterOperandFlags> registerFlagsFromRaw(std::uint32_t Raw) {
@@ -303,10 +300,11 @@ findCapturedOperand(const std::vector<MatchedRule::CapturedOperand> &Captures, s
   return Position == Captures.end() ? nullptr : &*Position;
 }
 
-[[nodiscard]] Expected<std::optional<MatchedRule>> matchRule(const GIMatchTable &Table,
-                                                             const GIMatchTableRule &Rule,
-                                                             const MachineInstr &Instruction,
-                                                             const MachineFunction &Function) {
+template <typename TargetPredicateMatcher>
+[[nodiscard]] Expected<std::optional<MatchedRule>>
+matchRule(const GIMatchTable &Table, const GIMatchTableRule &Rule, const MachineInstr &Instruction,
+          const MachineFunction &Function, const TargetDescription &Target,
+          TargetPredicateMatcher &&MatchesTargetPredicate) {
   const std::size_t First = Rule.FirstInstruction;
   const std::size_t Count = Rule.InstructionCount;
   if (First > Table.Instructions.size() || Count > Table.Instructions.size() - First) {
@@ -376,6 +374,30 @@ findCapturedOperand(const std::vector<MatchedRule::CapturedOperand> &Captures, s
         return std::optional<MatchedRule>{};
       break;
     }
+    case GIMatchTableOpcode::CheckTargetPredicate: {
+      if (Applying || Step.Symbol.empty()) {
+        return makeUnexpected(ErrorCode::VerificationError,
+                              "generated match table has an invalid target predicate");
+      }
+      auto Matched = MatchesTargetPredicate(Step.Symbol, Instruction, Function, Target);
+      if (!Matched)
+        return std::unexpected(Matched.error());
+      if (!*Matched)
+        return std::optional<MatchedRule>{};
+      break;
+    }
+    case GIMatchTableOpcode::CheckOperandImmediateRanges: {
+      if (Applying) {
+        return makeUnexpected(ErrorCode::VerificationError,
+                              "generated match table performs a predicate after an action");
+      }
+      auto Matched = matchesOperandImmediateRanges(Table, Instruction, Step);
+      if (!Matched)
+        return std::unexpected(Matched.error());
+      if (!*Matched)
+        return std::optional<MatchedRule>{};
+      break;
+    }
     case GIMatchTableOpcode::CheckOperandDefConstantRanges: {
       if (Applying) {
         return makeUnexpected(ErrorCode::VerificationError,
@@ -1175,6 +1197,15 @@ Expected<void> GIMatchTableExecutor::finalizeGeneratedInstruction(MachineInstr &
   return {};
 }
 
+Expected<bool> GIMatchTableExecutor::matchesTargetPredicate(std::string_view Name,
+                                                            const MachineInstr &,
+                                                            const MachineFunction &,
+                                                            const TargetDescription &) const {
+  return makeUnexpected(ErrorCode::VerificationError,
+                        "generated match table uses unknown target predicate '" +
+                            std::string(Name) + "'");
+}
+
 Expected<std::optional<GIMatchComplexOperand>>
 GIMatchTableExecutor::renderComplexOperand(std::string_view, const MachineInstr &,
                                            const MachineFunction &, const TargetDescription &,
@@ -1187,7 +1218,12 @@ Expected<bool> GIMatchTableExecutor::match(const GIMatchTable &Table,
                                            const MachineFunction &Function,
                                            const TargetDescription &Target) const {
   for (const GIMatchTableRule &Rule : Table.Rules) {
-    auto Match = matchRule(Table, Rule, Instruction, Function);
+    auto Match = matchRule(Table, Rule, Instruction, Function, Target,
+                           [this](std::string_view Name, const MachineInstr &Input,
+                                  const MachineFunction &InputFunction,
+                                  const TargetDescription &InputTarget) {
+                             return matchesTargetPredicate(Name, Input, InputFunction, InputTarget);
+                           });
     if (!Match)
       return std::unexpected(Match.error());
     if (!*Match)
@@ -1218,7 +1254,12 @@ Expected<bool> GIMatchTableExecutor::select(const GIMatchTable &Table, MachineIn
                                             MachineFunction &Function,
                                             const TargetDescription &Target) const {
   for (const GIMatchTableRule &Rule : Table.Rules) {
-    auto Match = matchRule(Table, Rule, Instruction, Function);
+    auto Match = matchRule(Table, Rule, Instruction, Function, Target,
+                           [this](std::string_view Name, const MachineInstr &Input,
+                                  const MachineFunction &InputFunction,
+                                  const TargetDescription &InputTarget) {
+                             return matchesTargetPredicate(Name, Input, InputFunction, InputTarget);
+                           });
     if (!Match)
       return std::unexpected(Match.error());
     if (!*Match)
diff --git a/backend/lib/TableGen/Parser.cpp b/backend/lib/TableGen/Parser.cpp
index b4ff73c..7c6a8de 100644
--- a/backend/lib/TableGen/Parser.cpp
+++ b/backend/lib/TableGen/Parser.cpp
@@ -1173,6 +1173,92 @@ private:
     return true;
   }
 
+  [[nodiscard]] bool
+  parseSelectionTargetPredicates(const yaml::Node &Node,
+                                 std::vector<SourceSelectionTargetPredicate> &Result) {
+    if (!requireSequence(&Node, "selection target predicates"))
+      return false;
+    for (const yaml::NodeID ElementID : Node.sequence()) {
+      const yaml::Node *Element = Input.node(ElementID);
+      if (Element == nullptr)
+        return false;
+      const auto Symbol = scalar(*Element, "selection target predicate");
+      if (!Symbol)
+        return false;
+      Result.push_back({
+          .Symbol = *Symbol,
+          .Range = Element->range(),
+      });
+    }
+    return true;
+  }
+
+  [[nodiscard]] bool parseSelectionConstantRanges(const yaml::Node &Node,
+                                                  std::vector<SourceSelectionConstantRange> &Result,
+                                                  std::string_view Context) {
+    if (!requireSequence(&Node, Context) || Node.sequence().empty())
+      return false;
+    for (const yaml::NodeID RangeID : Node.sequence()) {
+      const yaml::Node *Range = Input.node(RangeID);
+      const std::string RangeContext = std::string(Context) + " range";
+      if (!requireMapping(Range, RangeContext) ||
+          !checkKeys(*Range, {"min", "max"}, RangeContext)) {
+        return false;
+      }
+      const yaml::Node *MinimumNode = requiredNode(*Range, "min", RangeContext);
+      const yaml::Node *MaximumNode = requiredNode(*Range, "max", RangeContext);
+      if (MinimumNode == nullptr || MaximumNode == nullptr)
+        return false;
+      const auto Minimum = signedInteger(*MinimumNode, "selection constant range minimum");
+      const auto Maximum = signedInteger(*MaximumNode, "selection constant range maximum");
+      if (!Minimum || !Maximum || *Minimum > *Maximum) {
+        if (Minimum && Maximum)
+          fail(Range->range(), "selection constant range is invalid");
+        return false;
+      }
+      Result.push_back({
+          .Minimum = *Minimum,
+          .Maximum = *Maximum,
+          .SourceRange = Range->range(),
+      });
+    }
+    return true;
+  }
+
+  [[nodiscard]] bool
+  parseSelectionOperandImmediates(const yaml::Node &Node,
+                                  std::vector<SourceSelectionOperandImmediatePredicate> &Result) {
+    if (!requireSequence(&Node, "selection operand immediates"))
+      return false;
+    for (const yaml::NodeID ElementID : Node.sequence()) {
+      const yaml::Node *Element = Input.node(ElementID);
+      if (!requireMapping(Element, "selection operand immediate") ||
+          !checkKeys(*Element, {"index", "ranges"}, "selection operand immediate")) {
+        return false;
+      }
+      const yaml::Node *IndexNode = requiredNode(*Element, "index", "selection operand immediate");
+      const yaml::Node *RangesNode =
+          requiredNode(*Element, "ranges", "selection operand immediate");
+      if (IndexNode == nullptr || RangesNode == nullptr)
+        return false;
+      const auto Index = unsignedInteger32(*IndexNode, "selection operand immediate index");
+      if (!Index)
+        return false;
+
+      SourceSelectionOperandImmediatePredicate Predicate{
+          .OperandIndex = *Index,
+          .Ranges = {},
+          .Range = Element->range(),
+      };
+      if (!parseSelectionConstantRanges(*RangesNode, Predicate.Ranges,
+                                        "selection operand immediate ranges")) {
+        return false;
+      }
+      Result.push_back(std::move(Predicate));
+    }
+    return true;
+  }
+
   [[nodiscard]] bool parseSelectionOperandDefConstants(
       const yaml::Node &Node, std::vector<SourceSelectionOperandDefConstantPredicate> &Result) {
     if (!requireSequence(&Node, "selection operand definition constants"))
@@ -1191,8 +1277,7 @@ private:
         return false;
       const auto Index =
           unsignedInteger32(*IndexNode, "selection operand definition constant index");
-      if (!Index || !requireSequence(RangesNode, "selection operand definition constant ranges") ||
-          RangesNode->sequence().empty()) {
+      if (!Index) {
         return false;
       }
 
@@ -1201,30 +1286,9 @@ private:
           .Ranges = {},
           .Range = Element->range(),
       };
-      for (const yaml::NodeID RangeID : RangesNode->sequence()) {
-        const yaml::Node *Range = Input.node(RangeID);
-        if (!requireMapping(Range, "selection operand definition constant range") ||
-            !checkKeys(*Range, {"min", "max"}, "selection operand definition constant range")) {
-          return false;
-        }
-        const yaml::Node *MinimumNode =
-            requiredNode(*Range, "min", "selection operand definition constant range");
-        const yaml::Node *MaximumNode =
-            requiredNode(*Range, "max", "selection operand definition constant range");
-        if (MinimumNode == nullptr || MaximumNode == nullptr)
-          return false;
-        const auto Minimum = signedInteger(*MinimumNode, "selection constant range minimum");
-        const auto Maximum = signedInteger(*MaximumNode, "selection constant range maximum");
-        if (!Minimum || !Maximum || *Minimum > *Maximum) {
-          if (Minimum && Maximum)
-            fail(Range->range(), "selection operand definition constant range is invalid");
-          return false;
-        }
-        Predicate.Ranges.push_back({
-            .Minimum = *Minimum,
-            .Maximum = *Maximum,
-            .SourceRange = Range->range(),
-        });
+      if (!parseSelectionConstantRanges(*RangesNode, Predicate.Ranges,
+                                        "selection operand definition constant ranges")) {
+        return false;
       }
       Result.push_back(std::move(Predicate));
     }
@@ -1472,8 +1536,9 @@ private:
   [[nodiscard]] bool parseSelectionRuleMatch(const yaml::Node &Node, SourceSelectionRule &Result) {
     if (!requireMapping(&Node, "selection rule match") ||
         !checkKeys(Node,
-                   {"opcode", "result_type", "integer_predicate", "operand_types", "operand_banks",
-                    "operand_classes", "operand_def_constants", "def_use"},
+                   {"opcode", "result_type", "integer_predicate", "target_predicates",
+                    "operand_types", "operand_banks", "operand_classes", "operand_immediates",
+                    "operand_def_constants", "def_use"},
                    "selection rule match")) {
       return false;
     }
@@ -1491,6 +1556,10 @@ private:
       if (!parseSelectionIntegerPredicate(*Predicate, Result))
         return false;
     }
+    if (const yaml::Node *TargetPredicates = optionalNode(Node, "target_predicates")) {
+      if (!parseSelectionTargetPredicates(*TargetPredicates, Result.TargetPredicates))
+        return false;
+    }
     if (const yaml::Node *OperandTypes = optionalNode(Node, "operand_types")) {
       if (!parseSelectionOperandTypes(*OperandTypes, Result.OperandTypes))
         return false;
@@ -1503,6 +1572,10 @@ private:
       if (!parseSelectionOperandClasses(*OperandClasses, Result.OperandClasses))
         return false;
     }
+    if (const yaml::Node *OperandImmediates = optionalNode(Node, "operand_immediates")) {
+      if (!parseSelectionOperandImmediates(*OperandImmediates, Result.OperandImmediates))
+        return false;
+    }
     if (const yaml::Node *OperandDefConstants = optionalNode(Node, "operand_def_constants")) {
       if (!parseSelectionOperandDefConstants(*OperandDefConstants, Result.OperandDefConstants))
         return false;
diff --git a/backend/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp b/backend/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp
index 64cbf47..898b366 100644
--- a/backend/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp
+++ b/backend/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp
@@ -147,6 +147,16 @@ public:
                            const RISCVRegisterBankInfo &RegisterBanks)
       : Subtarget(Subtarget), RegisterBanks(RegisterBanks) {}
 
+  [[nodiscard]] Expected<bool> matchesTargetPredicate(std::string_view Name, const MachineInstr &,
+                                                      const MachineFunction &,
+                                                      const TargetDescription &) const override {
+    if (Name == "zba")
+      return Subtarget.features().has(RISCVExtension::Zba);
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "RISC-V generated GlobalISel rule uses unknown target predicate '" +
+                              std::string(Name) + "'");
+  }
+
   [[nodiscard]] Expected<void>
   constrainSelectedInstruction(MachineInstr &Instruction, MachineFunction &Function,
                                const TargetDescription &Target) const override {
@@ -375,9 +385,6 @@ public:
     case GenericOpcode::G_PTRTOINT:
     case GenericOpcode::G_INTTOPTR:
     case GenericOpcode::G_TRUNC:
-    case GenericOpcode::G_SEXT:
-    case GenericOpcode::G_SEXT_INREG:
-    case GenericOpcode::G_ZEXT:
     case GenericOpcode::G_FCMP:
     case GenericOpcode::G_SMULH:
     case GenericOpcode::G_UMULH:
@@ -456,13 +463,9 @@ public:
     case GenericOpcode::G_INTTOPTR:
       return Instruction.setOpcode(opcode(GenericOpcode::COPY));
     case GenericOpcode::G_ANYEXT:
-    case GenericOpcode::G_SEXT:
-    case GenericOpcode::G_ZEXT:
-      return selectIntegerExtension(Instruction, Function, Target, Generic, Type);
+      return Instruction.setOpcode(opcode(GenericOpcode::COPY));
     case GenericOpcode::G_TRUNC:
       return Instruction.setOpcode(opcode(GenericOpcode::COPY));
-    case GenericOpcode::G_SEXT_INREG:
-      return selectSignExtendInRegister(Instruction, Function, Target, Type);
     case GenericOpcode::G_FCMP:
       return selectFloatingCompare(Instruction, Function, Target);
     case GenericOpcode::G_SELECT:
@@ -890,7 +893,6 @@ private:
     return {};
   }
 
-
   [[nodiscard]] Expected<void> selectAtomicRMW(MachineInstr &Instruction, MachineFunction &Function,
                                                const TargetDescription &Target,
                                                std::optional<LowLevelType> Type,
@@ -975,123 +977,6 @@ private:
     return Instruction.setOpcode(instructionOpcode(Selected));
   }
 
-  [[nodiscard]] Expected<void> selectIntegerExtension(MachineInstr &Instruction,
-                                                      MachineFunction &Function,
-                                                      const TargetDescription &Target,
-                                                      GenericOpcode Generic,
-                                                      std::optional<LowLevelType> Type) const {
-    const std::optional<LowLevelType> SourceType = operandType(Function, Instruction, 1);
-    const MachineOperand *DestinationOperand = Instruction.operandAt(0);
-    const MachineOperand *SourceOperand = Instruction.operandAt(1);
-    const Register *Destination =
-        DestinationOperand == nullptr ? nullptr : DestinationOperand->getIf<Register>();
-    const Register *Source = SourceOperand == nullptr ? nullptr : SourceOperand->getIf<Register>();
-    if (!Type || !SourceType || Destination == nullptr || Source == nullptr)
-      return makeUnexpected(ErrorCode::VerificationError,
-                            "RISC-V integer extension has invalid register operands");
-    const Register DestinationRegister = *Destination;
-    if (Generic == GenericOpcode::G_ANYEXT)
-      return Instruction.setOpcode(opcode(GenericOpcode::COPY));
-
-    const std::uint32_t DestinationBits = Type->scalarBitWidth();
-    const std::uint32_t SourceBits = SourceType->scalarBitWidth();
-    if (Generic == GenericOpcode::G_SEXT && DestinationBits == 64 && SourceBits == 32) {
-      auto ImmediateResult = appendImmediate(Instruction, 0);
-      if (!ImmediateResult)
-        return std::unexpected(ImmediateResult.error());
-      return Instruction.setOpcode(instructionOpcode(InstructionOpcode::ADDIW));
-    }
-
-    const std::uint32_t ContainerBits = DestinationBits <= 32 ? 32 : 64;
-    const std::int64_t Shift = static_cast<std::int64_t>(ContainerBits - SourceBits);
-    const InstructionOpcode ShiftLeft =
-        ContainerBits == 32 ? InstructionOpcode::SLLIW : InstructionOpcode::SLLI;
-    const InstructionOpcode ShiftRight =
-        Generic == GenericOpcode::G_SEXT
-            ? (ContainerBits == 32 ? InstructionOpcode::SRAIW : InstructionOpcode::SRAI)
-            : (ContainerBits == 32 ? InstructionOpcode::SRLIW : InstructionOpcode::SRLI);
-    auto Temporary = createTemporary(Function, *Type, Target);
-    if (!Temporary)
-      return std::unexpected(Temporary.error());
-    auto SetDestination = Instruction.operandAt(0)->setRegister(*Temporary);
-    if (!SetDestination)
-      return std::unexpected(SetDestination.error());
-    auto SetOpcode = Instruction.setOpcode(instructionOpcode(ShiftLeft));
-    if (!SetOpcode)
-      return std::unexpected(SetOpcode.error());
-    auto ShiftImmediate = appendImmediate(Instruction, Shift);
-    if (!ShiftImmediate)
-      return std::unexpected(ShiftImmediate.error());
-
-    MachineIRBuilder Builder(Function, Target);
-    auto Shifted = buildImmediateAfter(Builder, Instruction, ShiftRight, DestinationRegister,
-                                       *Temporary, Shift);
-    if (!Shifted)
-      return std::unexpected(Shifted.error());
-    return {};
-  }
-
-  [[nodiscard]] Expected<void> selectSignExtendInRegister(MachineInstr &Instruction,
-                                                          MachineFunction &Function,
-                                                          const TargetDescription &Target,
-                                                          std::optional<LowLevelType> Type) const {
-    const MachineOperand *DestinationOperand = Instruction.operandAt(0);
-    const MachineOperand *SourceOperand = Instruction.operandAt(1);
-    const Register *Destination =
-        DestinationOperand == nullptr ? nullptr : DestinationOperand->getIf<Register>();
-    const Register *Source = SourceOperand == nullptr ? nullptr : SourceOperand->getIf<Register>();
-    std::optional<std::int64_t> Width;
-    if (const MachineOperand *WidthOperand = Instruction.operandAt(2); WidthOperand != nullptr) {
-      auto ParsedWidth = signedImmediate(*WidthOperand);
-      if (ParsedWidth)
-        Width = *ParsedWidth;
-    }
-    if (!Type || Destination == nullptr || Source == nullptr || !Width || *Width <= 0 ||
-        static_cast<std::uint64_t>(*Width) > Type->scalarBitWidth())
-      return makeUnexpected(ErrorCode::VerificationError,
-                            "RISC-V sign-extension-in-register has invalid operands");
-    const Register DestinationRegister = *Destination;
-    if (static_cast<std::uint64_t>(*Width) == Type->scalarBitWidth()) {
-      auto EraseWidth = Instruction.eraseOperand(2);
-      if (!EraseWidth)
-        return std::unexpected(EraseWidth.error());
-      return Instruction.setOpcode(opcode(GenericOpcode::COPY));
-    }
-
-    const std::uint32_t TypeBits = Type->scalarBitWidth();
-    if (TypeBits == 64 && *Width == 32) {
-      auto ImmediateResult = replaceImmediate(Instruction, 2, 0);
-      if (!ImmediateResult)
-        return std::unexpected(ImmediateResult.error());
-      return Instruction.setOpcode(instructionOpcode(InstructionOpcode::ADDIW));
-    }
-    const std::uint32_t ContainerBits = TypeBits <= 32 ? 32 : 64;
-    const std::int64_t Shift = static_cast<std::int64_t>(ContainerBits) - *Width;
-    const InstructionOpcode ShiftLeft =
-        ContainerBits == 32 ? InstructionOpcode::SLLIW : InstructionOpcode::SLLI;
-    const InstructionOpcode ShiftRight =
-        ContainerBits == 32 ? InstructionOpcode::SRAIW : InstructionOpcode::SRAI;
-    auto Temporary = createTemporary(Function, *Type, Target);
-    if (!Temporary)
-      return std::unexpected(Temporary.error());
-    auto SetDestination = Instruction.operandAt(0)->setRegister(*Temporary);
-    if (!SetDestination)
-      return std::unexpected(SetDestination.error());
-    auto SetOpcode = Instruction.setOpcode(instructionOpcode(ShiftLeft));
-    if (!SetOpcode)
-      return std::unexpected(SetOpcode.error());
-    auto ShiftImmediate = replaceImmediate(Instruction, 2, Shift);
-    if (!ShiftImmediate)
-      return std::unexpected(ShiftImmediate.error());
-
-    MachineIRBuilder Builder(Function, Target);
-    auto Shifted = buildImmediateAfter(Builder, Instruction, ShiftRight, DestinationRegister,
-                                       *Temporary, Shift);
-    if (!Shifted)
-      return std::unexpected(Shifted.error());
-    return {};
-  }
-
   [[nodiscard]] Expected<void> selectFloatingCompare(MachineInstr &Instruction,
                                                      MachineFunction &Function,
                                                      const TargetDescription &Target) const {
diff --git a/backend/lib/Target/RISCV/RISCVGISelInteger.rccdesc b/backend/lib/Target/RISCV/RISCVGISelInteger.rccdesc
index e096993..1667f66 100644
--- a/backend/lib/Target/RISCV/RISCVGISelInteger.rccdesc
+++ b/backend/lib/Target/RISCV/RISCVGISelInteger.rccdesc
@@ -1,5 +1,198 @@
 selection_rules:
   -
+    id: select_g_sext_s64_s32
+    priority: 500
+    match:
+      opcode: G_SEXT
+      result_type: { kind: scalar, bits: 64 }
+      operand_types:
+        - { index: 1, kind: scalar, bits: 32 }
+      operand_banks:
+        - { index: 0, bank: GPR }
+        - { index: 1, bank: GPR }
+    apply:
+      instructions:
+        - opcode: ADDIW
+          root: true
+          operands: [{ operand: 0 }, { operand: 1 }, { immediate: 0 }]
+  -
+    id: select_g_sext_s64_s16
+    priority: 400
+    match:
+      opcode: G_SEXT
+      result_type: { kind: scalar, bits: 64 }
+      operand_types:
+        - { index: 1, kind: scalar, bits: 16 }
+      operand_banks:
+        - { index: 0, bank: GPR }
+        - { index: 1, bank: GPR }
+    apply:
+      temporaries:
+        - { id: extended, type_from_operand: 0 }
+      instructions:
+        - opcode: SLLI
+          root: true
+          operands:
+            - { temporary: extended, flags: [define] }
+            - { operand: 1 }
+            - { immediate: 48 }
+        - opcode: SRAI
+          root: false
+          operands:
+            - { operand: 0 }
+            - { temporary: extended, flags: [use] }
+            - { immediate: 48 }
+  -
+    id: select_g_sext_s32_s16
+    priority: 400
+    match:
+      opcode: G_SEXT
+      result_type: { kind: scalar, bits: 32 }
+      operand_types:
+        - { index: 1, kind: scalar, bits: 16 }
+      operand_banks:
+        - { index: 0, bank: GPR }
+        - { index: 1, bank: GPR }
+    apply:
+      temporaries:
+        - { id: extended, type_from_operand: 0 }
+      instructions:
+        - opcode: SLLIW
+          root: true
+          operands:
+            - { temporary: extended, flags: [define] }
+            - { operand: 1 }
+            - { immediate: 16 }
+        - opcode: SRAIW
+          root: false
+          operands:
+            - { operand: 0 }
+            - { temporary: extended, flags: [use] }
+            - { immediate: 16 }
+  -
+    id: select_g_zext_s64_s32_zba
+    priority: 500
+    match:
+      opcode: G_ZEXT
+      result_type: { kind: scalar, bits: 64 }
+      target_predicates: [zba]
+      operand_types:
+        - { index: 1, kind: scalar, bits: 32 }
+      operand_banks:
+        - { index: 0, bank: GPR }
+        - { index: 1, bank: GPR }
+    apply:
+      instructions:
+        - opcode: ADD_UW
+          root: true
+          operands:
+            - { operand: 0 }
+            - { operand: 1 }
+            - { physical_register: X0, flags: [use] }
+  -
+    id: select_g_zext_s64_s32
+    priority: 400
+    match:
+      opcode: G_ZEXT
+      result_type: { kind: scalar, bits: 64 }
+      operand_types:
+        - { index: 1, kind: scalar, bits: 32 }
+      operand_banks:
+        - { index: 0, bank: GPR }
+        - { index: 1, bank: GPR }
+    apply:
+      temporaries:
+        - { id: extended, type_from_operand: 0 }
+      instructions:
+        - opcode: SLLI
+          root: true
+          operands:
+            - { temporary: extended, flags: [define] }
+            - { operand: 1 }
+            - { immediate: 32 }
+        - opcode: SRLI
+          root: false
+          operands:
+            - { operand: 0 }
+            - { temporary: extended, flags: [use] }
+            - { immediate: 32 }
+  -
+    id: select_g_zext_s64_s16
+    priority: 400
+    match:
+      opcode: G_ZEXT
+      result_type: { kind: scalar, bits: 64 }
+      operand_types:
+        - { index: 1, kind: scalar, bits: 16 }
+      operand_banks:
+        - { index: 0, bank: GPR }
+        - { index: 1, bank: GPR }
+    apply:
+      temporaries:
+        - { id: extended, type_from_operand: 0 }
+      instructions:
+        - opcode: SLLI
+          root: true
+          operands:
+            - { temporary: extended, flags: [define] }
+            - { operand: 1 }
+            - { immediate: 48 }
+        - opcode: SRLI
+          root: false
+          operands:
+            - { operand: 0 }
+            - { temporary: extended, flags: [use] }
+            - { immediate: 48 }
+  -
+    id: select_g_zext_s32_s16
+    priority: 400
+    match:
+      opcode: G_ZEXT
+      result_type: { kind: scalar, bits: 32 }
+      operand_types:
+        - { index: 1, kind: scalar, bits: 16 }
+      operand_banks:
+        - { index: 0, bank: GPR }
+        - { index: 1, bank: GPR }
+    apply:
+      temporaries:
+        - { id: extended, type_from_operand: 0 }
+      instructions:
+        - opcode: SLLIW
+          root: true
+          operands:
+            - { temporary: extended, flags: [define] }
+            - { operand: 1 }
+            - { immediate: 16 }
+        - opcode: SRLIW
+          root: false
+          operands:
+            - { operand: 0 }
+            - { temporary: extended, flags: [use] }
+            - { immediate: 16 }
+  -
+    id: select_g_sext_inreg_s64_32
+    priority: 500
+    match:
+      opcode: G_SEXT_INREG
+      result_type: { kind: scalar, bits: 64 }
+      operand_types:
+        - { index: 1, kind: scalar, bits: 64 }
+      operand_banks:
+        - { index: 0, bank: GPR }
+        - { index: 1, bank: GPR }
+      operand_immediates:
+        - index: 2
+          ranges: [{ min: 32, max: 32 }]
+    apply:
+      instructions:
+        - opcode: ADDIW
+          root: true
+          operands:
+            - { operand: 0 }
+            - { operand: 1 }
+            - { immediate: 0 }
+  -
     id: select_g_ptr_add_immediate
     priority: 400
     match:
diff --git a/backend/lib/Target/RISCV/RISCVSubtarget.cpp b/backend/lib/Target/RISCV/RISCVSubtarget.cpp
index e38a574..b6370fd 100644
--- a/backend/lib/Target/RISCV/RISCVSubtarget.cpp
+++ b/backend/lib/Target/RISCV/RISCVSubtarget.cpp
@@ -32,6 +32,7 @@ constexpr std::array ExtensionNames{
     ExtensionName{"f", RISCVExtension::F},
     ExtensionName{"d", RISCVExtension::D},
     ExtensionName{"c", RISCVExtension::C},
+    ExtensionName{"zba", RISCVExtension::Zba},
     ExtensionName{"zbb", RISCVExtension::Zbb},
     ExtensionName{"zicsr", RISCVExtension::Zicsr},
     ExtensionName{"zifencei", RISCVExtension::Zifencei},