[AArch64][GlobalISel] Legalize ctpop for v2s64, v2s32, v4s32, v4s16, v8s16

[jroelofs]

Summary

https://llvm.godbolt.org/z/nTTK6M5qe

[paquette]

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

1213

doesn't SmallVector need an initial size?

1233

I think CTPOP is only used here? Maybe it's a little cleaner to move it and related stuff down near it.

1240

I don't think you should have to do this here, since we're before register bank selection?

1246

Logic here is getting a little complex. Maybe this?

if (Ty.isScalar() && Size == 64)
  MIRBuilder.buildZExt(Dst, UADD);
else
  UADD->getOperand(0).setReg(Dst);

1254

This shouldn't be necessary since we haven't selected regbanks yet.

[jroelofs]

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

1213

Sean Silva recently added a default to it: https://reviews.llvm.org/D92522

1233

I wanted to keep HAdds near where it was used.

1240

MachineVerifier seems to expect them just after legalization:

*** Bad machine code: Virtual register does not match instruction constraint ***
- function:    s32_lower
- basic block: %bb.0  (0x7fba998570c8)
- instruction: %ctpop:_(s32) = UADDLVv8i8v %4:_(<8 x s8>)
- operand 0:   %ctpop:_
Expect register class FPR16 but got nothing

1246

👍

[paquette]

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

1240

ah, yeah I guess it's cause it's not a generic instr. :/

Maybe it would make sense to add a G_UADDLV generic instruction to avoid emitting target instrs.

IIRC we import tablegen patterns for the uaddlv intrinsic allowing us to select it easily. Maybe if we had a generic instruction, we could do something similar.

[jroelofs]

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

1240

maybe I ought to just use the intrinsics then

[paquette]

LGTM

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

getActionDefinitionsBuilder(G_ROTR)
       })
       .lower();
   getActionDefinitionsBuilder(G_ROTL).lower();
 
   getActionDefinitionsBuilder({G_SBFX, G_UBFX})
       .customFor({{s32, s32}, {s64, s64}});
 
   // TODO: Custom legalization for s128
-  // TODO: v2s64, v2s32, v4s32, v4s16, v8s16
   // TODO: Use generic lowering when custom lowering is not possible.
   auto always = [=](const LegalityQuery &Q) { return true; };
   getActionDefinitionsBuilder(G_CTPOP)
       .legalFor({{v8s8, v8s8}, {v16s8, v16s8}})
       .clampScalar(0, s32, s128)
       .widenScalarToNextPow2(0)
       .minScalarEltSameAsIf(always, 1, 0)
       .maxScalarEltSameAsIf(always, 1, 0)
