[AArch64][GlobalISel] Legalize scalar G_CTTZ + G_CTTZ_ZERO_UNDEF

[paquette]

Summary

This adds legalization for scalar G_CTTZ and G_CTTZ_ZERO_UNDEF. Vector support requires handling vector G_BITREVERSE, which I haven't gotten around to yet.

For G_CTTZ_ZERO_UNDEF, we just lower it to G_CTTZ.

For G_CTTZ, we match SelectionDAG's lowering to a G_BITREVERSE + G_CTLZ.

e.g. https://godbolt.org/z/nPEseYh1s

(With this patch, we have slightly worse codegen than SDAG for types smaller than s32; it seems like we're missing a combine.)

Also, this adds in a function to build G_BITREVERSE to MachineIRBuilder.

[jroelofs]

LGTM

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

   }
 
   /// Build and insert \p Dst = G_ROTL \p Src, \p Amt
   MachineInstrBuilder buildRotateLeft(const DstOp &Dst, const SrcOp &Src,
                                       const SrcOp &Amt) {
     return buildInstr(TargetOpcode::G_ROTL, {Dst}, {Src, Amt});
   }
 
+  /// Build and insert \p Dst = G_BITREVERSE \p Src
+  MachineInstrBuilder buildBitReverse(const DstOp &Dst, const SrcOp &Src) {
+    return buildInstr(TargetOpcode::G_BITREVERSE, {Dst}, {Src});
+  }
+
   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/lib/Target/AArch64/GISel/AArch64LegalizerInfo.h

private:
   bool legalizeBitfieldExtract(MachineInstr &MI, MachineRegisterInfo &MRI,
                                LegalizerHelper &Helper) const;
   bool legalizeRotate(MachineInstr &MI, MachineRegisterInfo &MRI,
                       LegalizerHelper &Helper) const;
   bool legalizeCTPOP(MachineInstr &MI, MachineRegisterInfo &MRI,
                      LegalizerHelper &Helper) const;
   bool legalizeAtomicCmpxchg128(MachineInstr &MI, MachineRegisterInfo &MRI,
                                 LegalizerHelper &Helper) const;
+  bool legalizeCTTZ(MachineInstr &MI, LegalizerHelper &Helper) const;
   const AArch64Subtarget *ST;
 };
 } // End llvm namespace.
 #endif

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

getActionDefinitionsBuilder(G_CTLZ)
       .legalForCartesianProduct(
           {s32, s64, v8s8, v16s8, v4s16, v8s16, v2s32, v4s32})
       .scalarize(1);
   getActionDefinitionsBuilder(G_CTLZ_ZERO_UNDEF).lower();
 
   // TODO: Custom lowering for v2s32, v4s32, v2s64.
   getActionDefinitionsBuilder(G_BITREVERSE).legalFor({s32, s64, v8s8, v16s8});
 
