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

[paquette]

Summary

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 denote

• 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.

[arsenm]

llvm/include/llvm/Target/GenericOpcodes.td

1364

These can be registers like a normal op. AMDGPU has registers/variable inputs for the offset and width

[aemerson]

If we have to allow variable width operands for these, it doesn't really help AArch64 that much. We have to match the constant pattern to a target specific opcode, and then lower the rest of these to shifts by default. I can see the benefit of AMDGPU, but with variable extracts this just puts us back where we started for AArch64.

[paquette]

• Make the LSB + width registers
• Update verifier
• Add a MIRBuilder test

[paquette]

In D98464#2622698, @aemerson wrote:

If we have to allow variable width operands for these, it doesn't really help AArch64 that much. We have to match the constant pattern to a target specific opcode, and then lower the rest of these to shifts by default. I can see the benefit of AMDGPU, but with variable extracts this just puts us back where we started for AArch64.

Would a allowsVariableWidthExtracts hook of some sort help? Then we could match the constants-only pattern in AArch64, but allow other targets to have variable-width operands.

[paquette]

I think something like this would work post-legalization:

// We'd need these functions...
if (isBeforeLegalizer() || !isLegal(...))
  return false;

Register LshrLHS, LshrRHS;
int64_t Mask;
if (!mi_match(Reg, MRI,
              m_GAnd(m_GLShr(m_Reg(LshrLHS), m_Reg(LshrRHS)), m_ICst(Mask))))
  return false;
if (/*Mask isn't a mask...*/)
  return false;

if (TLI.allowsRegisterExtractOps(/*...*/)) {
  // ... Do stuff ...
  return true;
}

// Need immediates for LSB + width.
int64_t LshrImm;
if (!mi_match(LshrRHS, MRI, m_ICst(LshrImm)))
  return false;

// ... Do stuff ...
return true;

[aemerson]

In D98464#2622885, @paquette wrote:

I think something like this would work post-legalization:

// We'd need these functions...
if (isBeforeLegalizer() || !isLegal(...))
  return false;

Register LshrLHS, LshrRHS;
int64_t Mask;
if (!mi_match(Reg, MRI,
              m_GAnd(m_GLShr(m_Reg(LshrLHS), m_Reg(LshrRHS)), m_ICst(Mask))))
  return false;
if (/*Mask isn't a mask...*/)
  return false;

if (TLI.allowsRegisterExtractOps(/*...*/)) {
  // ... Do stuff ...
  return true;
}

// Need immediates for LSB + width.
int64_t LshrImm;
if (!mi_match(LshrRHS, MRI, m_ICst(LshrImm)))
  return false;

// ... Do stuff ...
return true;

So we form these post-legalize only if we can see that the sources are constant for arm64. Sounds reasonable to me. @arsenm that ok for AMDGPU?

[arsenm]

In D98464#2622698, @aemerson wrote:

If we have to allow variable width operands for these, it doesn't really help AArch64 that much. We have to match the constant pattern to a target specific opcode, and then lower the rest of these to shifts by default. I can see the benefit of AMDGPU, but with variable extracts this just puts us back where we started for AArch64.

In D98464#2627521, @aemerson wrote:

In D98464#2622885, @paquette wrote:

I think something like this would work post-legalization:

// We'd need these functions...
if (isBeforeLegalizer() || !isLegal(...))
  return false;

Register LshrLHS, LshrRHS;
int64_t Mask;
if (!mi_match(Reg, MRI,
              m_GAnd(m_GLShr(m_Reg(LshrLHS), m_Reg(LshrRHS)), m_ICst(Mask))))
  return false;
if (/*Mask isn't a mask...*/)
  return false;

if (TLI.allowsRegisterExtractOps(/*...*/)) {
  // ... Do stuff ...
  return true;
}

// Need immediates for LSB + width.
int64_t LshrImm;
if (!mi_match(LshrRHS, MRI, m_ICst(LshrImm)))
  return false;

// ... Do stuff ...
return true;

So we form these post-legalize only if we can see that the sources are constant for arm64. Sounds reasonable to me. @arsenm that ok for AMDGPU?