-      .customFor({{s32, s32}, {s64, s64}});
+      .customFor({{s32, s32},
+                  {s64, s64},
+                  {v2s64, v2s64},
+                  {v2s32, v2s32},
+                  {v4s32, v4s32},
+                  {v4s16, v4s16},
+                  {v8s16, v8s16}});
 
   getLegacyLegalizerInfo().computeTables();
   verify(*ST.getInstrInfo());
 }
 
 bool AArch64LegalizerInfo::legalizeCustom(LegalizerHelper &Helper,
                                           MachineInstr &MI) const {
   MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;
bool AArch64LegalizerInfo::legalizeCTPOP(MachineInstr &MI,
   // While there is no integer popcount instruction, it can
   // be more efficiently lowered to the following sequence that uses
   // AdvSIMD registers/instructions as long as the copies to/from
   // the AdvSIMD registers are cheap.
   //  FMOV    D0, X0        // copy 64-bit int to vector, high bits zero'd
   //  CNT     V0.8B, V0.8B  // 8xbyte pop-counts
   //  ADDV    B0, V0.8B     // sum 8xbyte pop-counts
   //  UMOV    X0, V0.B[0]   // copy byte result back to integer reg
+  //
+  // For 128 bit vector popcounts, we lower to the following sequence:
+  //  cnt.16b   v0, v0  // v8s16, v4s32, v2s64
+  //  uaddlp.8h v0, v0  // v8s16, v4s32, v2s64
+  //  uaddlp.4s v0, v0  //        v4s32, v2s64
+  //  uaddlp.2d v0, v0  //               v2s64
+  //
+  // For 64 bit vector popcounts, we lower to the following sequence:
+  //  cnt.8b    v0, v0  // v4s16, v2s32
+  //  uaddlp.4h v0, v0  // v4s16, v2s32
+  //  uaddlp.2s v0, v0  //        v2s32
+
   if (!ST->hasNEON() ||
       MI.getMF()->getFunction().hasFnAttribute(Attribute::NoImplicitFloat))
     return false;
   MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;
   Register Dst = MI.getOperand(0).getReg();
   Register Val = MI.getOperand(1).getReg();
   LLT Ty = MRI.getType(Val);
 
-  // TODO: Handle vector types.
-  assert(!Ty.isVector() && "Vector types not handled yet!");
   assert(Ty == MRI.getType(Dst) &&
          "Expected src and dst to have the same type!");
-  // TODO: Handle s128.
   unsigned Size = Ty.getSizeInBits();
+
+  // Pre-conditioning: widen Val up to the nearest vector type.
+  // s32,s64,v4s16,v2s32 -> v8i8
+  // v8s16,v4s32,v2s64 -> v16i8
+  LLT VTy = Size == 128 ? LLT::fixed_vector(16, 8) : LLT::fixed_vector(8, 8);
+  if (Ty.isScalar()) {
+    // TODO: Handle s128.
     assert((Size == 32 || Size == 64) && "Expected only 32 or 64 bit scalars!");
-  if (Size == 32)
+    if (Size == 32) {
       Val = MIRBuilder.buildZExt(LLT::scalar(64), Val).getReg(0);
-  const LLT V8S8 = LLT::fixed_vector(8, LLT::scalar(8));
-  Val = MIRBuilder.buildBitcast(V8S8, Val).getReg(0);
-  auto CTPOP = MIRBuilder.buildCTPOP(V8S8, Val);
-  auto UADDLV =
-      MIRBuilder
-          .buildIntrinsic(Intrinsic::aarch64_neon_uaddlv, {LLT::scalar(32)},
-                          /*HasSideEffects = */ false)
-          .addUse(CTPOP.getReg(0));
-  if (Size == 64)
-    MIRBuilder.buildZExt(Dst, UADDLV);
+    }
+  }
+  Val = MIRBuilder.buildBitcast(VTy, Val).getReg(0);
+
+  // Count bits in each byte-sized lane.
+  auto CTPOP = MIRBuilder.buildCTPOP(VTy, Val);
+
+  // Sum across lanes.
+  Register HSum = CTPOP.getReg(0);
+  unsigned Opc;
+  SmallVector<LLT> HAddTys;
+  if (Ty.isScalar()) {
+    Opc = Intrinsic::aarch64_neon_uaddlv;
+    HAddTys.push_back(LLT::scalar(32));
+  } else if (Ty == LLT::fixed_vector(8, 16)) {
+    Opc = Intrinsic::aarch64_neon_uaddlp;
+    HAddTys.push_back(LLT::fixed_vector(8, 16));
+  } else if (Ty == LLT::fixed_vector(4, 32)) {
+    Opc = Intrinsic::aarch64_neon_uaddlp;
+    HAddTys.push_back(LLT::fixed_vector(8, 16));
+    HAddTys.push_back(LLT::fixed_vector(4, 32));
+  } else if (Ty == LLT::fixed_vector(2, 64)) {
+    Opc = Intrinsic::aarch64_neon_uaddlp;
+    HAddTys.push_back(LLT::fixed_vector(8, 16));
+    HAddTys.push_back(LLT::fixed_vector(4, 32));
+    HAddTys.push_back(LLT::fixed_vector(2, 64));
+  } else if (Ty == LLT::fixed_vector(4, 16)) {
+    Opc = Intrinsic::aarch64_neon_uaddlp;
+    HAddTys.push_back(LLT::fixed_vector(4, 16));
+  } else if (Ty == LLT::fixed_vector(2, 32)) {
+    Opc = Intrinsic::aarch64_neon_uaddlp;
+    HAddTys.push_back(LLT::fixed_vector(4, 16));
+    HAddTys.push_back(LLT::fixed_vector(2, 32));
+  } else
+    llvm_unreachable("unexpected vector shape");
+  MachineInstrBuilder UADD;
+  for (LLT HTy : HAddTys) {
+    UADD = MIRBuilder.buildIntrinsic(Opc, {HTy}, /*HasSideEffects =*/false)
+                     .addUse(HSum);
+    HSum = UADD.getReg(0);
+  }
+
+  // Post-conditioning.
+  if (Ty.isScalar() && Size == 64)
+    MIRBuilder.buildZExt(Dst, UADD);
   else
-    UADDLV->getOperand(0).setReg(Dst);
+    UADD->getOperand(0).setReg(Dst);
   MI.eraseFromParent();
   return true;
 }
 
 bool AArch64LegalizerInfo::legalizeAtomicCmpxchg128(
     MachineInstr &MI, MachineRegisterInfo &MRI, LegalizerHelper &Helper) const {
   MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;
   LLT s64 = LLT::scalar(64);
   auto Addr = MI.getOperand(1).getReg();
   auto DesiredI = MIRBuilder.buildUnmerge({s64, s64}, MI.getOperand(2));
   auto NewI = MIRBuilder.buildUnmerge({s64, s64}, MI.getOperand(3));
   auto DstLo = MRI.createGenericVirtualRegister(s64);
   auto DstHi = MRI.createGenericVirtualRegister(s64);
 
   MachineInstrBuilder CAS;
   if (ST->hasLSE()) {
     // We have 128-bit CASP instructions taking XSeqPair registers, which are
     // s128. We need the merge/unmerge to bracket the expansion and pair up with
     // the rest of the MIR so we must reassemble the extracted registers into a
     // 128-bit known-regclass one with code like this:
     //
     //     %in1 = REG_SEQUENCE Lo, Hi    ; One for each input
     //     %out = CASP %in1, ...
     //     %OldLo = G_EXTRACT %out, 0
     //     %OldHi = G_EXTRACT %out, 64
     auto Ordering = (*MI.memoperands_begin())->getMergedOrdering();
     unsigned Opcode;
     switch (Ordering) {
     case AtomicOrdering::Acquire:
       Opcode = AArch64::CASPAX;
       break;
     case AtomicOrdering::Release:
       Opcode = AArch64::CASPLX;
       break;
     case AtomicOrdering::AcquireRelease:
     case AtomicOrdering::SequentiallyConsistent:
       Opcode = AArch64::CASPALX;
       break;
     default:
       Opcode = AArch64::CASPX;
       break;
     }
 
     LLT s128 = LLT::scalar(128);
     auto CASDst = MRI.createGenericVirtualRegister(s128);
     auto CASDesired = MRI.createGenericVirtualRegister(s128);
     auto CASNew = MRI.createGenericVirtualRegister(s128);
     MIRBuilder.buildInstr(TargetOpcode::REG_SEQUENCE, {CASDesired}, {})
         .addUse(DesiredI->getOperand(0).getReg())
         .addImm(AArch64::sube64)
         .addUse(DesiredI->getOperand(1).getReg())
         .addImm(AArch64::subo64);
     MIRBuilder.buildInstr(TargetOpcode::REG_SEQUENCE, {CASNew}, {})
         .addUse(NewI->getOperand(0).getReg())
         .addImm(AArch64::sube64)
         .addUse(NewI->getOperand(1).getReg())
         .addImm(AArch64::subo64);
 
     CAS = MIRBuilder.buildInstr(Opcode, {CASDst}, {CASDesired, CASNew, Addr});
 
     MIRBuilder.buildExtract({DstLo}, {CASDst}, 0);
     MIRBuilder.buildExtract({DstHi}, {CASDst}, 64);
   } else {
     // The -O0 CMP_SWAP_128 is friendlier to generate code for because LDXP/STXP

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

bb.0:
     ; CHECK: $w0 = COPY %ext(s32)
     ; CHECK: RET_ReallyLR implicit $w0
     %copy:_(s32) = COPY $w0
     %trunc:_(s8) = G_TRUNC %copy(s32)
     %ctpop:_(s16) = G_CTPOP %trunc(s8)
     %ext:_(s32) = G_ANYEXT %ctpop(s16)
     $w0 = COPY %ext(s32)
     RET_ReallyLR implicit $w0
+
+...
+---
+name:            custom_8x16
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $q0
+
+    ; CHECK-LABEL: name: custom_8x16
+    ; CHECK: liveins: $q0
+    ; CHECK: [[COPY:%[0-9]+]]:_(<8 x s16>) = COPY $q0
+    ; CHECK: [[BITCAST:%[0-9]+]]:_(<16 x s8>) = G_BITCAST [[COPY]](<8 x s16>)
+    ; CHECK: [[CTPOP:%[0-9]+]]:_(<16 x s8>) = G_CTPOP [[BITCAST]](<16 x s8>)
+    ; CHECK: [[INT:%[0-9]+]]:_(<8 x s16>) = G_INTRINSIC intrinsic(@llvm.aarch64.neon.uaddlp), [[CTPOP]](<16 x s8>)
+    ; CHECK: $q0 = COPY [[INT]](<8 x s16>)
+    ; CHECK: RET_ReallyLR implicit $q0
+    %0:_(<8 x s16>) = COPY $q0
+    %1:_(<8 x s16>) = G_CTPOP %0(<8 x s16>)
+    $q0 = COPY %1(<8 x s16>)
+    RET_ReallyLR implicit $q0
+
+...
+---
+name:            custom_4x32
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $q0
+
+    ; CHECK-LABEL: name: custom_4x32
+    ; CHECK: liveins: $q0
+    ; CHECK: [[COPY:%[0-9]+]]:_(<4 x s32>) = COPY $q0
+    ; CHECK: [[BITCAST:%[0-9]+]]:_(<16 x s8>) = G_BITCAST [[COPY]](<4 x s32>)
+    ; CHECK: [[CTPOP:%[0-9]+]]:_(<16 x s8>) = G_CTPOP [[BITCAST]](<16 x s8>)
+    ; CHECK: [[INT:%[0-9]+]]:_(<8 x s16>) = G_INTRINSIC intrinsic(@llvm.aarch64.neon.uaddlp), [[CTPOP]](<16 x s8>)
+    ; CHECK: [[INT1:%[0-9]+]]:_(<4 x s32>) = G_INTRINSIC intrinsic(@llvm.aarch64.neon.uaddlp), [[INT]](<8 x s16>)
+    ; CHECK: $q0 = COPY [[INT1]](<4 x s32>)
+    ; CHECK: RET_ReallyLR implicit $q0
+    %0:_(<4 x s32>) = COPY $q0
+    %1:_(<4 x s32>) = G_CTPOP %0(<4 x s32>)
+    $q0 = COPY %1(<4 x s32>)
+    RET_ReallyLR implicit $q0
+
+...
+---
+name:            custom_2x64
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $q0
+
+    ; CHECK-LABEL: name: custom_2x64
+    ; CHECK: liveins: $q0
+    ; CHECK: [[COPY:%[0-9]+]]:_(<2 x s64>) = COPY $q0
+    ; CHECK: [[BITCAST:%[0-9]+]]:_(<16 x s8>) = G_BITCAST [[COPY]](<2 x s64>)
+    ; CHECK: [[CTPOP:%[0-9]+]]:_(<16 x s8>) = G_CTPOP [[BITCAST]](<16 x s8>)
+    ; CHECK: [[INT:%[0-9]+]]:_(<8 x s16>) = G_INTRINSIC intrinsic(@llvm.aarch64.neon.uaddlp), [[CTPOP]](<16 x s8>)
+    ; CHECK: [[INT1:%[0-9]+]]:_(<4 x s32>) = G_INTRINSIC intrinsic(@llvm.aarch64.neon.uaddlp), [[INT]](<8 x s16>)
+    ; CHECK: [[INT2:%[0-9]+]]:_(<2 x s64>) = G_INTRINSIC intrinsic(@llvm.aarch64.neon.uaddlp), [[INT1]](<4 x s32>)
+    ; CHECK: $q0 = COPY [[INT2]](<2 x s64>)
+    ; CHECK: RET_ReallyLR implicit $q0
+    %0:_(<2 x s64>) = COPY $q0
+    %1:_(<2 x s64>) = G_CTPOP %0(<2 x s64>)
+    $q0 = COPY %1(<2 x s64>)
+    RET_ReallyLR implicit $q0
+
+...
+---
+name:            custom_4x16
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $d0
+
+    ; CHECK-LABEL: name: custom_4x16
+    ; CHECK: liveins: $d0
+    ; CHECK: [[COPY:%[0-9]+]]:_(<4 x s16>) = COPY $d0
+    ; CHECK: [[BITCAST:%[0-9]+]]:_(<8 x s8>) = G_BITCAST [[COPY]](<4 x s16>)
+    ; CHECK: [[CTPOP:%[0-9]+]]:_(<8 x s8>) = G_CTPOP [[BITCAST]](<8 x s8>)
+    ; CHECK: [[INT:%[0-9]+]]:_(<4 x s16>) = G_INTRINSIC intrinsic(@llvm.aarch64.neon.uaddlp), [[CTPOP]](<8 x s8>)
+    ; CHECK: $d0 = COPY [[INT]](<4 x s16>)
+    ; CHECK: RET_ReallyLR implicit $d0
+    %0:_(<4 x s16>) = COPY $d0
+    %1:_(<4 x s16>) = G_CTPOP %0(<4 x s16>)
+    $d0 = COPY %1(<4 x s16>)
+    RET_ReallyLR implicit $d0
+
+...
+---
+name:            custom_2x32
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $d0
+
+    ; CHECK-LABEL: name: custom_2x32
+    ; CHECK: liveins: $d0
+    ; CHECK: [[COPY:%[0-9]+]]:_(<2 x s32>) = COPY $d0
+    ; CHECK: [[BITCAST:%[0-9]+]]:_(<8 x s8>) = G_BITCAST [[COPY]](<2 x s32>)
+    ; CHECK: [[CTPOP:%[0-9]+]]:_(<8 x s8>) = G_CTPOP [[BITCAST]](<8 x s8>)
+    ; CHECK: [[INT:%[0-9]+]]:_(<4 x s16>) = G_INTRINSIC intrinsic(@llvm.aarch64.neon.uaddlp), [[CTPOP]](<8 x s8>)
+    ; CHECK: [[INT1:%[0-9]+]]:_(<2 x s32>) = G_INTRINSIC intrinsic(@llvm.aarch64.neon.uaddlp), [[INT]](<4 x s16>)
+    ; CHECK: $d0 = COPY [[INT1]](<2 x s32>)
+    ; CHECK: RET_ReallyLR implicit $d0
+    %0:_(<2 x s32>) = COPY $d0
+    %1:_(<2 x s32>) = G_CTPOP %0(<2 x s32>)
+    $d0 = COPY %1(<2 x s32>)
+    RET_ReallyLR implicit $d0