diff --git a/src/broadcom/compiler/vir_opt_dead_code.c b/src/broadcom/compiler/vir_opt_dead_code.c
index 55469402e7a..64c762c88db 100644
--- a/src/broadcom/compiler/vir_opt_dead_code.c
+++ b/src/broadcom/compiler/vir_opt_dead_code.c
@@ -149,25 +149,30 @@ check_first_ldunifa(struct v3d_compile *c,
 }
 
 static bool
-increment_unifa_address(struct v3d_compile *c, struct qinst *unifa)
+increment_unifa_address(struct v3d_compile *c, struct qblock *block, struct qinst *unifa)
 {
+        struct qblock *current_block = c->cur_block;
         if (unifa->qpu.type == V3D_QPU_INSTR_TYPE_ALU &&
             unifa->qpu.alu.mul.op == V3D_QPU_M_MOV) {
                 c->cursor = vir_after_inst(unifa);
+                c->cur_block = block;
                 struct qreg unifa_reg = vir_reg(QFILE_MAGIC, V3D_QPU_WADDR_UNIFA);
                 vir_ADD_dest(c, unifa_reg, unifa->src[0], vir_uniform_ui(c, 4u));
                 vir_remove_instruction(c, unifa);
+                c->cur_block = current_block;
                 return true;
         }
 
         if (unifa->qpu.type == V3D_QPU_INSTR_TYPE_ALU &&
             unifa->qpu.alu.add.op == V3D_QPU_A_ADD) {
                 c->cursor = vir_after_inst(unifa);
+                c->cur_block = block;
                 struct qreg unifa_reg = vir_reg(QFILE_MAGIC, V3D_QPU_WADDR_UNIFA);
                 struct qreg tmp =
                         vir_ADD(c, unifa->src[1], vir_uniform_ui(c, 4u));
                 vir_ADD_dest(c, unifa_reg, unifa->src[0], tmp);
                 vir_remove_instruction(c, unifa);
+                c->cur_block = current_block;
                 return true;
         }
 
@@ -271,7 +276,7 @@ vir_opt_dead_code(struct v3d_compile *c)
                          */
                         if (is_first_ldunifa) {
                                 assert(unifa);
-                                if (!increment_unifa_address(c, unifa))
+                                if (!increment_unifa_address(c, block, unifa))
                                         continue;
                         }