[AArch64][GlobalISel] Legalize bswap <2 x i16>

[jroelofs]

marked as "WIP" because there are a few extraneous copies:

$ cat bswap.ll
; RUN: llc -mtriple=arm64-apple-ios < %s | FileCheck %s

define void @test1(<2 x i16>* %p) {
  %in = load <2 x i16>, <2 x i16>* %p
  %out = call <2 x i16> @llvm.bswap.v2i16(<2 x i16> %in)
  store <2 x i16> %out, <2 x i16>* %p
  ret void
}

declare <2 x i16> @llvm.bswap.v2i16(<2 x i16>) nounwind readnone

GISel:

$ ./bin/llc -global-isel=1 -global-isel-abort=1 bswap.ll -mtriple=arm64-apple-ios -o -
	.section	__TEXT,__text,regular,pure_instructions
	.globl	_test1                          ; -- Begin function test1
	.p2align	2
_test1:                                 ; @test1
	.cfi_startproc
; %bb.0:
	ldr	h0, [x0]
	ldr	h1, [x0, #2]
	mov.h	v0[1], v1[0]
	fmov	w8, s0
	mov.s	v0[0], w8
	rev32.8b	v0, v0
	ushr.2s	v0, v0, #16
	fmov	s0, s0
	mov	h1, v0[1]
	str	h0, [x0]
	str	h1, [x0, #2]
	ret
	.cfi_endproc
                                        ; -- End function
.subsections_via_symbols

SDAG:

$ ./bin/llc -global-isel=0 -global-isel-abort=1 bswap.ll -mtriple=arm64-apple-ios -o -
	.section	__TEXT,__text,regular,pure_instructions
	.globl	_test1                          ; -- Begin function test1
	.p2align	2
_test1:                                 ; @test1
	.cfi_startproc
; %bb.0:
	ld1.h	{ v0 }[0], [x0]
	add	x8, x0, #2                      ; =2
	ld1.h	{ v0 }[2], [x8]
	rev32.8b	v0, v0
	ushr.2s	v0, v0, #16
	mov.s	w8, v0[1]
	fmov	w9, s0
	strh	w9, [x0]
	strh	w8, [x0, #2]
	ret
	.cfi_endproc
                                        ; -- End function
.subsections_via_symbols

[paquette]

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

1017

Can you use G_REV32 here instead? (see: AArch64InstrGISel.td)

I think it's preferable to keep the results of legalization as generic as possible.

1022

Is it possible to use a generic instruction here?

[jroelofs]

Closer.

Use more generic opcodes during legalization.

[jroelofs]

fixup test checks

[paquette]

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

3999

This looks similar to getSubRegForClass? Maybe it's possible to share some code there?

[paquette]

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

3999

Also, it'd be nice to have a specific test for G_UNMERGE_VALUES selection. I think that this part + that test could be moved into an independent patch.

[jroelofs]

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

3999

good idea

3999

https://reviews.llvm.org/D106007

[jroelofs]

Getting closer:

SDAG:

$ ./bin/llc -global-isel=0 -global-isel-abort=1 bswap.ll -mtriple=arm64-apple-ios -o -
	.section	__TEXT,__text,regular,pure_instructions
	.globl	_test1                          ; -- Begin function test1
	.p2align	2
_test1:                                 ; @test1
	.cfi_startproc
; %bb.0:
	ld1.h	{ v0 }[0], [x0]
	add	x8, x0, #2                      ; =2
	ld1.h	{ v0 }[2], [x8]
	rev32.8b	v0, v0
	ushr.2s	v0, v0, #16
	mov.s	w8, v0[1]
	fmov	w9, s0
	strh	w9, [x0]
	strh	w8, [x0, #2]
	ret
	.cfi_endproc
                                        ; -- End function
.subsections_via_symbols

GISel:

$ ./bin/llc -global-isel=1 -global-isel-abort=1 bswap.ll -mtriple=arm64-apple-ios -o -
	.section	__TEXT,__text,regular,pure_instructions
	.globl	_test1                          ; -- Begin function test1
	.p2align	2
_test1:                                 ; @test1
	.cfi_startproc
; %bb.0:
	ldr	h0, [x0]
	ldr	h1, [x0, #2]
	mov.h	v0[1], v1[0]
	mov.s	v0[0], v0[0]
	rev32.8b	v0, v0
	ushr.2s	v0, v0, #16
	mov	h1, v0[1]
	str	h0, [x0]
	str	h1, [x0, #2]
	ret
	.cfi_endproc
                                        ; -- End function
.subsections_via_symbols

[jroelofs]

SDAG:

$ ./bin/llc -global-isel=0 -global-isel-abort=1 bswap.ll -mtriple=arm64-apple-ios -o -
	.section	__TEXT,__text,regular,pure_instructions
	.globl	_test1                          ; -- Begin function test1
	.p2align	2
_test1:                                 ; @test1
	.cfi_startproc
; %bb.0:
	ld1.h	{ v0 }[0], [x0]
	add	x8, x0, #2                      ; =2
	ld1.h	{ v0 }[2], [x8]
	rev32.8b	v0, v0
	ushr.2s	v0, v0, #16
	mov.s	w8, v0[1]
	fmov	w9, s0
	strh	w9, [x0]
	strh	w8, [x0, #2]
	ret
	.cfi_endproc
                                        ; -- End function
.subsections_via_symbols

GISel:

$ ./bin/llc -global-isel=1 -global-isel-abort=1 bswap.ll -mtriple=arm64-apple-ios -o -
	.section	__TEXT,__text,regular,pure_instructions
	.globl	_test1                          ; -- Begin function test1
	.p2align	2
_test1:                                 ; @test1
	.cfi_startproc
; %bb.0:
	ldr	h0, [x0]
	ldr	h1, [x0, #2]
	mov.h	v0[1], v1[0]
	rev32.8b	v0, v0
	ushr.2s	v0, v0, #16
	mov	h1, v0[1]
	str	h0, [x0]
	str	h1, [x0, #2]
	ret
	.cfi_endproc
                                        ; -- End function
.subsections_via_symbols

[paquette]

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

1009

Do you think you could add an explanation somewhere for why the v2s16 case is special?

[paquette]

Aside from the comment, I think this looks pretty good at this point? The codegen differences between SDAG and GISel seem unrelated to the bswap at this point.

[jroelofs]

Explain why <2 x half> is weird, and why we're not directly selecting the instructions we want during legalization.

[paquette]

Thanks!

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

     for (unsigned Idx = 0; Idx < NumInsertRegs; ++Idx) {
       Register ImpDefReg = MRI.createVirtualRegister(&AArch64::FPR128RegClass);
       MachineInstr &ImpDefMI =
           *BuildMI(MBB, I, I.getDebugLoc(), TII.get(TargetOpcode::IMPLICIT_DEF),
                    ImpDefReg);
 
       // Now, create the subregister insert from SrcReg.
       Register InsertReg = MRI.createVirtualRegister(&AArch64::FPR128RegClass);
+      auto SubImm = [&]() {
+        switch (WideTy.getScalarSizeInBits() * NumElts) {
+        default:
+          llvm_unreachable("unexpected subreg size");
+        case 16:
+          return AArch64::hsub;
+        case 32:
+          return AArch64::ssub;
+        case 64:
+          return AArch64::dsub;
+        }
+      }();
       MachineInstr &InsMI =
           *BuildMI(MBB, I, I.getDebugLoc(),
                    TII.get(TargetOpcode::INSERT_SUBREG), InsertReg)
                .addUse(ImpDefReg)
                .addUse(SrcReg)
-               .addImm(AArch64::dsub);
+               .addImm(SubImm);
 
       constrainSelectedInstRegOperands(ImpDefMI, TII, TRI, RBI);
       constrainSelectedInstRegOperands(InsMI, TII, TRI, RBI);
 
       // Save the register so that we can copy from it after.
       InsertRegs.push_back(InsertReg);
     }
   }

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

public:
   AArch64LegalizerInfo(const AArch64Subtarget &ST);
 
   bool legalizeCustom(LegalizerHelper &Helper, MachineInstr &MI) const override;
 
   bool legalizeIntrinsic(LegalizerHelper &Helper,
                          MachineInstr &MI) const override;
 
 private:
+  bool legalizeBSwap(MachineInstr &MI, MachineRegisterInfo &MRI,
+                     MachineIRBuilder &MIRBuilder) const;
   bool legalizeVaArg(MachineInstr &MI, MachineRegisterInfo &MRI,
                      MachineIRBuilder &MIRBuilder) const;
   bool legalizeLoadStore(MachineInstr &MI, MachineRegisterInfo &MRI,
                          MachineIRBuilder &MIRBuilder,
                          GISelChangeObserver &Observer) const;
   bool legalizeShlAshrLshr(MachineInstr &MI, MachineRegisterInfo &MRI,
                            MachineIRBuilder &MIRBuilder,
                            GISelChangeObserver &Observer) const;

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

AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
   getActionDefinitionsBuilder(G_PHI).legalFor({p0, s16, s32, s64})
       .legalFor(PackedVectorAllTypeList)
       .clampScalar(0, s16, s64)
       .widenScalarToNextPow2(0);
 
   getActionDefinitionsBuilder(G_BSWAP)
       .legalFor({s32, s64, v4s32, v2s32, v2s64})
       .clampScalar(0, s32, s64)
-      .widenScalarToNextPow2(0);
+      .widenScalarToNextPow2(0)
+      .customIf(typeIs(0, v2s16));
 
   getActionDefinitionsBuilder({G_ADD, G_SUB, G_MUL, G_AND, G_OR, G_XOR})
       .legalFor({s32, s64, v2s32, v4s32, v4s16, v8s16, v16s8, v8s8})
       .scalarizeIf(
           [=](const LegalityQuery &Query) {
             return Query.Opcode == G_MUL && Query.Types[0] == v2s64;
           },
           0)
bool AArch64LegalizerInfo::legalizeCustom(LegalizerHelper &Helper,
   default:
     // No idea what to do.
     return false;
   case TargetOpcode::G_VAARG:
     return legalizeVaArg(MI, MRI, MIRBuilder);
   case TargetOpcode::G_LOAD:
   case TargetOpcode::G_STORE:
     return legalizeLoadStore(MI, MRI, MIRBuilder, Observer);
+  case TargetOpcode::G_BSWAP:
+    return legalizeBSwap(MI, MRI, MIRBuilder);
   case TargetOpcode::G_SHL:
   case TargetOpcode::G_ASHR:
   case TargetOpcode::G_LSHR:
     return legalizeShlAshrLshr(MI, MRI, MIRBuilder, Observer);
   case TargetOpcode::G_GLOBAL_VALUE:
     return legalizeSmallCMGlobalValue(MI, MRI, MIRBuilder, Observer);
   case TargetOpcode::G_TRUNC:
     return legalizeVectorTrunc(MI, Helper);
bool AArch64LegalizerInfo::legalizeLoadStore(
   } else {
     auto NewLoad = MIRBuilder.buildLoad(NewTy, MI.getOperand(1), MMO);
     MIRBuilder.buildBitcast(ValReg, NewLoad);
   }
   MI.eraseFromParent();
   return true;
 }
 
+bool AArch64LegalizerInfo::legalizeBSwap(MachineInstr &MI,
+                                         MachineRegisterInfo &MRI,
+                                         MachineIRBuilder &MIRBuilder) const {
+  assert(MI.getOpcode() == TargetOpcode::G_BSWAP);
+
+  Register ValReg = MI.getOperand(1).getReg();
+  const LLT v4s16 = LLT::fixed_vector(4, 16);
+  const LLT v2s32 = LLT::fixed_vector(2, 32);
+  const LLT v8s8 = LLT::fixed_vector(8, 8);
+  const LLT s32 = LLT::scalar(32);
+  assert(LLT::fixed_vector(2, 16) == MRI.getType(ValReg));
+
+  auto Undef = MIRBuilder.buildUndef(v2s32);
+  auto Bitcast = MIRBuilder.buildBitcast(s32, ValReg);
+  auto Zero = MIRBuilder.buildConstant(s32, 0);
+  auto Insert =
+      MIRBuilder.buildInsertVectorElement(v2s32, Undef, Bitcast, Zero);
+  auto Bitcast3 = MIRBuilder.buildBitcast(v8s8, Insert);
+  auto Rev32 = MIRBuilder.buildInstr(AArch64::G_REV32, {v8s8}, {Bitcast3});
+  auto Bitcast2 = MIRBuilder.buildBitcast(v2s32, Rev32);
+  auto Amt = MIRBuilder.buildConstant(v2s32, 16);
+  auto UShr =
+      MIRBuilder.buildInstr(TargetOpcode::G_LSHR, {v2s32}, {Bitcast2, Amt});
+  auto Extract = MIRBuilder.buildExtractVectorElement(s32, UShr, Zero);
+  MIRBuilder.buildBitcast({MI.getOperand(0).getReg()}, Extract);
+  MI.eraseFromParent();
+  return true;
+}
+
 bool AArch64LegalizerInfo::legalizeVaArg(MachineInstr &MI,
                                          MachineRegisterInfo &MRI,
                                          MachineIRBuilder &MIRBuilder) const {
   MachineFunction &MF = MIRBuilder.getMF();
   Align Alignment(MI.getOperand(2).getImm());
   Register Dst = MI.getOperand(0).getReg();
   Register ListPtr = MI.getOperand(1).getReg();
 

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

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -march=aarch64 -run-pass=legalizer %s -o - -verify-machineinstrs | FileCheck %s
 --- |
   target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
   target triple = "aarch64"
 
   declare i16 @llvm.bswap.i16(i16) #0
 
   define i16 @bswap_s16(i16 %a) { ret i16 0 }
 
+  define <2 x i16> @bswap_2xi16(<2 x i16> %a) { ret <2 x i16> <i16 0, i16 0> }
+
   attributes #0 = { nounwind readnone speculatable willreturn }
 
 ...
 ---
 name:            bswap_s16
 alignment:       4
 tracksRegLiveness: true
 liveins:
bb.1:
     %1:_(s32) = COPY $w0
     %0:_(s16) = G_TRUNC %1(s32)
     %2:_(s16) = G_BSWAP %0
     %3:_(s32) = G_ANYEXT %2(s16)
     $w0 = COPY %3(s32)
     RET_ReallyLR implicit $w0
 
 ...
+---
+name:            bswap_2xi16
+alignment:       4
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: _ }
+  - { id: 1, class: _ }
+liveins:
+  - { reg: '$s0' }
+frameInfo:
+  maxAlignment:    1
+machineFunctionInfo: {}
+body:             |
+  bb.1:
+    liveins: $s0
+
+    ; CHECK-LABEL: name: bswap_2xi16
+    ; CHECK: liveins: $s0
+    ; CHECK: [[COPY:%[0-9]+]]:_(<2 x s16>) = COPY $s0
+    ; CHECK: [[DEF:%[0-9]+]]:_(<2 x s32>) = G_IMPLICIT_DEF
+    ; CHECK: [[BITCAST:%[0-9]+]]:_(s32) = G_BITCAST [[COPY]](<2 x s16>)
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK: [[IVEC:%[0-9]+]]:_(<2 x s32>) = G_INSERT_VECTOR_ELT [[DEF]], [[BITCAST]](s32), [[C]](s32)
+    ; CHECK: [[BITCAST1:%[0-9]+]]:_(<8 x s8>) = G_BITCAST [[IVEC]](<2 x s32>)
+    ; CHECK: [[REV32_:%[0-9]+]]:_(<8 x s8>) = G_REV32 [[BITCAST1]]
+    ; CHECK: [[BITCAST2:%[0-9]+]]:_(<2 x s32>) = G_BITCAST [[REV32_]](<8 x s8>)
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 16
+    ; CHECK: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x s32>) = G_BUILD_VECTOR [[C1]](s32), [[C1]](s32)
+    ; CHECK: [[LSHR:%[0-9]+]]:_(<2 x s32>) = G_LSHR [[BITCAST2]], [[BUILD_VECTOR]](<2 x s32>)
+    ; CHECK: [[C2:%[0-9]+]]:_(s64) = G_CONSTANT i64 0
+    ; CHECK: [[EVEC:%[0-9]+]]:_(s32) = G_EXTRACT_VECTOR_ELT [[LSHR]](<2 x s32>), [[C2]](s64)
+    ; CHECK: [[BITCAST3:%[0-9]+]]:_(<2 x s16>) = G_BITCAST [[EVEC]](s32)
+    ; CHECK: $s0 = COPY [[BITCAST3]](<2 x s16>)
+    ; CHECK: RET_ReallyLR
+    %0:_(<2 x s16>) = COPY $s0
+    %1:_(<2 x s16>) = G_BSWAP %0
+    $s0 = COPY %1(<2 x s16>)
+    RET_ReallyLR
+
+...