nir: handle bit shifts by constants in nir_def_bits_used

useful for open-coded bitfield extracts that are not using ubfe/ibfe

Reviewed-by: Georg Lehmann <dadschoorse@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/34489>
This commit is contained in:
Marek Olšák 2025-04-10 20:49:50 -04:00 committed by Marge Bot
parent 7d24a9b649
commit 7e7ef7b8b7
2 changed files with 52 additions and 2 deletions

View file

@ -2183,8 +2183,26 @@ ssa_def_bits_used(const nir_def *def, int recur)
case nir_op_ishl:
case nir_op_ishr:
case nir_op_ushr:
if (src_idx == 1) {
bits_used |= (nir_src_bit_size(use_alu->src[0].src) - 1);
if (src_idx == 0 && nir_src_is_const(use_alu->src[1].src)) {
unsigned bit_size = def->bit_size;
unsigned shift = nir_alu_src_as_uint(use_alu->src[1]) & (bit_size - 1);
uint64_t def_bits_used = ssa_def_bits_used(&use_alu->def, recur);
/* If one of the sign-extended bits is used, set the "last src
* bit before shifting" as used.
*/
if (use_alu->op == nir_op_ishr &&
def_bits_used & ~(all_bits >> shift))
def_bits_used |= BITFIELD64_BIT(bit_size - 1 - shift);
/* Reverse the shift to get the bits before shifting. */
if (use_alu->op == nir_op_ushr || use_alu->op == nir_op_ishr)
bits_used |= (def_bits_used << shift) & all_bits;
else
bits_used |= def_bits_used >> shift;
break;
} else if (src_idx == 1) {
bits_used |= use_alu->def.bit_size - 1;
break;
} else {
return all_bits;

View file

@ -324,3 +324,35 @@ TEST_F(ssa_def_bits_used_test, ubfe_iand)
EXPECT_EQ(nir_def_bits_used(load), BITFIELD_BIT(15));
}
TEST_F(ssa_def_bits_used_test, ishr_signed)
{
nir_def *load1 = nir_load_global(b, nir_undef(b, 1, 64), 4, 1, 32);
nir_def *load2 = nir_load_global(b, nir_undef(b, 1, 64), 4, 1, 32);
nir_def *alu1 = nir_iand_imm(b, nir_ishr_imm(b, load1, 13), 0x80000000);
nir_def *alu2 = nir_iand_imm(b, nir_ishr_imm(b, load2, 13), 0x8000);
nir_store_global(b, nir_undef(b, 1, 64), 4, alu1, 0x1);
nir_store_global(b, nir_undef(b, 1, 64), 4, alu2, 0x1);
EXPECT_EQ(nir_def_bits_used(load1), BITFIELD_BIT(31)); /* last bit */
EXPECT_EQ(nir_def_bits_used(load2), BITFIELD_BIT(15 + 13)); /* not last bit */
}
TEST_F(ssa_def_bits_used_test, ushr_ishr_ishl)
{
nir_def *load1 = nir_load_global(b, nir_undef(b, 1, 64), 4, 1, 32);
nir_def *load2 = nir_load_global(b, nir_undef(b, 1, 64), 4, 1, 32);
nir_def *load3 = nir_load_global(b, nir_undef(b, 1, 64), 4, 1, 32);
nir_def *alu1 = nir_ushr_imm(b, load1, 7);
nir_def *alu2 = nir_ishr_imm(b, load2, 11);
nir_def *alu3 = nir_ishl_imm(b, load3, 13);
nir_store_global(b, nir_undef(b, 1, 64), 4, alu1, 0x1);
nir_store_global(b, nir_undef(b, 1, 64), 4, alu2, 0x1);
nir_store_global(b, nir_undef(b, 1, 64), 4, alu3, 0x1);
EXPECT_EQ(nir_def_bits_used(load1), BITFIELD_RANGE(7, 32 - 7));
EXPECT_EQ(nir_def_bits_used(load2), BITFIELD_RANGE(11, 32 - 11));
EXPECT_EQ(nir_def_bits_used(load3), BITFIELD_RANGE(0, 32 - 13));
}