[AMDGPU] Make bfe patterns divergence-aware

[foad]

Summary

This tends to increase code size but more importantly it reduces vgpr
usage, and could avoid costly readfirstlanes if the result needs to be
in an sgpr.

[foad]

Rebase.

llvm/lib/Target/AMDGPU/SIInstructions.td

 }]>;
 
 def IMMPopCount : SDNodeXForm<imm, [{
   return CurDAG->getTargetConstant(countPopulation(N->getZExtValue()), SDLoc(N),
                                    MVT::i32);
 }]>;
 
 def : AMDGPUPat <
-  (and (i32 (srl i32:$src, i32:$rshift)), IMMZeroBasedBitfieldMask:$mask),
+  (DivergentBinFrag<and> (i32 (srl i32:$src, i32:$rshift)),
+                         IMMZeroBasedBitfieldMask:$mask),
   (V_BFE_U32 $src, $rshift, (S_MOV_B32 (i32 (IMMPopCount $mask))))
 >;
 
 // x & ((1 << y) - 1)
 def : AMDGPUPat <
-  (and i32:$src, (add_oneuse (shl_oneuse 1, i32:$width), -1)),
+  (DivergentBinFrag<and> i32:$src, (add_oneuse (shl_oneuse 1, i32:$width), -1)),
   (V_BFE_U32 $src, (S_MOV_B32 (i32 0)), $width)
 >;
 
 // x & ~(-1 << y)
 def : AMDGPUPat <
-  (and i32:$src, (xor_oneuse (shl_oneuse -1, i32:$width), -1)),
+  (DivergentBinFrag<and> i32:$src,
+                         (xor_oneuse (shl_oneuse -1, i32:$width), -1)),
   (V_BFE_U32 $src, (S_MOV_B32 (i32 0)), $width)
 >;
 
 // x & (-1 >> (bitwidth - y))
 def : AMDGPUPat <
-  (and i32:$src, (srl_oneuse -1, (sub 32, i32:$width))),
+  (DivergentBinFrag<and> i32:$src, (srl_oneuse -1, (sub 32, i32:$width))),
   (V_BFE_U32 $src, (S_MOV_B32 (i32 0)), $width)
 >;
 
 // x << (bitwidth - y) >> (bitwidth - y)
 def : AMDGPUPat <
-  (srl (shl_oneuse i32:$src, (sub 32, i32:$width)), (sub 32, i32:$width)),
+  (DivergentBinFrag<srl> (shl_oneuse i32:$src, (sub 32, i32:$width)),
+                         (sub 32, i32:$width)),
   (V_BFE_U32 $src, (S_MOV_B32 (i32 0)), $width)
 >;
 
 def : AMDGPUPat <
-  (sra (shl_oneuse i32:$src, (sub 32, i32:$width)), (sub 32, i32:$width)),
+  (DivergentBinFrag<sra> (shl_oneuse i32:$src, (sub 32, i32:$width)),
+                         (sub 32, i32:$width)),
   (V_BFE_I32 $src, (S_MOV_B32 (i32 0)), $width)
 >;
 
 // SHA-256 Ma patterns
 
 // ((x & z) | (y & (x | z))) -> BFI (XOR x, y), z, y
 def : AMDGPUPat <
   (DivergentBinFrag<or> (and i32:$x, i32:$z),

llvm/test/CodeGen/AMDGPU/amdgpu.private-memory.ll

 ; R600-VECT: MOVA_INT
 
 ; SI-ALLOCA-DAG: buffer_store_short v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], 0 offset:6 ; encoding: [0x06,0x00,0x68,0xe0
 ; SI-ALLOCA-DAG: buffer_store_short v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], 0 offset:4 ; encoding: [0x04,0x00,0x68,0xe0
 ; Loaded value is 0 or 1, so sext will become zext, so we get buffer_load_ushort instead of buffer_load_sshort.
 ; SI-ALLOCA: buffer_load_sshort v{{[0-9]+}}, v{{[0-9]+}}, s[{{[0-9]+:[0-9]+}}], 0
 
 ; SI-PROMOTE-VECT: s_load_dword [[IDX:s[0-9]+]]
-; SI-PROMOTE-VECT: s_mov_b32 [[SREG:s[0-9]+]], 0x10000
 ; SI-PROMOTE-VECT: s_lshl_b32 [[SCALED_IDX:s[0-9]+]], [[IDX]], 4
-; SI-PROMOTE-VECT: v_mov_b32_e32 [[VREG:v[0-9]+]], [[SCALED_IDX]]
-; SI-PROMOTE-VECT: v_bfe_u32 v{{[0-9]+}}, [[SREG]], [[VREG]], 16
+; SI-PROMOTE-VECT: s_lshr_b32 [[SREG:s[0-9]+]], 0x10000, [[SCALED_IDX]]
+; SI-PROMOTE-VECT: s_and_b32 s{{[0-9]+}}, [[SREG]], 0xffff
 define amdgpu_kernel void @short_array(i32 addrspace(1)* %out, i32 %index) #0 {
 entry:
   %0 = alloca [2 x i16], addrspace(5)
   %1 = getelementptr inbounds [2 x i16], [2 x i16] addrspace(5)* %0, i32 0, i32 0
   %2 = getelementptr inbounds [2 x i16], [2 x i16] addrspace(5)* %0, i32 0, i32 1
   store i16 0, i16 addrspace(5)* %1
   store i16 1, i16 addrspace(5)* %2
   %3 = getelementptr inbounds [2 x i16], [2 x i16] addrspace(5)* %0, i32 0, i32 %index

llvm/test/CodeGen/AMDGPU/bfe-patterns.ll

define amdgpu_kernel void @v_ubfe_sub_multi_use_shl_i32(i32 addrspace(1)* %out, i32 addrspace(1)* %in0, i32 addrspace(1)* %in1) #1 {
   %bfe = lshr i32 %shl, %sub
   store i32 %bfe, i32 addrspace(1)* %out.gep
   store volatile i32 %shl, i32 addrspace(1)* undef
   ret void
 }
 
 ; GCN-LABEL: {{^}}s_ubfe_sub_i32:
 ; GCN: s_load_dwordx2 s{{\[}}[[SRC:[0-9]+]]:[[WIDTH:[0-9]+]]{{\]}}, s[0:1], {{0xb|0x2c}}
