globalisel][aarch64] Make G_SEXT_INREG legal from all bit sizes to s32/s64

[dsanders]

Summary

I had hoped to demonstrate selectively marking certain bit sizes legal but
it turns out AArch64 supports them all via SBFM/UBFM. Added a unit test to
demonstrate that instead. Targets with more limited support for bit sizes
can use .legalForTypeWithImm({{s32, 8}, {s32, 16}}).lower() and similar to
keep G_SEXT_INREG for certain sizes and lower to shifts for others.

Depends on D61289

[dsanders]

Rebase onto updated D61289

[dsanders]

Rebased to master

[aemerson]

LGTM.

llvm/include/llvm/CodeGen/GlobalISel/LegalizationArtifactCombiner.h

bool tryCombineSExt(MachineInstr &MI,
     Builder.setInstr(MI);
     unsigned DstReg = MI.getOperand(0).getReg();
     unsigned SrcReg = lookThroughCopyInstrs(MI.getOperand(1).getReg());
 
     // sext(trunc x) - > (sext_inreg (aext/copy/trunc x), c)
     unsigned TruncSrc;
     if (mi_match(SrcReg, MRI, m_GTrunc(m_Reg(TruncSrc)))) {
       LLT DstTy = MRI.getType(DstReg);
-      if (isInstUnsupported({TargetOpcode::G_SEXT_INREG, {DstTy}}))
+      LLT SrcTy = MRI.getType(SrcReg);
+      if (isInstUnsupported({TargetOpcode::G_SEXT_INREG,
+                             {DstTy},
+                             {},
+                             {SrcTy.getScalarSizeInBits()}}))
         return false;
       LLVM_DEBUG(dbgs() << ".. Combine MI: " << MI;);
-      LLT SrcTy = MRI.getType(SrcReg);
       unsigned SizeInBits = SrcTy.getScalarSizeInBits();
       Builder.buildInstr(
           TargetOpcode::G_SEXT_INREG, {DstReg},
           {Builder.buildAnyExtOrTrunc(DstTy, TruncSrc), APInt(64, SizeInBits)});
       markInstAndDefDead(MI, *MRI.getVRegDef(SrcReg), DeadInsts);
       return true;
     }
     return tryFoldImplicitDef(MI, DeadInsts);

llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h

 
 /// The LegalityQuery object bundles together all the information that's needed
 /// to decide whether a given operation is legal or not.
 /// For efficiency, it doesn't make a copy of Types so care must be taken not
 /// to free it before using the query.
 struct LegalityQuery {
   unsigned Opcode;
   ArrayRef<LLT> Types;
+  ArrayRef<int64_t> Imms;
 
   struct MemDesc {
     uint64_t SizeInBits;
     uint64_t AlignInBits;
     AtomicOrdering Ordering;
   };
 
   /// Operations which require memory can use this to place requirements on the
   /// memory type for each MMO.
   ArrayRef<MemDesc> MMODescrs;
 
   constexpr LegalityQuery(unsigned Opcode, const ArrayRef<LLT> Types,
+                          const ArrayRef<MemDesc> MMODescrs,
+                          const ArrayRef<int64_t> Imms)
+      : Opcode(Opcode), Types(Types), Imms(Imms), MMODescrs(MMODescrs) {}
+  constexpr LegalityQuery(unsigned Opcode, const ArrayRef<LLT> Types,
                           const ArrayRef<MemDesc> MMODescrs)
-      : Opcode(Opcode), Types(Types), MMODescrs(MMODescrs) {}
+      : LegalityQuery(Opcode, Types, MMODescrs, {}) {}
   constexpr LegalityQuery(unsigned Opcode, const ArrayRef<LLT> Types)
-      : LegalityQuery(Opcode, Types, {}) {}
+      : LegalityQuery(Opcode, Types, {}, {}) {}
 
   raw_ostream &print(raw_ostream &OS) const;
 };
 
 /// The result of a query. It either indicates a final answer of Legal or
 /// Unsupported or describes an action that must be taken to make an operation
 /// more legal.
 struct LegalizeActionStep {

 LegalityPredicate
 typePairInSet(unsigned TypeIdx0, unsigned TypeIdx1,
               std::initializer_list<std::pair<LLT, LLT>> TypesInit);
 /// True iff the given types for the given pair of type indexes is one of the
 /// specified type pairs.
 LegalityPredicate typePairAndMemDescInSet(
     unsigned TypeIdx0, unsigned TypeIdx1, unsigned MMOIdx,
     std::initializer_list<TypePairAndMemDesc> TypesAndMemDescInit);
+/// True iff the given type and immediate is one of the specified pairs.
+LegalityPredicate
+typeAndImmInSet(unsigned TypeIdx0, unsigned ImmIdx,
+                std::initializer_list<std::pair<LLT, int64_t>> TypesAndImmInit);
 /// True iff the specified type index is a scalar.
 LegalityPredicate isScalar(unsigned TypeIdx);
 /// True iff the specified type index is a vector.
 LegalityPredicate isVector(unsigned TypeIdx);
 /// True iff the specified type index is a pointer (with any address space).
 LegalityPredicate isPointer(unsigned TypeIdx);
 /// True iff the specified type index is a pointer with the specified address
 /// space.
public:
   /// The instruction is legal when type indexes 0, 1, and 2 are both their
   /// respective lists.
   LegalizeRuleSet &legalForCartesianProduct(std::initializer_list<LLT> Types0,
                                             std::initializer_list<LLT> Types1,
                                             std::initializer_list<LLT> Types2) {
     return actionForCartesianProduct(LegalizeAction::Legal, Types0, Types1,
                                      Types2);
   }
+  /// The instruction is legal when type index 0 along with the immediate 0
+  /// match any given type and immediate tuple in the given list.
+  LegalizeRuleSet &legalForTypeWithImm(
+      std::initializer_list<std::pair<LLT, int64_t>> TypeAndImm) {
+    return actionIf(LegalizeAction::Legal,
+                    LegalityPredicates::typeAndImmInSet(
+                        typeIdx(0), /*ImmIdx*/ 0, TypeAndImm));
+  }
 
   LegalizeRuleSet &alwaysLegal() {
     using namespace LegalizeMutations;
     markAllTypeIdxsAsCovered();
     return actionIf(LegalizeAction::Legal, always);
   }
 
   /// The instruction is lowered.

llvm/lib/CodeGen/GlobalISel/LegalityPredicates.cpp

return [=](const LegalityQuery &Query) {
     return std::find_if(
       TypesAndMemDesc.begin(), TypesAndMemDesc.end(),
       [=](const TypePairAndMemDesc &Entry) ->bool {
         return Match.isCompatible(Entry);
       }) != TypesAndMemDesc.end();
   };
 }
 
+LegalityPredicate LegalityPredicates::typeAndImmInSet(
+    unsigned TypeIdx0, unsigned ImmIdx,
+    std::initializer_list<std::pair<LLT, int64_t>> TypeAndImmInit) {
+  SmallVector<std::pair<LLT, int64_t>, 4> TypeAndImm = TypeAndImmInit;
+  return [=](const LegalityQuery &Query) {
+    std::pair<LLT, int64_t> Match = {Query.Types[TypeIdx0], Query.Imms[ImmIdx]};
+    return std::find(TypeAndImm.begin(), TypeAndImm.end(), Match) !=
+           TypeAndImm.end();
+  };
+}
+
 LegalityPredicate LegalityPredicates::isScalar(unsigned TypeIdx) {
   return [=](const LegalityQuery &Query) {
     return Query.Types[TypeIdx].isScalar();
   };
 }
 
 LegalityPredicate LegalityPredicates::isVector(unsigned TypeIdx) {
   return [=](const LegalityQuery &Query) {

llvm/lib/CodeGen/GlobalISel/LegalizerInfo.cpp

raw_ostream &llvm::operator<<(raw_ostream &OS, LegalizeAction Action) {
   return OS;
 }
 
 raw_ostream &LegalityQuery::print(raw_ostream &OS) const {
   OS << Opcode << ", Tys={";
   for (const auto &Type : Types) {
     OS << Type << ", ";
   }
-  OS << "}, Opcode=";
+  OS << "}";
 
-  OS << Opcode << ", MMOs={";
+  OS << ", MMOs={";
   for (const auto &MMODescr : MMODescrs) {
     OS << MMODescr.SizeInBits << ", ";
   }
   OS << "}";
 
+  OS << ", Imms={";
+  for (const auto &Imm : Imms) {
+    OS << Imm << ", ";
+  }
+  OS << "}";
+
   return OS;
 }
 
 #ifndef NDEBUG
 // Make sure the rule won't (trivially) loop forever.
 static bool hasNoSimpleLoops(const LegalizeRule &Rule, const LegalityQuery &Q,
                              const std::pair<unsigned, LLT> &Mutation) {
   switch (Rule.getAction()) {
LegalizerInfo::getAction(const LegalityQuery &Query) const {
   LLVM_DEBUG(dbgs() << ".. (legacy) Legal\n");
   return {Legal, 0, LLT{}};
 }
 
 LegalizeActionStep
 LegalizerInfo::getAction(const MachineInstr &MI,
                          const MachineRegisterInfo &MRI) const {
   SmallVector<LLT, 2> Types;
+  SmallVector<int64_t, 2> Imms;
   SmallBitVector SeenTypes(8);
   const MCOperandInfo *OpInfo = MI.getDesc().OpInfo;
   // FIXME: probably we'll need to cache the results here somehow?
   for (unsigned i = 0; i < MI.getDesc().getNumOperands(); ++i) {
+    if (MI.getOperand(i).isImm())
+      Imms.push_back(MI.getOperand(i).getImm());
     if (!OpInfo[i].isGenericType())
       continue;
 
     // We must only record actions once for each TypeIdx; otherwise we'd
     // try to legalize operands multiple times down the line.
     unsigned TypeIdx = OpInfo[i].getGenericTypeIndex();
     if (SeenTypes[TypeIdx])
       continue;
 
     SeenTypes.set(TypeIdx);
 
     LLT Ty = getTypeFromTypeIdx(MI, MRI, i, TypeIdx);
     Types.push_back(Ty);
   }
 
   SmallVector<LegalityQuery::MemDesc, 2> MemDescrs;
   for (const auto &MMO : MI.memoperands())
     MemDescrs.push_back({8 * MMO->getSize() /* in bits */,
                          8 * MMO->getAlignment(),
                          MMO->getOrdering()});
 
-  return getAction({MI.getOpcode(), Types, MemDescrs});
+  return getAction({MI.getOpcode(), Types, MemDescrs, Imms});
 }
 
 bool LegalizerInfo::isLegal(const MachineInstr &MI,
                             const MachineRegisterInfo &MRI) const {
   return getAction(MI, MRI).Action == Legal;
 }
 
 bool LegalizerInfo::isLegalOrCustom(const MachineInstr &MI,

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

if (DstSize > 32) {
           .addUse(SrcReg)
           .addImm(AArch64::sub_32);
       I.getOperand(1).setReg(ExtSrc);
     }
     return selectCopy(I, TII, MRI, TRI, RBI);
   }
 
   case TargetOpcode::G_ZEXT:
-  case TargetOpcode::G_SEXT: {
+  case TargetOpcode::G_SEXT:
+  case TargetOpcode::G_SEXT_INREG: {
     unsigned Opcode = I.getOpcode();
-    const LLT DstTy = MRI.getType(I.getOperand(0).getReg()),
-              SrcTy = MRI.getType(I.getOperand(1).getReg());
-    const bool isSigned = Opcode == TargetOpcode::G_SEXT;
+    const LLT DstTy = MRI.getType(I.getOperand(0).getReg());
+    const bool isSigned =
+        Opcode == TargetOpcode::G_SEXT || Opcode == TargetOpcode::G_SEXT_INREG;
     const unsigned DefReg = I.getOperand(0).getReg();
     const unsigned SrcReg = I.getOperand(1).getReg();
     const RegisterBank &RB = *RBI.getRegBank(DefReg, MRI, TRI);
 
+    unsigned SrcSizeInBits =
+        Opcode == TargetOpcode::G_SEXT_INREG
+            ? I.getOperand(2).getImm()
+            : MRI.getType(I.getOperand(1).getReg()).getSizeInBits();
+
     if (RB.getID() != AArch64::GPRRegBankID) {
       LLVM_DEBUG(dbgs() << TII.getName(I.getOpcode()) << " on bank: " << RB
                         << ", expected: GPR\n");
       return false;
     }
 
     MachineInstr *ExtI;
     if (DstTy == LLT::scalar(64)) {
if (DstTy == LLT::scalar(64)) {
           .addUse(SrcReg)
           .addImm(AArch64::sub_32);
 
       const unsigned NewOpc = isSigned ? AArch64::SBFMXri : AArch64::UBFMXri;
       ExtI = BuildMI(MBB, I, I.getDebugLoc(), TII.get(NewOpc))
                  .addDef(DefReg)
                  .addUse(SrcXReg)
                  .addImm(0)
-                 .addImm(SrcTy.getSizeInBits() - 1);
+                 .addImm(SrcSizeInBits - 1);
     } else if (DstTy.isScalar() && DstTy.getSizeInBits() <= 32) {
       const unsigned NewOpc = isSigned ? AArch64::SBFMWri : AArch64::UBFMWri;
       ExtI = BuildMI(MBB, I, I.getDebugLoc(), TII.get(NewOpc))
                  .addDef(DefReg)
                  .addUse(SrcReg)
                  .addImm(0)
-                 .addImm(SrcTy.getSizeInBits() - 1);
+                 .addImm(SrcSizeInBits - 1);
     } else {
       return false;
     }
 
     constrainSelectedInstRegOperands(*ExtI, TII, TRI, RBI);
 
     I.eraseFromParent();
     return true;

llvm/lib/Target/AArch64/AArch64LegalizerInfo.cpp

getActionDefinitionsBuilder({G_ZEXT, G_SEXT, G_ANYEXT})
         if (SrcSize < 8 || !isPowerOf2_32(SrcSize))
           return false;
 
         return true;
       });
 
   getActionDefinitionsBuilder(G_TRUNC).alwaysLegal();
 
-  getActionDefinitionsBuilder(G_SEXT_INREG).lower();
+  getActionDefinitionsBuilder(G_SEXT_INREG)
+      // It happens that SBFM/UBFM can handle all possible extensions for each
+      // type so we don't need to check Query.Imm[0].
+      .legalFor({s32, s64})
+      .clampScalar(0, s32, s64)
+      .widenScalarToNextPow2(0)
+      .lower();
 
   // FP conversions
   getActionDefinitionsBuilder(G_FPTRUNC).legalFor(
       {{s16, s32}, {s16, s64}, {s32, s64}, {v4s16, v4s32}, {v2s32, v2s64}});
   getActionDefinitionsBuilder(G_FPEXT).legalFor(
       {{s32, s16}, {s64, s16}, {s64, s32}, {v4s32, v4s16}, {v2s64, v2s32}});
 
   // Conversions

llvm/test/CodeGen/AArch64/GlobalISel/legalize-div.mir

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -O0 -march=aarch64 -run-pass=legalizer %s -o - | FileCheck %s
 ---
 name:            test_div
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: test_div
     ; CHECK-DAG: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
     ; CHECK-DAG: [[COPY1:%[0-9]+]]:_(s64) = COPY $x1
-    ; CHECK-DAG: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 24
     ; CHECK-DAG: [[TRUNC:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
-    ; CHECK-DAG: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[TRUNC]], [[C]](s32)
-    ; CHECK-DAG: [[ASHR:%[0-9]+]]:_(s32) = G_ASHR [[SHL]], [[C]](s32)
+    ; CHECK-DAG: [[SEXT_INREG:%[0-9]+]]:_(s32) = G_SEXT_INREG [[TRUNC]], 8
     ; CHECK-DAG: [[TRUNC1:%[0-9]+]]:_(s32) = G_TRUNC [[COPY1]](s64)
-    ; CHECK-DAG: [[SHL1:%[0-9]+]]:_(s32) = G_SHL [[TRUNC1]], [[C]](s32)
-    ; CHECK-DAG: [[ASHR1:%[0-9]+]]:_(s32) = G_ASHR [[SHL1]], [[C]](s32)
-    ; CHECK-DAG: [[SDIV:%[0-9]+]]:_(s32) = G_SDIV [[ASHR]], [[ASHR1]]
+    ; CHECK-DAG: [[SEXT_INREG1:%[0-9]+]]:_(s32) = G_SEXT_INREG [[TRUNC1]], 8
+    ; CHECK-DAG: [[SDIV:%[0-9]+]]:_(s32) = G_SDIV [[SEXT_INREG]], [[SEXT_INREG1]]
     ; CHECK-DAG: [[COPY2:%[0-9]+]]:_(s32) = COPY [[SDIV]](s32)
     ; CHECK-DAG: $w0 = COPY [[COPY2]](s32)
     ; CHECK-DAG: [[C2:%[0-9]+]]:_(s32) = G_CONSTANT i32 255
     ; CHECK-DAG: [[TRUNC2:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
     ; CHECK-DAG: [[AND:%[0-9]+]]:_(s32) = G_AND [[TRUNC2]], [[C2]]
     ; CHECK-DAG: [[TRUNC3:%[0-9]+]]:_(s32) = G_TRUNC [[COPY1]](s64)
     ; CHECK-DAG: [[AND1:%[0-9]+]]:_(s32) = G_AND [[TRUNC3]], [[C2]]
     ; CHECK-DAG: [[UDIV:%[0-9]+]]:_(s32) = G_UDIV [[AND]], [[AND1]]

llvm/test/CodeGen/AArch64/GlobalISel/legalize-ext.mir

bb.0.entry:
     ; CHECK-DAG: [[COPY1:%[0-9]+]]:_(s64) = COPY [[COPY]](s64)
     ; CHECK-DAG: $x0 = COPY [[COPY1]](s64)
     ; CHECK-DAG: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 255
     ; CHECK-DAG: [[COPY2:%[0-9]+]]:_(s64) = COPY [[COPY]](s64)
     ; CHECK-DAG: [[AND:%[0-9]+]]:_(s64) = G_AND [[COPY2]], [[C]]
     ; CHECK-DAG: $x0 = COPY [[AND]](s64)
     ; CHECK-DAG: [[COPY3:%[0-9]+]]:_(s64) = COPY [[COPY]](s64)
     ; CHECK-DAG: $x0 = COPY [[COPY3]](s64)
-    ; CHECK-DAG: [[C1:%[0-9]+]]:_(s64) = G_CONSTANT i64 32
     ; CHECK-DAG: [[COPY4:%[0-9]+]]:_(s64) = COPY [[COPY]](s64)
-    ; CHECK-DAG: [[SHL:%[0-9]+]]:_(s64) = G_SHL [[COPY4]], [[C1]]
-    ; CHECK-DAG: [[ASHR:%[0-9]+]]:_(s64) = G_ASHR [[SHL]], [[C1]]
-    ; CHECK-DAG: $x0 = COPY [[ASHR]](s64)
-    ; CHECK-DAG: [[C2:%[0-9]+]]:_(s32) = G_CONSTANT i32 31
+    ; CHECK-DAG: [[SEXT_INREG:%[0-9]+]]:_(s64) = G_SEXT_INREG [[COPY4]], 32
+    ; CHECK-DAG: $x0 = COPY [[SEXT_INREG]](s64)
     ; CHECK-DAG: [[TRUNC4:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
-    ; CHECK-DAG: [[SHL1:%[0-9]+]]:_(s32) = G_SHL [[TRUNC4]], [[C2]]
-    ; CHECK-DAG: [[ASHR1:%[0-9]+]]:_(s32) = G_ASHR [[SHL1]], [[C2]]
-    ; CHECK-DAG: $w0 = COPY [[ASHR1]](s32)
+    ; CHECK-DAG: [[SEXT_INREG1:%[0-9]+]]:_(s32) = G_SEXT_INREG [[TRUNC4]], 1
+    ; CHECK-DAG: $w0 = COPY [[SEXT_INREG1]](s32)
     ; CHECK-DAG: [[C3:%[0-9]+]]:_(s32) = G_CONSTANT i32 255
     ; CHECK-DAG: [[TRUNC5:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
     ; CHECK-DAG: [[AND1:%[0-9]+]]:_(s32) = G_AND [[TRUNC5]], [[C3]]
     ; CHECK-DAG: $w0 = COPY [[AND1]](s32)
     ; CHECK-DAG: [[TRUNC6:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
     ; CHECK-DAG: $w0 = COPY [[TRUNC6]](s32)
     ; CHECK-DAG: [[C4:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
     ; CHECK-DAG: [[TRUNC7:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
     ; CHECK-DAG: [[AND2:%[0-9]+]]:_(s32) = G_AND [[TRUNC7]], [[C4]]
     ; CHECK-DAG: $w0 = COPY [[AND2]](s32)
     ; CHECK-DAG: [[TRUNC8:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
     ; CHECK-DAG: $w0 = COPY [[TRUNC8]](s32)
-    ; CHECK-DAG: [[C5:%[0-9]+]]:_(s32) = G_CONSTANT i32 16
     ; CHECK-DAG: [[TRUNC9:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
-    ; CHECK-DAG: [[SHL2:%[0-9]+]]:_(s32) = G_SHL [[TRUNC9]], [[C5]]
-    ; CHECK-DAG: [[ASHR2:%[0-9]+]]:_(s32) = G_ASHR [[SHL2]], [[C5]]
-    ; CHECK-DAG: $w0 = COPY [[ASHR2]](s32)
+    ; CHECK-DAG: [[SEXT_INREG2:%[0-9]+]]:_(s32) = G_SEXT_INREG [[TRUNC9]], 16
+    ; CHECK-DAG: $w0 = COPY [[SEXT_INREG2]](s32)
     ; CHECK-DAG: [[TRUNC10:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
     ; CHECK-DAG: [[AND3:%[0-9]+]]:_(s32) = G_AND [[TRUNC10]], [[C4]]
     ; CHECK-DAG: $w0 = COPY [[AND3]](s32)
     ; CHECK-DAG: [[TRUNC11:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
     ; CHECK-DAG: $w0 = COPY [[TRUNC11]](s32)
     ; CHECK-DAG: [[TRUNC12:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
     ; CHECK-DAG: $w0 = COPY [[TRUNC12]](s32)
     ; CHECK-DAG: [[FPEXT:%[0-9]+]]:_(s64) = G_FPEXT [[TRUNC12]](s32)

 ---
 name:            test_anyext_sext
 body:             |
   bb.0:
     liveins: $w0
 
     ; CHECK-LABEL: name: test_anyext_sext
     ; CHECK-DAG: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
-    ; CHECK-DAG: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 31
     ; CHECK-DAG: [[COPY1:%[0-9]+]]:_(s32) = COPY [[COPY]](s32)
-    ; CHECK-DAG: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[COPY1]], [[C]]
-    ; CHECK-DAG: [[ASHR:%[0-9]+]]:_(s32) = G_ASHR [[SHL]], [[C]]
-    ; CHECK-DAG: $w0 = COPY [[ASHR]](s32)
+    ; CHECK-DAG: [[SEXT_INREG:%[0-9]+]]:_(s32) = G_SEXT_INREG [[COPY1]], 1
+    ; CHECK-DAG: $w0 = COPY [[SEXT_INREG]](s32)
     %0:_(s32) = COPY $w0
     %1:_(s1) = G_TRUNC %0(s32)
     %2:_(s8) = G_SEXT %1(s1)
     %3:_(s32) = G_ANYEXT %2(s8)
     $w0 = COPY %3(s32)
 
 ...
 ---

llvm/test/CodeGen/AArch64/GlobalISel/legalize-gep.mir

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -march=aarch64 -run-pass=legalizer %s -o - | FileCheck %s
 ---
 name:            test_gep_small
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: test_gep_small
     ; CHECK-DAG: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
     ; CHECK-DAG: [[COPY1:%[0-9]+]]:_(s64) = COPY $x1
-    ; CHECK-DAG: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 56
     ; CHECK-DAG: [[COPY2:%[0-9]+]]:_(s64) = COPY [[COPY1]](s64)
-    ; CHECK-DAG: [[SHL:%[0-9]+]]:_(s64) = G_SHL [[COPY2]], [[C]]
-    ; CHECK-DAG: [[ASHR:%[0-9]+]]:_(s64) = G_ASHR [[SHL]], [[C]]
-    ; CHECK-DAG: [[GEP:%[0-9]+]]:_(p0) = G_GEP [[COPY]], [[ASHR]](s64)
+    ; CHECK-DAG: [[SEXT_INREG:%[0-9]+]]:_(s64) = G_SEXT_INREG [[COPY2]], 8
+    ; CHECK-DAG: [[GEP:%[0-9]+]]:_(p0) = G_GEP [[COPY]], [[SEXT_INREG]](s64)
     ; CHECK-DAG: $x0 = COPY [[GEP]](p0)
     %0:_(p0) = COPY $x0
     %1:_(s64) = COPY $x1
     %2:_(s8) = G_TRUNC %1(s64)
     %3:_(p0) = G_GEP %0, %2(s8)
     $x0 = COPY %3(p0)
 
 ...

llvm/test/CodeGen/AArch64/GlobalISel/legalize-itofp.mir

 
 ---
 name:            test_sitofp_s32_s1
 body: |
   bb.0:
     liveins: $w0
     ; CHECK-LABEL: name: test_sitofp_s32_s1
     ; CHECK-DAG: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
-    ; CHECK-DAG: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 31
     ; CHECK-DAG: [[COPY1:%[0-9]+]]:_(s32) = COPY [[COPY]](s32)
-    ; CHECK-DAG: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[COPY1]], [[C]](s32)
-    ; CHECK-DAG: [[ASHR:%[0-9]+]]:_(s32) = G_ASHR [[SHL]], [[C]](s32)
-    ; CHECK-DAG: [[SITOFP:%[0-9]+]]:_(s32) = G_SITOFP [[ASHR]](s32)
+    ; CHECK-DAG: [[SEXT_INREG:%[0-9]+]]:_(s32) = G_SEXT_INREG [[COPY1]], 1
+    ; CHECK-DAG: [[SITOFP:%[0-9]+]]:_(s32) = G_SITOFP [[SEXT_INREG]](s32)
     ; CHECK-DAG: $w0 = COPY [[SITOFP]](s32)
     %0:_(s32) = COPY $w0
     %1:_(s1) = G_TRUNC %0
     %2:_(s32) = G_SITOFP %1
     $w0 = COPY %2
 ...
 
 ---

 
 ---
 name:            test_sitofp_s64_s8
 body: |
   bb.0:
     liveins: $w0
     ; CHECK-LABEL: name: test_sitofp_s64_s8
     ; CHECK-DAG: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
-    ; CHECK-DAG: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 24
     ; CHECK-DAG: [[COPY1:%[0-9]+]]:_(s32) = COPY [[COPY]](s32)
-    ; CHECK-DAG: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[COPY1]], [[C]](s32)
-    ; CHECK-DAG: [[ASHR:%[0-9]+]]:_(s32) = G_ASHR [[SHL]], [[C]](s32)
-    ; CHECK-DAG: [[SITOFP:%[0-9]+]]:_(s64) = G_SITOFP [[ASHR]](s32)
+    ; CHECK-DAG: [[SEXT_INREG:%[0-9]+]]:_(s32) = G_SEXT_INREG [[COPY1]], 8
+    ; CHECK-DAG: [[SITOFP:%[0-9]+]]:_(s64) = G_SITOFP [[SEXT_INREG]](s32)
     ; CHECK-DAG: $x0 = COPY [[SITOFP]](s64)
     %0:_(s32) = COPY $w0
     %1:_(s8) = G_TRUNC %0
     %2:_(s64) = G_SITOFP %1
     $x0 = COPY %2
 ...
 
 ---

 
 ---
 name:            test_sitofp_s32_s16
 body: |
   bb.0:
     liveins: $w0
     ; CHECK-LABEL: name: test_sitofp_s32_s16
     ; CHECK-DAG: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
-    ; CHECK-DAG: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 16
     ; CHECK-DAG: [[COPY1:%[0-9]+]]:_(s32) = COPY [[COPY]](s32)
-    ; CHECK-DAG: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[COPY1]], [[C]](s32)
-    ; CHECK-DAG: [[ASHR:%[0-9]+]]:_(s32) = G_ASHR [[SHL]], [[C]](s32)
-    ; CHECK-DAG: [[SITOFP:%[0-9]+]]:_(s32) = G_SITOFP [[ASHR]](s32)
+    ; CHECK-DAG: [[SEXT_INREG:%[0-9]+]]:_(s32) = G_SEXT_INREG [[COPY1]], 16
+    ; CHECK-DAG: [[SITOFP:%[0-9]+]]:_(s32) = G_SITOFP [[SEXT_INREG]](s32)
     ; CHECK-DAG: $w0 = COPY [[SITOFP]](s32)
     %0:_(s32) = COPY $w0
     %1:_(s16) = G_TRUNC %0
     %2:_(s32) = G_SITOFP %1
     $w0 = COPY %2
 ...
 
 ---

llvm/test/CodeGen/AArch64/GlobalISel/legalize-rem.mir

 ...
 ---
 name:            test_srem_8
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: test_srem_8
     ; CHECK-DAG: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
     ; CHECK-DAG: [[COPY1:%[0-9]+]]:_(s64) = COPY $x1
-    ; CHECK-DAG: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 24
     ; CHECK-DAG: [[TRUNC:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
-    ; CHECK-DAG: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[TRUNC]], [[C]]
-    ; CHECK-DAG: [[ASHR:%[0-9]+]]:_(s32) = G_ASHR [[SHL]], [[C]]
+    ; CHECK-DAG: [[SEXT_INREG:%[0-9]+]]:_(s32) = G_SEXT_INREG [[TRUNC]], 8
     ; CHECK-DAG: [[TRUNC1:%[0-9]+]]:_(s32) = G_TRUNC [[COPY1]](s64)
-    ; CHECK-DAG: [[SHL1:%[0-9]+]]:_(s32) = G_SHL [[TRUNC1]], [[C]]
-    ; CHECK-DAG: [[ASHR1:%[0-9]+]]:_(s32) = G_ASHR [[SHL1]], [[C]]
-    ; CHECK-DAG: [[SDIV:%[0-9]+]]:_(s32) = G_SDIV [[ASHR]], [[ASHR1]]
+    ; CHECK-DAG: [[SEXT_INREG1:%[0-9]+]]:_(s32) = G_SEXT_INREG [[TRUNC1]], 8
+    ; CHECK-DAG: [[SDIV:%[0-9]+]]:_(s32) = G_SDIV [[SEXT_INREG]], [[SEXT_INREG1]]
     ; CHECK-DAG: [[COPY2:%[0-9]+]]:_(s32) = COPY [[SDIV]](s32)
     ; CHECK-DAG: [[TRUNC2:%[0-9]+]]:_(s32) = G_TRUNC [[COPY1]](s64)
     ; CHECK-DAG: [[MUL:%[0-9]+]]:_(s32) = G_MUL [[COPY2]], [[TRUNC2]]
     ; CHECK-DAG: [[TRUNC3:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
     ; CHECK-DAG: [[COPY3:%[0-9]+]]:_(s32) = COPY [[MUL]](s32)
     ; CHECK-DAG: [[SUB:%[0-9]+]]:_(s32) = G_SUB [[TRUNC3]], [[COPY3]]
     ; CHECK-DAG: [[COPY4:%[0-9]+]]:_(s32) = COPY [[SUB]](s32)
     ; CHECK-DAG: $w0 = COPY [[COPY4]](s32)

llvm/test/CodeGen/AArch64/GlobalISel/legalize-sext.mir

 # RUN: llc -march=aarch64 -run-pass=legalizer %s -o - | FileCheck %s
 ---
 name:            test_sext_inreg
 body: |
   bb.0.entry:
     liveins: $w0, $w1
     ; CHECK-LABEL: name: test_sext_inreg
     ; CHECK-DAG: [[COPY:%[0-9]+]]:_(s32) = COPY $w1
-    ; CHECK-DAG: [[I25:%[0-9]+]]:_(s32) = G_CONSTANT i32 25
-    ; CHECK-DAG: [[SEXT1:%[0-9]+]]:_(s32) = G_SHL [[COPY]], [[I25]]
-    ; CHECK-DAG: [[SEXT2:%[0-9]+]]:_(s32) = G_ASHR [[SEXT1]], [[I25]]
+    ; CHECK-DAG: [[SEXT2:%[0-9]+]]:_(s32) = G_SEXT_INREG [[COPY]], 7
     ; CHECK-DAG: $w0 = COPY [[SEXT2]](s32)
     %0:_(s32) = COPY $w1
     %2:_(s32) = G_SEXT_INREG %0(s32), 7
     $w0 = COPY %2(s32)
 ...
+---
+name:            test_sext_inreg_16
+body: |
+  bb.0.entry:
+    liveins: $w0, $w1
+    ; CHECK-LABEL: name: test_sext_inreg_16
+    ; CHECK-DAG: [[COPY:%[0-9]+]]:_(s32) = COPY $w1
+    ; CHECK-DAG: [[SEXT:%[0-9]+]]:_(s32) = G_SEXT_INREG [[COPY]], 16
+    ; CHECK-DAG: $w0 = COPY [[SEXT]](s32)
+    %0:_(s32) = COPY $w1
+    %2:_(s32) = G_SEXT_INREG %0(s32), 16
+    $w0 = COPY %2(s32)
+...

llvm/test/CodeGen/AArch64/GlobalISel/legalize-shift.mir

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -O0 -march=aarch64 -run-pass=legalizer %s -o - | FileCheck %s
 ---
 name:            test_shift
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: test_shift
     ; CHECK-DAG: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
     ; CHECK-DAG: [[COPY1:%[0-9]+]]:_(s64) = COPY $x1
     ; CHECK-DAG: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 255
     ; CHECK-DAG: [[TRUNC:%[0-9]+]]:_(s32) = G_TRUNC [[COPY1]](s64)
     ; CHECK-DAG: [[AND:%[0-9]+]]:_(s32) = G_AND [[TRUNC]], [[C]]
-    ; CHECK-DAG: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 24
     ; CHECK-DAG: [[TRUNC1:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
-    ; CHECK-DAG: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[TRUNC1]], [[C1]](s32)
-    ; CHECK-DAG: [[ASHR:%[0-9]+]]:_(s32) = G_ASHR [[SHL]], [[C1]](s32)
-    ; CHECK-DAG: [[ASHR1:%[0-9]+]]:_(s32) = G_ASHR [[ASHR]], [[AND]](s32)
+    ; CHECK-DAG: [[SEXT_INREG:%[0-9]+]]:_(s32) = G_SEXT_INREG [[TRUNC1]], 8
+    ; CHECK-DAG: [[ASHR1:%[0-9]+]]:_(s32) = G_ASHR [[SEXT_INREG]], [[AND]](s32)
     ; CHECK-DAG: [[COPY2:%[0-9]+]]:_(s32) = COPY [[ASHR1]](s32)
     ; CHECK-DAG: $w0 = COPY [[COPY2]](s32)
     ; CHECK-DAG: [[TRUNC2:%[0-9]+]]:_(s32) = G_TRUNC [[COPY1]](s64)
     ; CHECK-DAG: [[AND1:%[0-9]+]]:_(s32) = G_AND [[TRUNC2]], [[C]]
     ; CHECK-DAG: [[TRUNC3:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
     ; CHECK-DAG: [[AND2:%[0-9]+]]:_(s32) = G_AND [[TRUNC3]], [[C]]
     ; CHECK-DAG: [[LSHR:%[0-9]+]]:_(s32) = G_LSHR [[AND2]], [[AND1]](s32)
     ; CHECK-DAG: [[COPY3:%[0-9]+]]:_(s32) = COPY [[LSHR]](s32)

llvm/test/CodeGen/AArch64/GlobalISel/legalize-undef.mir

 
 ---
 name: test_implicit_def_s3
 body: |
   bb.0:
     liveins:
 
     ; CHECK-LABEL: name: test_implicit_def_s3
-    ; CHECK-DAG: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 61
     ; CHECK-DAG: [[DEF:%[0-9]+]]:_(s64) = G_IMPLICIT_DEF
-    ; CHECK-DAG: [[SHL:%[0-9]+]]:_(s64) = G_SHL [[DEF]], [[C]](s64)
-    ; CHECK-DAG: [[ASHR:%[0-9]+]]:_(s64) = G_ASHR [[SHL]], [[C]](s64)
-    ; CHECK-DAG: $x0 = COPY [[ASHR]](s64)
+    ; CHECK-DAG: [[SEXT_INREG:%[0-9]+]]:_(s64) = G_SEXT_INREG [[DEF]], 3
+    ; CHECK-DAG: $x0 = COPY [[SEXT_INREG]](s64)
     %0:_(s3) = G_IMPLICIT_DEF
     %1:_(s64) = G_SEXT %0
     $x0 = COPY %1(s64)
 ...
 
 # FIXME: s2 not correctly handled
 
 ---

llvm/test/CodeGen/AArch64/GlobalISel/select-int-ext.mir

--- |
   define void @zext_s32_from_s8() { ret void }
   define void @zext_s16_from_s8() { ret void }
   define void @zext_v8s16_from_v8s8() { ret void }
   define void @zext_v4s32_from_v4s16() { ret void }
   define void @zext_v2s64_from_v2s32() { ret void }
 
   define void @sext_s64_from_s32() { ret void }
   define void @sext_s32_from_s16() { ret void }
+  define void @sext_s32_from_s16_inreg() { ret void }
+  define void @sext_s32_from_s3_inreg() { ret void }
   define void @sext_s32_from_s8() { ret void }
   define void @sext_s16_from_s8() { ret void }
   define void @sext_v8s16_from_v8s8() { ret void }
   define void @sext_v4s32_from_v4s16() { ret void }
   define void @sext_v2s64_from_v2s32() { ret void }
 ...
 
 ---
bb.0:
     ; CHECK: $w0 = COPY [[SBFMWri]]
     %2:gpr(s32) = COPY $w0
     %0(s16) = G_TRUNC %2
     %1(s32) = G_SEXT %0
     $w0 = COPY %1
 ...
 
 ---
+name:            sext_s32_from_s16_inreg
+legalized:       true
+regBankSelected: true
+
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; CHECK-LABEL: name: sext_s32_from_s16_inreg
+    ; CHECK: [[COPY:%[0-9]+]]:gpr32 = COPY $w0
+    ; CHECK: [[SBFMWri:%[0-9]+]]:gpr32 = SBFMWri [[COPY]], 0, 15
+    ; CHECK: $w0 = COPY [[SBFMWri]]
+    %0:gpr(s32) = COPY $w0
+    %1(s32) = G_SEXT_INREG %0, 16
+    $w0 = COPY %1
+...
+
+---
+name:            sext_s32_from_s3_inreg
+legalized:       true
+regBankSelected: true
+
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; CHECK-LABEL: name: sext_s32_from_s3_inreg
+    ; CHECK: [[COPY:%[0-9]+]]:gpr32 = COPY $w0
+    ; CHECK: [[SBFMWri:%[0-9]+]]:gpr32 = SBFMWri [[COPY]], 0, 2
+    ; CHECK: $w0 = COPY [[SBFMWri]]
+    %0:gpr(s32) = COPY $w0
+    %1(s32) = G_SEXT_INREG %0, 3
+    $w0 = COPY %1
+...
+
+---
 name:            sext_s32_from_s8
 legalized:       true
 regBankSelected: true
 
 registers:
   - { id: 0, class: gpr }
   - { id: 1, class: gpr }
 

llvm/test/CodeGen/AArch64/arm64-fastisel-gep-promote-before-add.ll

 ; fastisel should not fold add with non-pointer bitwidth
 ; sext(a) + sext(b) != sext(a + b)
+; FIXME: This test isn't fastisel anymore. It's GlobalISel
 ; RUN: llc -mtriple=arm64-apple-darwin %s -O0 -o - | FileCheck %s
 
 define zeroext i8 @gep_promotion(i8* %ptr) nounwind uwtable ssp {
 entry:
   %ptr.addr = alloca i8*, align 8
   %add = add i8 64, 64 ; 0x40 + 0x40
   %0 = load i8*, i8** %ptr.addr, align 8
 
   ; CHECK-LABEL: _gep_promotion:
-  ; CHECK: ldrb {{[a-z][0-9]+}}, {{\[[a-z][0-9]+\]}}
+  ; CHECK: ldurb {{[a-z][0-9]+}}, [{{[a-z][0-9]+}}, #-128]
   %arrayidx = getelementptr inbounds i8, i8* %0, i8 %add
 
   %1 = load i8, i8* %arrayidx, align 1
   ret i8 %1
 }
 

llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

TEST_F(GISelMITest, NarrowSEXTINREG) {
   // Check
   ASSERT_TRUE(CheckMachineFunction(*MF, CheckStr));
 }
 
 TEST_F(GISelMITest, NarrowSEXTINREG2) {
   if (!TM)
     return;
 
-  // Declare your legalization info, these aren't actually relevant to the test.
-  DefineLegalizerInfo(
-      A, { getActionDefinitionsBuilder(G_SEXT_INREG).legalFor({s64}); });
+  // Declare your legalization info, these aren't actually relevant to the
+  // helper test below.
+  DefineLegalizerInfo(A, {
+    getActionDefinitionsBuilder(G_SEXT_INREG)
+        .legalForTypeWithImm({{s64, 8}})
+        .lower();
+  });
+
   // Build Instr
   auto MIB = B.buildInstr(
       TargetOpcode::G_SEXT_INREG, {LLT::scalar(32)},
       {B.buildInstr(TargetOpcode::G_TRUNC, {LLT::scalar(32)}, {Copies[0]}),
        APInt(32, 9)});
   AInfo Info(MF->getSubtarget());
   DummyGISelObserver Observer;
   LegalizerHelper Helper(*MF, Info, Observer, B);
TEST_F(GISelMITest, NarrowSEXTINREG2) {
   ASSERT_TRUE(CheckMachineFunction(*MF, CheckStr));
 }
 
 TEST_F(GISelMITest, LowerSEXTINREG) {
   if (!TM)
     return;
 
   // Declare your legalization info, these aren't actually relevant to the test.
-  DefineLegalizerInfo(
-      A, { getActionDefinitionsBuilder(G_SEXT_INREG).legalFor({s64}); });
+  DefineLegalizerInfo(A, {
+    getActionDefinitionsBuilder(G_SEXT_INREG)
+        .legalForTypeWithImm({{s64, 8}})
+        .lower();
+  });
   // Build Instr
   auto MIB = B.buildInstr(
       TargetOpcode::G_SEXT_INREG, {LLT::scalar(32)},
       {B.buildInstr(TargetOpcode::G_TRUNC, {LLT::scalar(32)}, {Copies[0]}),
        APInt(32, 8)});
   AInfo Info(MF->getSubtarget());
   DummyGISelObserver Observer;
   LegalizerHelper Helper(*MF, Info, Observer, B);
   // Perform Legalization
   ASSERT_TRUE(Helper.lower(*MIB, 0, LLT()) ==
               LegalizerHelper::LegalizeResult::Legalized);
   MF->dump();
 
   auto CheckStr = R"(
   CHECK: [[T1:%[0-9]+]]:_(s32) = G_TRUNC
   CHECK: [[CST:%[0-9]+]]:_(s32) = G_CONSTANT i32 24
   CHECK: [[T2:%[0-9]+]]:_(s32) = G_SHL [[T1]]:_, [[CST]]:_
   CHECK: [[T3:%[0-9]+]]:_(s32) = G_ASHR [[T2]]:_, [[CST]]:_
   )";
 
   // Check
   ASSERT_TRUE(CheckMachineFunction(*MF, CheckStr));
+
+  LLT s64 = LLT::scalar(64);
+  ASSERT_TRUE(
+      Info.getAction({TargetOpcode::G_SEXT_INREG, {s64}, {}, {8}}).Action ==
+      LegalizeAction::Legal);
+  ASSERT_TRUE(
+      Info.getAction({TargetOpcode::G_SEXT_INREG, {s64}, {}, {9}}).Action ==
+      LegalizeAction::Lower);
 }
 } // namespace