[GlobalISel] Add G_SBFX + G_UBFX (bitfield extraction opcodes)

There is a bunch of similar bitfield extraction code throughout *ISelDAGToDAG.

E.g, ARMISelDAGToDAG, AArch64ISelDAGToDAG, and AMDGPUISelDAGToDAG all contain
code that matches a bitfield extract from an and + right shift.

Rather than duplicating code in the same way, this adds two opcodes:

- G_UBFX (unsigned bitfield extract)
- G_SBFX (signed bitfield extract)

They work like this

```
%x = G_UBFX %y, %lsb, %width
```

Where `lsb` and `width` are

- The least-significant bit of the extraction
- The width of the extraction

This will extract `width` bits from `%y`, starting at `lsb`. G_UBFX zero-extends
the result, while G_SBFX sign-extends the result.

This should allow us to use the combiner to match the bitfield extraction
patterns rather than duplicating pattern-matching code in each target.

Differential Revision: https://reviews.llvm.org/D98464

llvm/docs/GlobalISel/GenericOpcode.rst

@@ -233,6 +233,39 @@ Reverse the order of the bits in a scalar.
 
   %1:_(s32) = G_BITREVERSE %0:_(s32)
 
+G_SBFX, G_UBFX
+^^^^^^^^^^^^^^
+
+Extract a range of bits from a register.
+
+The source operands are registers as follows:
+
+- Source
+- The least-significant bit for the extraction
+- The width of the extraction
+
+G_SBFX sign-extends the result, while G_UBFX zero-extends the result.
+
+.. code-block:: none
+
+  ; Extract 5 bits starting at bit 1 from %x and store them in %a.
+  ; Sign-extend the result.
+  ;
+  ; Example:
+  ; %x = 0...0000[10110]1 ---> %a = 1...111111[10110]
+  %lsb_one = G_CONSTANT i32 1
+  %width_five = G_CONSTANT i32 5
+  %a:_(s32) = G_SBFX %x, %lsb_one, %width_five
+
+  ; Extract 3 bits starting at bit 2 from %x and store them in %b. Zero-extend
+  ; the result.
+  ;
+  ; Example:
+  ; %x = 1...11111[100]11 ---> %b = 0...00000[100]
+  %lsb_two = G_CONSTANT i32 2
+  %width_three = G_CONSTANT i32 3
+  %b:_(s32) = G_UBFX %x, %lsb_two, %width_three
+
 Integer Operations
 -------------------
 

llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

@@ -1831,6 +1831,18 @@ public:
                                 DstMMO, SrcMMO);
   }
 
+  /// Build and insert \p Dst = G_SBFX \p Src, \p LSB, \p Width.
+  MachineInstrBuilder buildSbfx(const DstOp &Dst, const SrcOp &Src,
+                                const SrcOp &LSB, const SrcOp &Width) {
+    return buildInstr(TargetOpcode::G_SBFX, {Dst}, {Src, LSB, Width});
+  }
+
+  /// Build and insert \p Dst = G_UBFX \p Src, \p LSB, \p Width.
+  MachineInstrBuilder buildUbfx(const DstOp &Dst, const SrcOp &Src,
+                                const SrcOp &LSB, const SrcOp &Width) {
+    return buildInstr(TargetOpcode::G_UBFX, {Dst}, {Src, LSB, Width});
+  }
+
   virtual MachineInstrBuilder buildInstr(unsigned Opc, ArrayRef<DstOp> DstOps,
                                          ArrayRef<SrcOp> SrcOps,
                                          Optional<unsigned> Flags = None);

llvm/include/llvm/Support/TargetOpcodes.def

@@ -749,10 +749,13 @@ HANDLE_TARGET_OPCODE(G_VECREDUCE_SMIN)
 HANDLE_TARGET_OPCODE(G_VECREDUCE_UMAX)
 HANDLE_TARGET_OPCODE(G_VECREDUCE_UMIN)
 
+HANDLE_TARGET_OPCODE(G_SBFX)
+HANDLE_TARGET_OPCODE(G_UBFX)
+
 /// Marker for the end of the generic opcode.
 /// This is used to check if an opcode is in the range of the
 /// generic opcodes.