-; GCN: v_mov_b32_e32 [[VWIDTH:v[0-9]+]], s[[WIDTH]]
-; GCN: v_bfe_u32 v{{[0-9]+}}, s[[SRC]], 0, [[VWIDTH]]
+; GCN: s_sub_i32 [[SUB:s[0-9]+]], 32, s[[WIDTH]]
+; GCN: s_lshl_b32 [[TMP:s[0-9]+]], s[[SRC]], [[SUB]]
+; GCN: s_lshr_b32 s{{[0-9]+}}, [[TMP]], [[SUB]]
 define amdgpu_kernel void @s_ubfe_sub_i32(i32 addrspace(1)* %out, i32 %src, i32 %width) #1 {
   %id.x = tail call i32 @llvm.amdgcn.workitem.id.x()
   %out.gep = getelementptr i32, i32 addrspace(1)* %out, i32 %id.x
   %sub = sub i32 32, %width
   %shl = shl i32 %src, %sub
   %bfe = lshr i32 %shl, %sub
   store i32 %bfe, i32 addrspace(1)* %out.gep
   ret void
define amdgpu_kernel void @v_sbfe_sub_multi_use_shl_i32(i32 addrspace(1)* %out, i32 addrspace(1)* %in0, i32 addrspace(1)* %in1) #1 {
   %bfe = ashr i32 %shl, %sub
   store i32 %bfe, i32 addrspace(1)* %out.gep
   store volatile i32 %shl, i32 addrspace(1)* undef
   ret void
 }
 
 ; GCN-LABEL: {{^}}s_sbfe_sub_i32:
 ; GCN: s_load_dwordx2 s{{\[}}[[SRC:[0-9]+]]:[[WIDTH:[0-9]+]]{{\]}}, s[0:1], {{0xb|0x2c}}
-; GCN: v_mov_b32_e32 [[VWIDTH:v[0-9]+]], s[[WIDTH]]
-; GCN: v_bfe_i32 v{{[0-9]+}}, s[[SRC]], 0, [[VWIDTH]]
+; GCN: s_sub_i32 [[SUB:s[0-9]+]], 32, s[[WIDTH]]
+; GCN: s_lshl_b32 [[TMP:s[0-9]+]], s[[SRC]], [[SUB]]
+; GCN: s_ashr_i32 s{{[0-9]+}}, [[TMP]], [[SUB]]
 define amdgpu_kernel void @s_sbfe_sub_i32(i32 addrspace(1)* %out, i32 %src, i32 %width) #1 {
   %id.x = tail call i32 @llvm.amdgcn.workitem.id.x()
   %out.gep = getelementptr i32, i32 addrspace(1)* %out, i32 %id.x
   %sub = sub i32 32, %width
   %shl = shl i32 %src, %sub
   %bfe = ashr i32 %shl, %sub
   store i32 %bfe, i32 addrspace(1)* %out.gep
   ret void

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

 }
 
 define amdgpu_kernel void @v_lshr_and(i32 addrspace(1)* %out, i32 %a, i32 %b) #0 {
 ; SI-LABEL: v_lshr_and:
 ; SI:       ; %bb.0:
 ; SI-NEXT:    s_load_dwordx2 s[2:3], s[0:1], 0xb
 ; SI-NEXT:    s_load_dwordx2 s[4:5], s[0:1], 0x9
 ; SI-NEXT:    s_mov_b32 s7, 0xf000
-; SI-NEXT:    s_mov_b32 s6, -1
 ; SI-NEXT:    s_waitcnt lgkmcnt(0)
-; SI-NEXT:    v_mov_b32_e32 v0, s3
-; SI-NEXT:    v_bfe_u32 v0, s2, v0, 3
+; SI-NEXT:    s_lshr_b32 s0, s2, s3
+; SI-NEXT:    s_and_b32 s0, s0, 7
+; SI-NEXT:    s_mov_b32 s6, -1
+; SI-NEXT:    v_mov_b32_e32 v0, s0
 ; SI-NEXT:    buffer_store_dword v0, off, s[4:7], 0
 ; SI-NEXT:    s_endpgm
 ;
 ; VI-LABEL: v_lshr_and:
 ; VI:       ; %bb.0:
 ; VI-NEXT:    s_load_dwordx2 s[4:5], s[0:1], 0x24
 ; VI-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x2c
 ; VI-NEXT:    s_mov_b32 s7, 0xf000
 ; VI-NEXT:    s_mov_b32 s6, -1
 ; VI-NEXT:    s_waitcnt lgkmcnt(0)
-; VI-NEXT:    v_mov_b32_e32 v0, s1
-; VI-NEXT:    v_bfe_u32 v0, s0, v0, 3
+; VI-NEXT:    s_lshr_b32 s0, s0, s1
+; VI-NEXT:    s_and_b32 s0, s0, 7
+; VI-NEXT:    v_mov_b32_e32 v0, s0
 ; VI-NEXT:    buffer_store_dword v0, off, s[4:7], 0
 ; VI-NEXT:    s_endpgm
   %c = lshr i32 %a, %b
   %d = and i32 %c, 7
   store i32 %d, i32 addrspace(1)* %out, align 8
   ret void
 }