[GlobalISel] Allow different tyeps for G_SBFX and G_UBFX operands

[bcahoon]

Summary

Change the definition of G_SBFX and G_UBFX so that the lsb and width
can have different types than the src and dst operands.

[arsenm]

Description is misleading. This should be allow a different type for the offset/width operands, which do not necessarily need to be constant

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

718

Maybe this should keep using i64 for the amounts? It probably doesn't really matter since it only handles constant cases

[bcahoon]

Keep operand types for lsb and width to be the same type as the src/dst for AArch64.

[arsenm]

Typo tyeps in commit message

[paquette]

LGTM

llvm/include/llvm/Target/GenericOpcodes.td

 
 //------------------------------------------------------------------------------
 // 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 InOperandList = (ins type0:$src, type1:$lsb, type1:$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 InOperandList = (ins type0:$src, type1:$lsb, type1:$width);
   let hasSideEffects = false;
 }
 
 //------------------------------------------------------------------------------
 // Optimization hints
 //------------------------------------------------------------------------------
 
 // Asserts that an operation has already been zero-extended from a specific

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

AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
   getActionDefinitionsBuilder(G_ROTR)
       .legalFor({{s32, s64}, {s64, s64}})
       .customIf([=](const LegalityQuery &Q) {
         return Q.Types[0].isScalar() && Q.Types[1].getScalarSizeInBits() < 64;
       })
       .lower();
   getActionDefinitionsBuilder(G_ROTL).lower();
 
