[AMDGPU] Enable divergence-driven 'ctpop' selection

This change adds the patterns and divergence predicates for the ctpop (bitcount) nodes
to make them selected according to the divergence.

Reviewed By: foad

Differential Revision: https://reviews.llvm.org/D116284

Author: alex-t

llvm/lib/Target/AMDGPU/SIInstructions.td

@@ -1011,7 +1011,7 @@ def : GCNPat <
 }
 
 def : GCNPat <
-  (i32 (ctpop i32:$popcnt)),
+  (i32 (DivergentUnaryFrag<ctpop> i32:$popcnt)),
   (V_BCNT_U32_B32_e64 VSrc_b32:$popcnt, (i32 0))
 >;
 
@@ -1020,6 +1020,14 @@ def : GCNPat <
   (V_BCNT_U32_B32_e64 $popcnt, $val)
 >;
 
+def : GCNPat <
+  (i64 (DivergentUnaryFrag<ctpop> i64:$src)),
+  (REG_SEQUENCE VReg_64,
+    (V_BCNT_U32_B32_e64 (i32 (EXTRACT_SUBREG i64:$src, sub1)),
+      (i32 (V_BCNT_U32_B32_e64 (i32 (EXTRACT_SUBREG i64:$src, sub0)), (i32 0)))), sub0,
+      (i32 (V_MOV_B32_e32 (i32 0))), sub1)
+>;
+
 /********** ============================================ **********/
 /********** Extraction, Insertion, Building and Casting  **********/
 /********** ============================================ **********/

llvm/lib/Target/AMDGPU/SOPInstructions.td

@@ -246,10 +246,10 @@ let Defs = [SCC] in {
 def S_BCNT0_I32_B32 : SOP1_32 <"s_bcnt0_i32_b32">;
 def S_BCNT0_I32_B64 : SOP1_32_64 <"s_bcnt0_i32_b64">;
 def S_BCNT1_I32_B32 : SOP1_32 <"s_bcnt1_i32_b32",
-  [(set i32:$sdst, (ctpop i32:$src0))]
+  [(set i32:$sdst, (UniformUnaryFrag<ctpop> i32:$src0))]
 >;
 def S_BCNT1_I32_B64 : SOP1_32_64 <"s_bcnt1_i32_b64",
-  [(set i32:$sdst, (ctpop i64:$src0))]
+  [(set i32:$sdst, (UniformUnaryFrag<ctpop> i64:$src0))]
 >;
 } // End Defs = [SCC]
 
@@ -1371,7 +1371,7 @@ def : GCNPat <
 >;
 
 def : GCNPat <
-  (i64 (ctpop i64:$src)),
+  (i64 (UniformUnaryFrag<ctpop> i64:$src)),
     (i64 (REG_SEQUENCE SReg_64,
      (i32 (COPY_TO_REGCLASS (S_BCNT1_I32_B64 $src), SReg_32)), sub0,
      (S_MOV_B32 (i32 0)), sub1))

llvm/test/CodeGen/AMDGPU/divergence-driven-ctpop.ll

@@ -0,0 +1,53 @@
+; RUN: llc -march=amdgcn -stop-after=amdgpu-isel < %s | FileCheck -check-prefix=GCN %s
+
+; GCN-LABEL: name:            s_ctpop_i32
+; GCN: S_BCNT1_I32_B32
+define amdgpu_kernel void @s_ctpop_i32(i32 addrspace(1)* noalias %out, i32 %val) nounwind {
+  %ctpop = call i32 @llvm.ctpop.i32(i32 %val) nounwind readnone
+  store i32 %ctpop, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
+; GCN-LABEL: name:            s_ctpop_i64
+; GCN: %[[BCNT:[0-9]+]]:sreg_32 = S_BCNT1_I32_B64
+; GCN: %[[SREG1:[0-9]+]]:sreg_32 = COPY %[[BCNT]]
+; GCN: %[[SREG2:[0-9]+]]:sreg_32 = S_MOV_B32 0
+; GCN: REG_SEQUENCE killed %[[SREG1]], %subreg.sub0, killed %[[SREG2]], %subreg.sub1
+define amdgpu_kernel void @s_ctpop_i64(i32 addrspace(1)* noalias %out, i64 %val) nounwind {
+  %ctpop = call i64 @llvm.ctpop.i64(i64 %val) nounwind readnone
+  %truncctpop = trunc i64 %ctpop to i32
+  store i32 %truncctpop, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
+; GCN-LABEL: name:            v_ctpop_i32
+; GCN: V_BCNT_U32_B32_e64
+define amdgpu_kernel void @v_ctpop_i32(i32 addrspace(1)* noalias %out, i32 addrspace(1)* noalias %in) nounwind {
+  %tid = call i32 @llvm.amdgcn.workitem.id.x()
+  %in.gep = getelementptr i32, i32 addrspace(1)* %in, i32 %tid
+  %val = load i32, i32 addrspace(1)* %in.gep, align 4
+  %ctpop = call i32 @llvm.ctpop.i32(i32 %val) nounwind readnone
+  store i32 %ctpop, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
+; GCN-LABEL: name:            v_ctpop_i64
+; GCN: %[[BCNT1:[0-9]+]]:vgpr_32 = V_BCNT_U32_B32_e64 killed %{{[0-9]+}}, 0, implicit $exec
+; GCN: %[[BCNT2:[0-9]+]]:vgpr_32 = V_BCNT_U32_B32_e64 killed %{{[0-9]+}}, killed %[[BCNT1]], implicit $exec
+; GCN: %[[VGPR1:[0-9]+]]:vgpr_32 = V_MOV_B32_e32 0, implicit $exec
+; GCN: REG_SEQUENCE killed %[[BCNT2]], %subreg.sub0, killed %[[VGPR1]], %subreg.sub1
+define amdgpu_kernel void @v_ctpop_i64(i32 addrspace(1)* noalias %out, i64 addrspace(1)* noalias %in) nounwind {
+  %tid = call i32 @llvm.amdgcn.workitem.id.x()
+  %in.gep = getelementptr i64, i64 addrspace(1)* %in, i32 %tid
+  %val = load i64, i64 addrspace(1)* %in.gep, align 8
+  %ctpop = call i64 @llvm.ctpop.i64(i64 %val) nounwind readnone
+  %truncctpop = trunc i64 %ctpop to i32
+  store i32 %truncctpop, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
+declare i64 @llvm.ctpop.i64(i64) nounwind readnone
+
+declare i32 @llvm.ctpop.i32(i32) nounwind readnone
+
+declare i32 @llvm.amdgcn.workitem.id.x() nounwind readnone