-HANDLE_TARGET_OPCODE_MARKER(PRE_ISEL_GENERIC_OPCODE_END, G_VECREDUCE_UMIN)
+HANDLE_TARGET_OPCODE_MARKER(PRE_ISEL_GENERIC_OPCODE_END, G_UBFX)
 
 /// BUILTIN_OP_END - This must be the last enum value in this list.
 /// The target-specific post-isel opcode values start here.

llvm/include/llvm/Target/GenericOpcodes.td

@@ -1354,6 +1354,24 @@ def G_MEMSET : GenericInstruction {
   let mayStore = true;
 }
 
+//------------------------------------------------------------------------------
+// Bitfield extraction.
+//------------------------------------------------------------------------------
+
+// Generic signed bitfield extraction.
+def G_SBFX : GenericInstruction {
+  let OutOperandList = (outs type0:$dst);
+  let InOperandList = (ins type0:$src, type0:$lsb, type0:$width);
+  let hasSideEffects = false;
+}
+
+// Generic unsigned bitfield extraction.
+def G_UBFX : GenericInstruction {
+  let OutOperandList = (outs type0:$dst);
+  let InOperandList = (ins type0:$src, type0:$lsb, type0:$width);
+  let hasSideEffects = false;
+}
+
 //------------------------------------------------------------------------------
 // Optimization hints
 //------------------------------------------------------------------------------

llvm/lib/CodeGen/MachineVerifier.cpp

@@ -1566,6 +1566,17 @@ void MachineVerifier::verifyPreISelGenericInstruction(const MachineInstr *MI) {
       report("Vector reduction requires vector source=", MI);
     break;
   }
+
+  case TargetOpcode::G_SBFX:
+  case TargetOpcode::G_UBFX: {
+    LLT DstTy = MRI->getType(MI->getOperand(0).getReg());
+    if (DstTy.isVector()) {
+      report("Bitfield extraction is not supported on vectors", MI);
+      break;
+    }
+    break;
+  }
+
   default:
     break;
   }

llvm/test/MachineVerifier/test_g_ubfx_sbfx.mir

@@ -0,0 +1,15 @@
+# RUN: not --crash llc -verify-machineinstrs -run-pass none -o /dev/null %s 2>&1 | FileCheck %s
+# REQUIRES: aarch64-registered-target
+
+name:            test
+body: |
+  bb.0:
+    %v1:_(<2 x s64>) = G_IMPLICIT_DEF
+    %v2:_(<2 x s64>) = G_IMPLICIT_DEF
+    %v3:_(<2 x s64>) = G_IMPLICIT_DEF
+
+    ; CHECK: *** Bad machine code: Bitfield extraction is not supported on vectors ***
+    %ubfx_vector:_(<2 x s64>) = G_UBFX %v1, %v2, %v3
+    ; CHECK: *** Bad machine code: Bitfield extraction is not supported on vectors ***
+    %sbfx_vector:_(<2 x s64>) = G_SBFX %v1, %v2, %v3
+...

llvm/unittests/CodeGen/GlobalISel/MachineIRBuilderTest.cpp

@@ -398,3 +398,25 @@ TEST_F(AArch64GISelMITest, BuildAddoSubo) {
 
   EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
 }
+
+TEST_F(AArch64GISelMITest, BuildBitfieldExtract) {
+  setUp();
+  if (!TM)
+    return;
+  LLT S64 = LLT::scalar(64);
+  SmallVector<Register, 4> Copies;
+  collectCopies(Copies, MF);
+
+  auto Ubfx = B.buildUbfx(S64, Copies[0], Copies[1], Copies[2]);
+  B.buildSbfx(S64, Ubfx, Copies[0], Copies[2]);
+
+  const auto *CheckStr = R"(
+  ; CHECK: [[COPY0:%[0-9]+]]:_(s64) = COPY $x0
+  ; CHECK: [[COPY1:%[0-9]+]]:_(s64) = COPY $x1
+  ; CHECK: [[COPY2:%[0-9]+]]:_(s64) = COPY $x2
+  ; CHECK: [[UBFX:%[0-9]+]]:_(s64) = G_UBFX [[COPY0]]:_, [[COPY1]]:_, [[COPY2]]:_
+  ; CHECK: [[SBFX:%[0-9]+]]:_(s64) = G_SBFX [[UBFX]]:_, [[COPY0]]:_, [[COPY2]]:_
+  )";
+
+  EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
+}