diff --git a/src/panfrost/bifrost/bir.c b/src/panfrost/bifrost/bir.c
index 397710c44c3..602ec968e94 100644
--- a/src/panfrost/bifrost/bir.c
+++ b/src/panfrost/bifrost/bir.c
@@ -93,22 +93,24 @@ bi_count_read_registers(bi_instr *ins, unsigned s)
 }
 
 unsigned
-bi_writemask(bi_instr *ins, unsigned d)
+bi_count_write_registers(bi_instr *ins, unsigned d)
 {
-        /* Assume we write a scalar */
-        unsigned mask = 0xF;
-
         if (d == 0 && bi_opcode_props[ins->op].sr_write) {
-                unsigned count = bi_count_staging_registers(ins);
-
                 /* TODO: this special case is even more special, TEXC has a
                  * generic write mask stuffed in the desc... */
                 if (ins->op == BI_OPCODE_TEXC)
-                        count = 4;
-
-                mask = (1 << (count * 4)) - 1;
+                        return 4;
+                else
+                        return bi_count_staging_registers(ins);
         }
 
+        return 1;
+}
+
+unsigned
+bi_writemask(bi_instr *ins, unsigned d)
+{
+        unsigned mask = BITFIELD_MASK(bi_count_write_registers(ins, d) * 4);
         unsigned shift = ins->dest[d].offset * 4; /* 32-bit words */
         return (mask << shift);
 }
diff --git a/src/panfrost/bifrost/compiler.h b/src/panfrost/bifrost/compiler.h
index e8bbcfc76ea..3021eea17ee 100644
--- a/src/panfrost/bifrost/compiler.h
+++ b/src/panfrost/bifrost/compiler.h
@@ -756,6 +756,7 @@ pan_next_block(pan_block *block)
 
 bool bi_has_arg(bi_instr *ins, bi_index arg);
 unsigned bi_count_read_registers(bi_instr *ins, unsigned src);
+unsigned bi_count_write_registers(bi_instr *ins, unsigned dest);
 unsigned bi_writemask(bi_instr *ins, unsigned dest);
 bi_clause * bi_next_clause(bi_context *ctx, pan_block *block, bi_clause *clause);
 bool bi_side_effects(enum bi_opcode op);