[AArch64] SelectionDag codegen for gpr CTZ instruction

[stuij]

Summary

When feature CSSC is available we should use instruction CTZ in SelectionDag
where applicable:

• CTTZ intrinsics are lowered to using the gpr CTZ instruction
• BITREVERSE -> CTLZ instruction pattern gets replaced by CTZ

spec:
https://developer.arm.com/documentation/ddi0602/2022-09/Base-Instructions/CTZ--Count-Trailing-Zeros-

[stuij]

removing comments

[stuij]

removed comment as it duplicates the inline comment

llvm/lib/Target/AArch64/AArch64InstrInfo.td

8499

There's a question here if we should move these into dagcombiner and implement likewise for GlobalIsel, or if we should leave this here and have them be picked up by SelDag and GlobalIsel:

[lenary]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

8499

We've talked already and my position is that these should be in dagcombiner/globalisel combines, but I'm interested in the views of e.g. @paquette @RKSimon. I realise combines like this might have to be behind hooks given that on some targets, cttz/ctlz of zero is undefined.

llvm/test/CodeGen/AArch64/gpr_cttz.ll

146

how does this testcase have no CHECK lines?

[RKSimon]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

8499

That's why we have support for cttz_zero_undef/G_CTTZ_ZERO_UNDEF etc - legalization should use the standard 'zero -> bw' case if the target supports it, but there's also the 'zero -> undef' opcode for targets that only support that.

[stuij]

