AMDGPU: mbcnt allow for non-zero src1 for known-bits

[dstuttard]

Summary

Src1 for mbcnt can be a non-zero literal or register. Take this into account
when calculating known bits.

[rampitec]

LGTM

[foad]

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

4606

Comment seems misleading since this can give useful results even when src1 is not an immediate.

Also you could make this slightly simpler and more accurate by constructing a KnownBits for 0..(wavefront size - 1) (i.e. the high 26 or 27 bits are known zero) and then calling computeForAddSub to combine it with the known bits of src1.

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

void AMDGPUTargetLowering::computeKnownBitsForTargetNode(
   }
   case ISD::INTRINSIC_WO_CHAIN: {
     unsigned IID = cast<ConstantSDNode>(Op.getOperand(0))->getZExtValue();
     switch (IID) {
     case Intrinsic::amdgcn_mbcnt_lo:
     case Intrinsic::amdgcn_mbcnt_hi: {
       const GCNSubtarget &ST =
           DAG.getMachineFunction().getSubtarget<GCNSubtarget>();
-      // These return at most the wavefront size - 1.
+      // These return at most the (wavefront size - 1) + src1
+      // As long as src1 is an immediate we can calc known bits
+      KnownBits Src1Known = DAG.computeKnownBits(Op.getOperand(2), Depth + 1);
+      unsigned Src1ValBits = Src1Known.countMaxActiveBits();
+      unsigned MaxActiveBits = std::max(Src1ValBits, ST.getWavefrontSizeLog2());
+      // Cater for potential carry
+      MaxActiveBits += Src1ValBits ? 1 : 0;
       unsigned Size = Op.getValueType().getSizeInBits();
-      Known.Zero.setHighBits(Size - ST.getWavefrontSizeLog2());
+      if (MaxActiveBits < Size)
+        Known.Zero.setHighBits(Size - MaxActiveBits);
       break;
     }
     case Intrinsic::amdgcn_workitem_id_x:
     case Intrinsic::amdgcn_workitem_id_y:
     case Intrinsic::amdgcn_workitem_id_z: {
       unsigned MaxValue = Subtarget->getMaxWorkitemID(
           DAG.getMachineFunction().getFunction(), workitemIntrinsicDim(IID));
       Known.Zero.setHighBits(countLeadingZeros(MaxValue));

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.mbcnt.ll

 ; RUN: llc -march=amdgcn -mcpu=verde -verify-machineinstrs < %s | FileCheck -check-prefix=GCN -check-prefix=SI %s
 ; RUN: llc -march=amdgcn -mcpu=tonga -verify-machineinstrs < %s | FileCheck -check-prefix=GCN -check-prefix=VI %s
 
 ; GCN-LABEL: {{^}}mbcnt_intrinsics:
 ; GCN: v_mbcnt_lo_u32_b32{{(_e64)*}} [[LO:v[0-9]+]], -1, 0
 ; SI: v_mbcnt_hi_u32_b32_e32 {{v[0-9]+}}, -1, [[LO]]
 ; VI: v_mbcnt_hi_u32_b32 {{v[0-9]+}}, -1, [[LO]]
 define amdgpu_ps void @mbcnt_intrinsics(<16 x i8> addrspace(4)* inreg %arg, <16 x i8> addrspace(4)* inreg %arg1, <32 x i8> addrspace(4)* inreg %arg2, i32 inreg %arg3) {
 main_body:
   %lo = call i32 @llvm.amdgcn.mbcnt.lo(i32 -1, i32 0) #0
   %hi = call i32 @llvm.amdgcn.mbcnt.hi(i32 -1, i32 %lo) #0
   %tmp = bitcast i32 %hi to float
   call void @llvm.amdgcn.exp.f32(i32 0, i32 15, float %tmp, float %tmp, float %tmp, float %tmp, i1 true, i1 true) #1
   ret void
 }
 
-; GCN-LABEL: {{^}}mbcnt_lo_known_bits:
+; GCN-LABEL: {{^}}mbcnt_lo_known_bits_1:
 ; GCN: v_mbcnt_lo_u32_b32
-; GCN-NOT: and
-define i32 @mbcnt_lo_known_bits(i32 %x, i32 %y) #0 {
+; GCN: v_and_b32_e32
+define i32 @mbcnt_lo_known_bits_1(i32 %x, i32 %y) #0 {
   %lo = call i32 @llvm.amdgcn.mbcnt.lo(i32 %x, i32 %y)
   %mask = and i32 %lo, 63
   ret i32 %mask
 }
 
-; GCN-LABEL: {{^}}mbcnt_hi_known_bits:
-; GCN: v_mbcnt_hi_u32_b32
+; GCN-LABEL: {{^}}mbcnt_lo_known_bits_2:
+; GCN: v_mbcnt_lo_u32_b32
+; GCN-NOT: and
+define i32 @mbcnt_lo_known_bits_2(i32 %x) #0 {
+  %lo = call i32 @llvm.amdgcn.mbcnt.lo(i32 %x, i32 0)
+  %mask = and i32 %lo, 63
+  ret i32 %mask
+}
+
+; GCN-LABEL: {{^}}mbcnt_lo_known_bits_3:
+; GCN: v_mbcnt_lo_u32_b32
 ; GCN-NOT: and
-define i32 @mbcnt_hi_known_bits(i32 %x, i32 %y) #0 {
+define i32 @mbcnt_lo_known_bits_3(i32 %x) #0 {
+  %lo = call i32 @llvm.amdgcn.mbcnt.lo(i32 %x, i32 15)
+  %mask = and i32 %lo, 127
+  ret i32 %mask
+}
+
+; GCN-LABEL: {{^}}mbcnt_lo_known_bits_4:
+; GCN: v_mbcnt_lo_u32_b32
+; GCN: v_and_b32_e32
+define i32 @mbcnt_lo_known_bits_4(i32 %x) #0 {
+  %lo = call i32 @llvm.amdgcn.mbcnt.lo(i32 %x, i32 15)
+  %mask = and i32 %lo, 63
+  ret i32 %mask
+}
+
+
+; GCN-LABEL: {{^}}mbcnt_hi_known_bits_1:
+; GCN: v_mbcnt_hi_u32_b32
+; GCN: v_and_b32_e32
+define i32 @mbcnt_hi_known_bits_1(i32 %x, i32 %y) #0 {
   %hi = call i32 @llvm.amdgcn.mbcnt.hi(i32 %x, i32 %y)
   %mask = and i32 %hi, 63
   ret i32 %mask
 }
 