Yes, forming them only if the target supports the variable case makes sense (although also having the legalize back to shifts path would be nice for consistency)

[aemerson]

LGTM.

llvm/docs/GlobalISel/GenericOpcode.rst

 ^^^^^^^^^^^^
 
 Reverse the order of the bits in a scalar.
 
 .. code-block:: none
 
   %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
 -------------------
 
 G_ADD, G_SUB, G_MUL, G_AND, G_OR, G_XOR, G_SDIV, G_UDIV, G_SREM, G_UREM
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 These each perform their respective integer arithmetic on a scalar.
 

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

 
   MachineInstrBuilder buildMemCpy(const SrcOp &DstPtr, const SrcOp &SrcPtr,
                                   const SrcOp &Size, MachineMemOperand &DstMMO,
                                   MachineMemOperand &SrcMMO) {
     return buildMemTransferInst(TargetOpcode::G_MEMCPY, DstPtr, SrcPtr, Size,
                                 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);
 };
 
 } // End namespace llvm.
 #endif // LLVM_CODEGEN_GLOBALISEL_MACHINEIRBUILDER_H

llvm/include/llvm/Support/TargetOpcodes.def

 HANDLE_TARGET_OPCODE(G_VECREDUCE_AND)
 HANDLE_TARGET_OPCODE(G_VECREDUCE_OR)
 HANDLE_TARGET_OPCODE(G_VECREDUCE_XOR)
 HANDLE_TARGET_OPCODE(G_VECREDUCE_SMAX)
 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.
 HANDLE_TARGET_OPCODE_MARKER(GENERIC_OP_END, PRE_ISEL_GENERIC_OPCODE_END)

llvm/include/llvm/Target/GenericOpcodes.td

 def G_MEMSET : GenericInstruction {
   let OutOperandList = (outs);
   let InOperandList = (ins ptype0:$dst_addr, type1:$value, type2:$size, untyped_imm_0:$tailcall);
   let hasSideEffects = false;
   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
 //------------------------------------------------------------------------------
 
 // Asserts that an operation has already been zero-extended from a specific
 // type.
 def G_ASSERT_ZEXT : GenericInstruction {
   let OutOperandList = (outs type0:$dst);
   let InOperandList = (ins type0:$src, untyped_imm_0:$sz);

llvm/lib/CodeGen/MachineVerifier.cpp

case TargetOpcode::G_VECREDUCE_UMIN: {
     LLT DstTy = MRI->getType(MI->getOperand(0).getReg());
     LLT SrcTy = MRI->getType(MI->getOperand(1).getReg());
     if (!DstTy.isScalar())
       report("Vector reduction requires a scalar destination type", MI);
     if (!SrcTy.isVector())
       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;
   }
 }
 
 void MachineVerifier::visitMachineInstrBefore(const MachineInstr *MI) {
   const MCInstrDesc &MCID = MI->getDesc();
   if (MI->getNumOperands() < MCID.getNumOperands()) {

llvm/test/MachineVerifier/test_g_ubfx_sbfx.mir

+# 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

TEST_F(AArch64GISelMITest, BuildAddoSubo) {
   ; CHECK: [[UADDE:%[0-9]+]]:_(s64), [[UADDE_FLAG:%[0-9]+]]:_(s1) = G_UADDE [[COPY0]]:_, [[COPY1]]:_, [[UADDO_FLAG]]
   ; CHECK: [[USUBE:%[0-9]+]]:_(s64), [[USUBE_FLAG:%[0-9]+]]:_(s1) = G_USUBE [[COPY0]]:_, [[COPY1]]:_, [[USUBO_FLAG]]
   ; CHECK: [[SADDE:%[0-9]+]]:_(s64), [[SADDE_FLAG:%[0-9]+]]:_(s1) = G_SADDE [[COPY0]]:_, [[COPY1]]:_, [[SADDO_FLAG]]
   ; CHECK: [[SSUBE:%[0-9]+]]:_(s64), [[SSUBE_FLAG:%[0-9]+]]:_(s1) = G_SSUBE [[COPY0]]:_, [[COPY1]]:_, [[SSUBO_FLAG]]
   )";
 
   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;
+}