+  getActionDefinitionsBuilder(G_CTTZ_ZERO_UNDEF).lower();
+
+  // TODO: Handle vector types.
+  getActionDefinitionsBuilder(G_CTTZ)
+      .clampScalar(0, s32, s64)
+      .scalarSameSizeAs(1, 0)
+      .customFor({s32, s64});
+
   getActionDefinitionsBuilder(G_SHUFFLE_VECTOR)
       .legalIf([=](const LegalityQuery &Query) {
         const LLT &DstTy = Query.Types[0];
         const LLT &SrcTy = Query.Types[1];
         // For now just support the TBL2 variant which needs the source vectors
         // to be the same size as the dest.
         if (DstTy != SrcTy)
           return false;
bool AArch64LegalizerInfo::legalizeCustom(LegalizerHelper &Helper,
   case TargetOpcode::G_UBFX:
     return legalizeBitfieldExtract(MI, MRI, Helper);
   case TargetOpcode::G_ROTR:
     return legalizeRotate(MI, MRI, Helper);
   case TargetOpcode::G_CTPOP:
     return legalizeCTPOP(MI, MRI, Helper);
   case TargetOpcode::G_ATOMIC_CMPXCHG:
     return legalizeAtomicCmpxchg128(MI, MRI, Helper);
+  case TargetOpcode::G_CTTZ:
+    return legalizeCTTZ(MI, Helper);
   }
 
   llvm_unreachable("expected switch to return");
 }
 
 bool AArch64LegalizerInfo::legalizeRotate(MachineInstr &MI,
                                           MachineRegisterInfo &MRI,
                                           LegalizerHelper &Helper) const {
bool AArch64LegalizerInfo::legalizeAtomicCmpxchg128(
   constrainSelectedInstRegOperands(*CAS, *ST->getInstrInfo(),
                                    *MRI.getTargetRegisterInfo(),
                                    *ST->getRegBankInfo());
 
   MIRBuilder.buildMerge(MI.getOperand(0), {DstLo, DstHi});
   MI.eraseFromParent();
   return true;
 }
+
+bool AArch64LegalizerInfo::legalizeCTTZ(MachineInstr &MI,
+                                        LegalizerHelper &Helper) const {
+  MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;
+  MachineRegisterInfo &MRI = *MIRBuilder.getMRI();
+  LLT Ty = MRI.getType(MI.getOperand(1).getReg());
+  auto BitReverse = MIRBuilder.buildBitReverse(Ty, MI.getOperand(1));
+  MIRBuilder.buildCTLZ(MI.getOperand(0).getReg(), BitReverse);
+  MI.eraseFromParent();
+  return true;
+}

llvm/test/CodeGen/AArch64/GlobalISel/legalize-cttz-zero-undef.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64 -global-isel -verify-machineinstrs -run-pass=legalizer %s -o - | FileCheck %s
+
+...
+---
+name:            s8
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; CHECK-LABEL: name: s8
+    ; CHECK: liveins: $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 256
+    ; CHECK: [[OR:%[0-9]+]]:_(s32) = G_OR [[C]], [[C1]]
+    ; CHECK: [[BITREVERSE:%[0-9]+]]:_(s32) = G_BITREVERSE [[OR]]
+    ; CHECK: [[CTLZ:%[0-9]+]]:_(s32) = G_CTLZ [[BITREVERSE]](s32)
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY [[CTLZ]](s32)
+    ; CHECK: %ext:_(s32) = COPY [[COPY]](s32)
+    ; CHECK: $w0 = COPY %ext(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %val:_(s8) = G_IMPLICIT_DEF
+    %cttz:_(s8) = G_CTTZ_ZERO_UNDEF %val(s8)
+    %ext:_(s32) = G_ANYEXT %cttz(s8)
+    $w0 = COPY %ext(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            s16
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; CHECK-LABEL: name: s16
+    ; CHECK: liveins: $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 65536
+    ; CHECK: [[OR:%[0-9]+]]:_(s32) = G_OR [[C]], [[C1]]
+    ; CHECK: [[BITREVERSE:%[0-9]+]]:_(s32) = G_BITREVERSE [[OR]]
+    ; CHECK: [[CTLZ:%[0-9]+]]:_(s32) = G_CTLZ [[BITREVERSE]](s32)
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY [[CTLZ]](s32)
+    ; CHECK: %ext:_(s32) = COPY [[COPY]](s32)
+    ; CHECK: $w0 = COPY %ext(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %val:_(s16) = G_IMPLICIT_DEF
+    %cttz:_(s16) = G_CTTZ_ZERO_UNDEF %val(s16)
+    %ext:_(s32) = G_ANYEXT %cttz(s16)
+    $w0 = COPY %ext(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            s32
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; CHECK-LABEL: name: s32
+    ; CHECK: liveins: $w0
+    ; CHECK: %val:_(s32) = COPY $w0
+    ; CHECK: [[BITREVERSE:%[0-9]+]]:_(s32) = G_BITREVERSE %val
+    ; CHECK: [[CTLZ:%[0-9]+]]:_(s32) = G_CTLZ [[BITREVERSE]](s32)
+    ; CHECK: $w0 = COPY [[CTLZ]](s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %val:_(s32) = COPY $w0
+    %1:_(s32) = G_CTTZ_ZERO_UNDEF %val(s32)
+    $w0 = COPY %1(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            s64
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+
+    ; CHECK-LABEL: name: s64
+    ; CHECK: liveins: $x0
+    ; CHECK: %val:_(s64) = COPY $x0
+    ; CHECK: [[BITREVERSE:%[0-9]+]]:_(s64) = G_BITREVERSE %val
+    ; CHECK: [[CTLZ:%[0-9]+]]:_(s64) = G_CTLZ [[BITREVERSE]](s64)
+    ; CHECK: $x0 = COPY [[CTLZ]](s64)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %val:_(s64) = COPY $x0
+    %1:_(s64) = G_CTTZ_ZERO_UNDEF %val(s64)
+    $x0 = COPY %1(s64)
+    RET_ReallyLR implicit $x0
+
+...

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

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64 -global-isel -verify-machineinstrs -run-pass=legalizer %s -o - | FileCheck %s
+
+...
+---
+name:            s8
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; CHECK-LABEL: name: s8
+    ; CHECK: liveins: $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 256
+    ; CHECK: [[OR:%[0-9]+]]:_(s32) = G_OR [[C]], [[C1]]
+    ; CHECK: [[BITREVERSE:%[0-9]+]]:_(s32) = G_BITREVERSE [[OR]]
+    ; CHECK: [[CTLZ:%[0-9]+]]:_(s32) = G_CTLZ [[BITREVERSE]](s32)
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY [[CTLZ]](s32)
+    ; CHECK: %ext:_(s32) = COPY [[COPY]](s32)
+    ; CHECK: $w0 = COPY %ext(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %val:_(s8) = G_IMPLICIT_DEF
+    %cttz:_(s8) = G_CTTZ %val(s8)
+    %ext:_(s32) = G_ANYEXT %cttz(s8)
+    $w0 = COPY %ext(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            s16
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; CHECK-LABEL: name: s16
+    ; CHECK: liveins: $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 65536
+    ; CHECK: [[OR:%[0-9]+]]:_(s32) = G_OR [[C]], [[C1]]
+    ; CHECK: [[BITREVERSE:%[0-9]+]]:_(s32) = G_BITREVERSE [[OR]]
+    ; CHECK: [[CTLZ:%[0-9]+]]:_(s32) = G_CTLZ [[BITREVERSE]](s32)
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY [[CTLZ]](s32)
+    ; CHECK: %ext:_(s32) = COPY [[COPY]](s32)
+    ; CHECK: $w0 = COPY %ext(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %val:_(s16) = G_IMPLICIT_DEF
+    %cttz:_(s16) = G_CTTZ %val(s16)
+    %ext:_(s32) = G_ANYEXT %cttz(s16)
+    $w0 = COPY %ext(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            s32
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; CHECK-LABEL: name: s32
+    ; CHECK: liveins: $w0
+    ; CHECK: %val:_(s32) = COPY $w0
+    ; CHECK: [[BITREVERSE:%[0-9]+]]:_(s32) = G_BITREVERSE %val
+    ; CHECK: [[CTLZ:%[0-9]+]]:_(s32) = G_CTLZ [[BITREVERSE]](s32)
+    ; CHECK: $w0 = COPY [[CTLZ]](s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %val:_(s32) = COPY $w0
+    %1:_(s32) = G_CTTZ %val(s32)
+    $w0 = COPY %1(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            s64
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+
+    ; CHECK-LABEL: name: s64
+    ; CHECK: liveins: $x0
+    ; CHECK: %val:_(s64) = COPY $x0
+    ; CHECK: [[BITREVERSE:%[0-9]+]]:_(s64) = G_BITREVERSE %val
+    ; CHECK: [[CTLZ:%[0-9]+]]:_(s64) = G_CTLZ [[BITREVERSE]](s64)
+    ; CHECK: $x0 = COPY [[CTLZ]](s64)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %val:_(s64) = COPY $x0
+    %1:_(s64) = G_CTTZ %val(s64)
+    $x0 = COPY %1(s64)
+    RET_ReallyLR implicit $x0

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

 # DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: G_EXTRACT_VECTOR_ELT (opcode {{[0-9]+}}): 3 type indices, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: G_SHUFFLE_VECTOR (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: G_CTTZ (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: .. type index coverage check SKIPPED: user-defined predicate detected
+# DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: G_CTTZ_ZERO_UNDEF (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: .. type index coverage check SKIPPED: user-defined predicate detected
+# DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: G_CTLZ (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_CTLZ_ZERO_UNDEF (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: G_CTPOP (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: user-defined predicate detected