removed instrInfo pattern match of (cttz (bitreverse -> cttz with equivalent target dagcombine

[stuij]

added missing testcase

[stuij]

llvm/test/CodeGen/AArch64/gpr_cttz.ll

146

I... don't know.
Added, thanks!

[lenary]

Cool, this works for me, and we can move the target dagcombine into target-independent dagcombiner later :)

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 //===-- AArch64ISelLowering.cpp - AArch64 DAG Lowering Implementation  ----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
   setOperationAction(ISD::UMUL_LOHI, MVT::i64, Expand);
   setOperationAction(ISD::SMUL_LOHI, MVT::i64, Expand);
 
   if (Subtarget->hasCSSC()) {
     setOperationAction(ISD::CTPOP, MVT::i32, Legal);
     setOperationAction(ISD::CTPOP, MVT::i64, Legal);
     setOperationAction(ISD::CTPOP, MVT::i128, Expand);
     setOperationAction(ISD::PARITY, MVT::i128, Expand);
+    setOperationAction(ISD::CTTZ, MVT::i32, Legal);
+    setOperationAction(ISD::CTTZ, MVT::i64, Legal);
+    setOperationAction(ISD::CTTZ, MVT::i128, Expand);
   } else {
     setOperationAction(ISD::CTPOP, MVT::i32, Custom);
     setOperationAction(ISD::CTPOP, MVT::i64, Custom);
     setOperationAction(ISD::CTPOP, MVT::i128, Custom);
 
     setOperationAction(ISD::PARITY, MVT::i64, Custom);
     setOperationAction(ISD::PARITY, MVT::i128, Custom);
   }

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 // 2022 Architecture Extensions:
 //===----------------------------===//
 
 //===----------------------------------------------------------------------===//
 // General Data-Processing Instructions (FEAT_V94_DP)
 //===----------------------------------------------------------------------===//
 defm ABS : OneOperandData<0b001000, "abs">, Requires<[HasCSSC]>;
 defm CNT : OneOperandData<0b000111, "cnt", ctpop>, Requires<[HasCSSC]>;
-defm CTZ : OneOperandData<0b000110, "ctz">, Requires<[HasCSSC]>;
+defm CTZ : OneOperandData<0b000110, "ctz", cttz>, Requires<[HasCSSC]>;
+
+let Predicates = [HasCSSC] in {
+  def : Pat<(i32 (ctlz (bitreverse GPR32:$Rn))),
+            (CTZWr GPR32:$Rn)>;
+  def : Pat<(i64 (ctlz (bitreverse GPR64:$Rn))),
+            (CTZXr GPR64:$Rn)>;
+}
 
 defm SMAX : ComparisonOp<0, 0, "smax">, Requires<[HasCSSC]>;
 defm SMIN : ComparisonOp<0, 1, "smin">, Requires<[HasCSSC]>;
 defm UMAX : ComparisonOp<1, 0, "umax">, Requires<[HasCSSC]>;
 defm UMIN : ComparisonOp<1, 1, "umin">, Requires<[HasCSSC]>;
 
 def RPRFM:
     I<(outs), (ins rprfop:$Rt, GPR64:$Rm, GPR64sp:$Rn),

llvm/test/CodeGen/AArch64/gpr_cttz.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -verify-machineinstrs -o - %s -mtriple=aarch64-linux-gnu | FileCheck %s
+; RUN: llc -verify-machineinstrs -o - %s -mtriple=aarch64-linux-gnu -mattr=+cssc | FileCheck %s -check-prefix=CHECK-CSSC
+
+define i4 @cttz4(i4 %x) {
+; CHECK-LABEL: cttz4:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    orr w8, w0, #0x10
+; CHECK-NEXT:    rbit w8, w8
+; CHECK-NEXT:    clz w0, w8
+; CHECK-NEXT:    ret
+;
+; CHECK-CSSC-LABEL: cttz4:
+; CHECK-CSSC:       // %bb.0:
+; CHECK-CSSC-NEXT:    orr w8, w0, #0x10
+; CHECK-CSSC-NEXT:    ctz w0, w8
+; CHECK-CSSC-NEXT:    ret
+  %ctz = tail call i4 @llvm.cttz.i4(i4 %x)
+  ret i4 %ctz
+}
+
+define i8 @cttz8(i8 %x) {
+; CHECK-LABEL: cttz8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    orr w8, w0, #0x100
+; CHECK-NEXT:    rbit w8, w8
+; CHECK-NEXT:    clz w0, w8
+; CHECK-NEXT:    ret
+;
+; CHECK-CSSC-LABEL: cttz8:
+; CHECK-CSSC:       // %bb.0:
+; CHECK-CSSC-NEXT:    orr w8, w0, #0x100
+; CHECK-CSSC-NEXT:    ctz w0, w8
+; CHECK-CSSC-NEXT:    ret
+  %ctz = tail call i8 @llvm.cttz.i8(i8 %x)
+  ret i8 %ctz
+}
+
+define i16 @cttz16(i16 %x) {
+; CHECK-LABEL: cttz16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    orr w8, w0, #0x10000
+; CHECK-NEXT:    rbit w8, w8
+; CHECK-NEXT:    clz w0, w8
+; CHECK-NEXT:    ret
+;
+; CHECK-CSSC-LABEL: cttz16:
+; CHECK-CSSC:       // %bb.0:
+; CHECK-CSSC-NEXT:    orr w8, w0, #0x10000
+; CHECK-CSSC-NEXT:    ctz w0, w8
+; CHECK-CSSC-NEXT:    ret
+  %ctz = tail call i16 @llvm.cttz.i16(i16 %x)
+  ret i16 %ctz
+}
+
+define i17 @cttz17(i17 %x) {
+; CHECK-LABEL: cttz17:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    orr w8, w0, #0x20000
+; CHECK-NEXT:    rbit w8, w8
+; CHECK-NEXT:    clz w0, w8
+; CHECK-NEXT:    ret
+;
+; CHECK-CSSC-LABEL: cttz17:
+; CHECK-CSSC:       // %bb.0:
+; CHECK-CSSC-NEXT:    orr w8, w0, #0x20000
+; CHECK-CSSC-NEXT:    ctz w0, w8
+; CHECK-CSSC-NEXT:    ret
+  %ctz = tail call i17 @llvm.cttz.i17(i17 %x)
+  ret i17 %ctz
+}
+
+define i32 @cttz32(i32 %x) nounwind readnone {
+; CHECK-LABEL: cttz32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    rbit w8, w0
+; CHECK-NEXT:    clz w0, w8
+; CHECK-NEXT:    ret
+;
+; CHECK-CSSC-LABEL: cttz32:
+; CHECK-CSSC:       // %bb.0:
+; CHECK-CSSC-NEXT:    ctz w0, w0
+; CHECK-CSSC-NEXT:    ret
+  %ctz = tail call i32 @llvm.cttz.i32(i32 %x)
+  ret i32 %ctz
+}
+
+define i64 @cttz64(i64 %x) nounwind readnone {
+; CHECK-LABEL: cttz64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    rbit x8, x0
+; CHECK-NEXT:    clz x0, x8
+; CHECK-NEXT:    ret
+;
+; CHECK-CSSC-LABEL: cttz64:
+; CHECK-CSSC:       // %bb.0:
+; CHECK-CSSC-NEXT:    ctz x0, x0
+; CHECK-CSSC-NEXT:    ret
+  %ctz = tail call i64 @llvm.cttz.i64(i64 %x)
+  ret i64 %ctz
+}
+
+define i128 @cttz128(i128 %x) nounwind readnone {
+; CHECK-LABEL: cttz128:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    rbit x9, x1
+; CHECK-NEXT:    rbit x8, x0
+; CHECK-NEXT:    clz x9, x9
+; CHECK-NEXT:    clz x8, x8
+; CHECK-NEXT:    add x9, x9, #64
+; CHECK-NEXT:    cmp x0, #0
+; CHECK-NEXT:    csel x0, x8, x9, ne
+; CHECK-NEXT:    mov x1, xzr
+; CHECK-NEXT:    ret
+;
+; CHECK-CSSC-LABEL: cttz128:
+; CHECK-CSSC:       // %bb.0:
+; CHECK-CSSC-NEXT:    ctz x9, x1
+; CHECK-CSSC-NEXT:    ctz x8, x0
+; CHECK-CSSC-NEXT:    add x9, x9, #64
+; CHECK-CSSC-NEXT:    cmp x0, #0
+; CHECK-CSSC-NEXT:    csel x0, x8, x9, ne
+; CHECK-CSSC-NEXT:    mov x1, xzr
+; CHECK-CSSC-NEXT:    ret
+  %ctz = tail call i128 @llvm.cttz.i128(i128 %x)
+  ret i128 %ctz
+}
+
+define i32 @cttz32combine(i32 %x) nounwind readnone {
+; CHECK-LABEL: cttz32combine:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    rbit w8, w0
+; CHECK-NEXT:    clz w0, w8
+; CHECK-NEXT:    ret
+;
+; CHECK-CSSC-LABEL: cttz32combine:
+; CHECK-CSSC:       // %bb.0:
+; CHECK-CSSC-NEXT:    ctz w0, w0
+; CHECK-CSSC-NEXT:    ret
+  %rev = tail call i32 @llvm.bitreverse.i32(i32 %x)
+  %ctz = tail call i32 @llvm.ctlz.i32(i32 %rev)
+  ret i32 %ctz
+}
+
+define i64 @cttz64combine(i64 %x) nounwind readnone {
+  %rev = tail call i64 @llvm.bitreverse.i64(i64 %x)
+  %ctz = tail call i64 @llvm.ctlz.i64(i64 %rev)
+  ret i64 %ctz
+}
+
+declare i4 @llvm.cttz.i4(i4 %x) nounwind readnone
+declare i8 @llvm.cttz.i8(i8 %x) nounwind readnone
+declare i16 @llvm.cttz.i16(i16 %x) nounwind readnone
+declare i17 @llvm.cttz.i17(i17 %x) nounwind readnone
+declare i32 @llvm.cttz.i32(i32) nounwind readnone
+declare i64 @llvm.cttz.i64(i64) nounwind readnone
+declare i128 @llvm.cttz.i128(i128) nounwind readnone
+
+declare i32 @llvm.ctlz.i32(i32 %x) nounwind readnone
+declare i32 @llvm.bitreverse.i32(i32 %x) nounwind readnone
+declare i64 @llvm.ctlz.i64(i64 %x) nounwind readnone
+declare i64 @llvm.bitreverse.i64(i64 %x) nounwind readnone