-  getActionDefinitionsBuilder({G_SBFX, G_UBFX}).customFor({s32, s64});
+  getActionDefinitionsBuilder({G_SBFX, G_UBFX}).customFor({{s32, s32}, {s64, s32}});
 
   computeTables();
   verify(*ST.getInstrInfo());
 }
 
 bool AArch64LegalizerInfo::legalizeCustom(LegalizerHelper &Helper,
                                           MachineInstr &MI) const {
   MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;

llvm/test/CodeGen/AArch64/GlobalISel/form-bitfield-extract-from-sextinreg.mir

 body:             |
   bb.0:
     liveins: $w0
     ; CHECK-LABEL: name: sextinreg_ashr_to_sbfx
     ; CHECK: liveins: $w0
     ; CHECK: %x:_(s32) = COPY $w0
     ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 5
     ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 14
-    ; CHECK: %sext_inreg:_(s32) = G_SBFX %x, [[C]], [[C1]]
+    ; CHECK: %sext_inreg:_(s32) = G_SBFX %x, [[C]](s32), [[C1]]
     ; CHECK: $w0 = COPY %sext_inreg(s32)
     ; CHECK: RET_ReallyLR implicit $w0
     %x:_(s32) = COPY $w0
     %lsb:_(s32) = G_CONSTANT i32 5
     %shift:_(s32) = G_ASHR %x, %lsb
     %sext_inreg:_(s32) = G_SEXT_INREG %shift, 10
     $w0 = COPY %sext_inreg
     RET_ReallyLR implicit $w0
 
 ...
 ---
 name:            sextinreg_lshr_to_sbfx
 tracksRegLiveness: true
 legalized: true
 body:             |
   bb.0:
     liveins: $w0
     ; CHECK-LABEL: name: sextinreg_lshr_to_sbfx
     ; CHECK: liveins: $w0
     ; CHECK: %x:_(s32) = COPY $w0
     ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 5
     ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 14
-    ; CHECK: %sext_inreg:_(s32) = G_SBFX %x, [[C]], [[C1]]
+    ; CHECK: %sext_inreg:_(s32) = G_SBFX %x, [[C]](s32), [[C1]]
     ; CHECK: $w0 = COPY %sext_inreg(s32)
     ; CHECK: RET_ReallyLR implicit $w0
     %x:_(s32) = COPY $w0
     %lsb:_(s32) = G_CONSTANT i32 5
     %shift:_(s32) = G_LSHR %x, %lsb
     %sext_inreg:_(s32) = G_SEXT_INREG %shift, 10
     $w0 = COPY %sext_inreg
     RET_ReallyLR implicit $w0

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

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -mtriple aarch64 -run-pass=legalizer -verify-machineinstrs %s -o - | FileCheck %s
 ...
 ---
 name:            s32
 body: |
   bb.0.entry:
     liveins: $w0
     ; CHECK-LABEL: name: s32
     ; CHECK: %copy:_(s32) = COPY $w0
     ; CHECK: %lsb:_(s32) = G_CONSTANT i32 1
     ; CHECK: %width:_(s32) = G_CONSTANT i32 2
-    ; CHECK: %sbfx:_(s32) = G_SBFX %copy, %lsb, %width
+    ; CHECK: %sbfx:_(s32) = G_SBFX %copy, %lsb(s32), %width
     ; CHECK: $w0 = COPY %sbfx(s32)
     %copy:_(s32) = COPY $w0
     %lsb:_(s32) = G_CONSTANT i32 1
     %width:_(s32) = G_CONSTANT i32 2
     %sbfx:_(s32) = G_SBFX %copy(s32), %lsb, %width
     $w0 = COPY %sbfx(s32)
 ...
 ---
 name:            s64
 body: |
   bb.0.entry:
     liveins: $x0
     ; CHECK-LABEL: name: s64
     ; CHECK: %copy:_(s64) = COPY $x0
-    ; CHECK: %lsb:_(s64) = G_CONSTANT i64 1
-    ; CHECK: %width:_(s64) = G_CONSTANT i64 2
-    ; CHECK: %sbfx:_(s64) = G_SBFX %copy, %lsb, %width
+    ; CHECK: %lsb:_(s32) = G_CONSTANT i32 1
+    ; CHECK: %width:_(s32) = G_CONSTANT i32 2
+    ; CHECK: %sbfx:_(s64) = G_SBFX %copy, %lsb(s32), %width
     ; CHECK: $x0 = COPY %sbfx(s64)
     %copy:_(s64) = COPY $x0
-    %lsb:_(s64) = G_CONSTANT i64 1
-    %width:_(s64) = G_CONSTANT i64 2
+    %lsb:_(s32) = G_CONSTANT i32 1
+    %width:_(s32) = G_CONSTANT i32 2
     %sbfx:_(s64) = G_SBFX %copy(s64), %lsb, %width
     $x0 = COPY %sbfx(s64)
 ...

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

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -mtriple aarch64 -run-pass=legalizer -verify-machineinstrs %s -o - | FileCheck %s
 ...
 ---
 name:            s32
 body: |
   bb.0.entry:
     liveins: $w0
     ; CHECK-LABEL: name: s32
     ; CHECK: %copy:_(s32) = COPY $w0
     ; CHECK: %lsb:_(s32) = G_CONSTANT i32 1
     ; CHECK: %width:_(s32) = G_CONSTANT i32 2
-    ; CHECK: %ubfx:_(s32) = G_UBFX %copy, %lsb, %width
+    ; CHECK: %ubfx:_(s32) = G_UBFX %copy, %lsb(s32), %width
     ; CHECK: $w0 = COPY %ubfx(s32)
     %copy:_(s32) = COPY $w0
     %lsb:_(s32) = G_CONSTANT i32 1
     %width:_(s32) = G_CONSTANT i32 2
     %ubfx:_(s32) = G_UBFX %copy(s32), %lsb, %width
     $w0 = COPY %ubfx(s32)
 ...
 ---
 name:            s64
 body: |
   bb.0.entry:
     liveins: $x0
     ; CHECK-LABEL: name: s64
     ; CHECK: %copy:_(s64) = COPY $x0
-    ; CHECK: %lsb:_(s64) = G_CONSTANT i64 1
-    ; CHECK: %width:_(s64) = G_CONSTANT i64 2
-    ; CHECK: %ubfx:_(s64) = G_UBFX %copy, %lsb, %width
+    ; CHECK: %lsb:_(s32) = G_CONSTANT i32 1
+    ; CHECK: %width:_(s32) = G_CONSTANT i32 2
+    ; CHECK: %ubfx:_(s64) = G_UBFX %copy, %lsb(s32), %width
     ; CHECK: $x0 = COPY %ubfx(s64)
     %copy:_(s64) = COPY $x0
-    %lsb:_(s64) = G_CONSTANT i64 1
-    %width:_(s64) = G_CONSTANT i64 2
+    %lsb:_(s32) = G_CONSTANT i32 1
+    %width:_(s32) = G_CONSTANT i32 2
     %ubfx:_(s64) = G_UBFX %copy(s64), %lsb, %width
     $x0 = COPY %ubfx(s64)
 ...

llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir

 # DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: G_VECREDUCE_UMAX (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: G_VECREDUCE_UMIN (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined
-# DEBUG-NEXT: G_SBFX (opcode {{[0-9]+}}): 1 type index, 0 imm indices
-# DEBUG-NEXT: .. the first uncovered type index: 1, OK
+# DEBUG-NEXT: G_SBFX (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
+# DEBUG-NEXT: .. the first uncovered type index: 2, OK
 # DEBUG-NEXT: .. the first uncovered imm index: 0, OK
-# DEBUG-NEXT: G_UBFX (opcode {{[0-9]+}}): 1 type index, 0 imm indices
+# DEBUG-NEXT: G_UBFX (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
 # DEBUG-NEXT: .. opcode {{[0-9]+}} is aliased to {{[0-9]+}}
-# DEBUG-NEXT: .. the first uncovered type index: 1, OK
+# DEBUG-NEXT: .. the first uncovered type index: 2, OK
 # DEBUG-NEXT: .. the first uncovered imm index: 0, OK
 # CHECK-NOT: ill-defined
 
 ---
 name: dummy
 body: |
   bb.0:
 ...

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

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
+    %cst1:gpr(s32) = G_CONSTANT i32 0
+    %cst2:gpr(s32) = G_CONSTANT i32 10
     %sbfx:gpr(s64) = G_SBFX %copy, %cst1, %cst2
     $x0 = COPY %sbfx
     RET_ReallyLR implicit $x0
 
 ...
 ---
 name:            sbfx_s32_31_31
 legalized:       true

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

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
+    %cst1:gpr(s32) = G_CONSTANT i32 0
+    %cst2:gpr(s32) = G_CONSTANT i32 10
     %ubfx:gpr(s64) = G_UBFX %copy, %cst1, %cst2
     $x0 = COPY %ubfx
     RET_ReallyLR implicit $x0
 
 ...
 ---
 name:            ubfx_s32_31_31
 legalized:       true

llvm/unittests/CodeGen/GlobalISel/MachineIRBuilderTest.cpp

TEST_F(AArch64GISelMITest, BuildAddoSubo) {
   EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
 }
 
 TEST_F(AArch64GISelMITest, BuildBitfieldExtract) {
   setUp();
   if (!TM)
     return;
   LLT S64 = LLT::scalar(64);
+  LLT S32 = LLT::scalar(32);
   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]);
+  auto C0 = B.buildConstant(S32, 1);
+  auto C1 = B.buildConstant(S32, 10);
+
+  auto Ubfx = B.buildUbfx(S64, Copies[0], C0, C1);
+  B.buildSbfx(S64, Ubfx, C0, C1);
 
   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]]:_
+  ; CHECK: [[CONST0:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+  ; CHECK: [[CONST1:%[0-9]+]]:_(s32) = G_CONSTANT i32 10
+  ; CHECK: [[UBFX:%[0-9]+]]:_(s64) = G_UBFX [[COPY0]]:_, [[CONST0]]:_(s32), [[CONST1]]:_
+  ; CHECK: [[SBFX:%[0-9]+]]:_(s64) = G_SBFX [[UBFX]]:_, [[CONST0]]:_(s32), [[CONST1]]:_
   )";
 
   EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
 }