[AArch64][GlobalISel] Select G_SBFX and G_UBFX

[paquette]

Summary

Add selection support for G_SBFX and G_UBFX and add a test.

These must always have a constant LSB and width.

Usually real SBFX/UBFX instructions have some requirements on the LSB/width operands (e.g. sbfx x, 31, 31 probably isn't allowed.) However, it seems like these are ultimately selected down to a sane instruction (e.g. asr) anyway. (See: https://godbolt.org/z/7KeeaerWG)

[aemerson]

LGTM.

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

if (selectImpl(I, *CoverageInfo))
     return true;
 
   LLT Ty =
       I.getOperand(0).isReg() ? MRI.getType(I.getOperand(0).getReg()) : LLT{};
 
   MachineIRBuilder MIB(I);
 
   switch (Opcode) {
+  case TargetOpcode::G_SBFX:
+  case TargetOpcode::G_UBFX: {
+    static const unsigned OpcTable[2][2] = {
+        {AArch64::UBFMWri, AArch64::UBFMXri},
+        {AArch64::SBFMWri, AArch64::SBFMXri}};
+    bool IsSigned = Opcode == TargetOpcode::G_SBFX;
+    unsigned Size = Ty.getSizeInBits();
+    unsigned Opc = OpcTable[IsSigned][Size == 64];
+    auto Cst1 =
+        getConstantVRegValWithLookThrough(I.getOperand(2).getReg(), MRI);
+    assert(Cst1 && "Should have gotten a constant for src 1?");
+    auto Cst2 =
+        getConstantVRegValWithLookThrough(I.getOperand(3).getReg(), MRI);
+    assert(Cst2 && "Should have gotten a constant for src 2?");
+    auto LSB = Cst1->Value.getZExtValue();
+    auto Width = Cst2->Value.getZExtValue();
+    MachineIRBuilder MIB(I);
+    auto BitfieldInst =
+        MIB.buildInstr(Opc, {I.getOperand(0)}, {I.getOperand(1)})
+            .addImm(LSB)
+            .addImm(Width);
+    I.eraseFromParent();
+    return constrainSelectedInstRegOperands(*BitfieldInst, TII, TRI, RBI);
+  }
   case TargetOpcode::G_BRCOND:
     return selectCompareBranch(I, MF, MRI);
 
   case TargetOpcode::G_BRINDIRECT: {
     I.setDesc(TII.get(AArch64::BR));
     return constrainSelectedInstRegOperands(I, TII, TRI, RBI);
   }
 

llvm/test/CodeGen/AArch64/GlobalISel/select-sbfx.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64 -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s
+
+...
+---
+name:            sbfx_s32
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; CHECK-LABEL: name: sbfx_s32
+    ; CHECK: liveins: $w0
+    ; CHECK: %copy:gpr32 = COPY $w0
+    ; CHECK: %sbfx:gpr32 = SBFMWri %copy, 0, 10
+    ; CHECK: $w0 = COPY %sbfx
+    ; CHECK: RET_ReallyLR implicit $w0
+    %copy:gpr(s32) = COPY $w0
+    %cst1:gpr(s32) = G_CONSTANT i32 0
+    %cst2:gpr(s32) = G_CONSTANT i32 10
+    %sbfx:gpr(s32) = G_SBFX %copy, %cst1, %cst2
+    $w0 = COPY %sbfx
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            sbfx_s64
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+    ; CHECK-LABEL: name: sbfx_s64
+    ; CHECK: liveins: $x0
+    ; CHECK: %copy:gpr64 = COPY $x0
+    ; CHECK: %sbfx:gpr64 = SBFMXri %copy, 0, 10
+    ; CHECK: $x0 = COPY %sbfx
+    ; CHECK: RET_ReallyLR implicit $x0
+    %copy:gpr(s64) = COPY $x0
+    %cst1:gpr(s64) = G_CONSTANT i64 0
+    %cst2:gpr(s64) = G_CONSTANT i64 10
+    %sbfx:gpr(s64) = G_SBFX %copy, %cst1, %cst2
+    $x0 = COPY %sbfx
+    RET_ReallyLR implicit $x0
+
+...
+---
+name:            sbfx_s32_31_31
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; This is just an asr, so it's okay.
+
+    ; CHECK-LABEL: name: sbfx_s32_31_31
+    ; CHECK: liveins: $w0
+    ; CHECK: %copy:gpr32 = COPY $w0
+    ; CHECK: %sbfx:gpr32 = SBFMWri %copy, 31, 31
+    ; CHECK: $w0 = COPY %sbfx
+    ; CHECK: RET_ReallyLR implicit $w0
+    %copy:gpr(s32) = COPY $w0
+    %cst1:gpr(s32) = G_CONSTANT i32 31
+    %cst2:gpr(s32) = G_CONSTANT i32 31
+    %sbfx:gpr(s32) = G_SBFX %copy, %cst1, %cst2
+    $w0 = COPY %sbfx
+    RET_ReallyLR implicit $w0

llvm/test/CodeGen/AArch64/GlobalISel/select-ubfx.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64 -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s
+
+...
+---
+name:            ubfx_s32
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; CHECK-LABEL: name: ubfx_s32
+    ; CHECK: liveins: $w0
+    ; CHECK: %copy:gpr32 = COPY $w0
+    ; CHECK: %ubfx:gpr32 = UBFMWri %copy, 0, 10
+    ; CHECK: $w0 = COPY %ubfx
+    ; CHECK: RET_ReallyLR implicit $w0
+    %copy:gpr(s32) = COPY $w0
+    %cst1:gpr(s32) = G_CONSTANT i32 0
+    %cst2:gpr(s32) = G_CONSTANT i32 10
+    %ubfx:gpr(s32) = G_UBFX %copy, %cst1, %cst2
+    $w0 = COPY %ubfx
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            ubfx_s64
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+    ; CHECK-LABEL: name: ubfx_s64
+    ; CHECK: liveins: $x0
+    ; CHECK: %copy:gpr64 = COPY $x0
+    ; CHECK: %ubfx:gpr64 = UBFMXri %copy, 0, 10
+    ; CHECK: $x0 = COPY %ubfx
+    ; CHECK: RET_ReallyLR implicit $x0
+    %copy:gpr(s64) = COPY $x0
+    %cst1:gpr(s64) = G_CONSTANT i64 0
+    %cst2:gpr(s64) = G_CONSTANT i64 10
+    %ubfx:gpr(s64) = G_UBFX %copy, %cst1, %cst2
+    $x0 = COPY %ubfx
+    RET_ReallyLR implicit $x0
+
+...
+---
+name:            ubfx_s32_31_31
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; This is just a lsr, so it's okay.
+
+    ; CHECK-LABEL: name: ubfx_s32_31_31
+    ; CHECK: liveins: $w0
+    ; CHECK: %copy:gpr32 = COPY $w0
+    ; CHECK: %ubfx:gpr32 = UBFMWri %copy, 31, 31
+    ; CHECK: $w0 = COPY %ubfx
+    ; CHECK: RET_ReallyLR implicit $w0
+    %copy:gpr(s32) = COPY $w0
+    %cst1:gpr(s32) = G_CONSTANT i32 31
+    %cst2:gpr(s32) = G_CONSTANT i32 31
+    %ubfx:gpr(s32) = G_UBFX %copy, %cst1, %cst2
+    $w0 = COPY %ubfx
+    RET_ReallyLR implicit $w0