+; GCN-LABEL: {{^}}mbcnt_hi_known_bits_2:
+; GCN: v_mbcnt_hi_u32_b32
+; GCN-NOT: and
+define i32 @mbcnt_hi_known_bits_2(i32 %x) #0 {
+  %hi = call i32 @llvm.amdgcn.mbcnt.hi(i32 %x, i32 0)
+  %mask = and i32 %hi, 63
+  ret i32 %mask
+}
+
+; GCN-LABEL: {{^}}mbcnt_hi_known_bits_3:
+; GCN: v_mbcnt_hi_u32_b32
+; GCN-NOT: and
+define i32 @mbcnt_hi_known_bits_3(i32 %x) #0 {
+  %hi = call i32 @llvm.amdgcn.mbcnt.hi(i32 %x, i32 15)
+  %mask = and i32 %hi, 127
+  ret i32 %mask
+}
+
+; GCN-LABEL: {{^}}mbcnt_hi_known_bits_4:
+; GCN: v_mbcnt_hi_u32_b32
+; GCN: v_and_b32_e32
+define i32 @mbcnt_hi_known_bits_4(i32 %x) #0 {
+  %hi = call i32 @llvm.amdgcn.mbcnt.hi(i32 %x, i32 15)
+  %mask = and i32 %hi, 63
+  ret i32 %mask
+}
+
 declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32) #0
 declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32) #0
 declare void @llvm.amdgcn.exp.f32(i32, i32, float, float, float, float, i1, i1) #1
 
 attributes #0 = { nounwind readnone }
 attributes #1 = { nounwind }