diff --git a/src/imagination/pco/pco_nir_compute.c b/src/imagination/pco/pco_nir_compute.c
index e54e4218d93..ab99d357f35 100644
--- a/src/imagination/pco/pco_nir_compute.c
+++ b/src/imagination/pco/pco_nir_compute.c
@@ -21,28 +21,16 @@
 #include <stdbool.h>
 #include <stdio.h>
 
-#define INST_CHK_FUNC "@pco_inst_chk"
-
 /**
  * \brief Inserts the instance check.
  *
- * \param[in,out] shader NIR shader.
+ * \param[in,out] entrypoint NIR shader entrypoint.
+ * \return The cursor position for where to re-insert the entrypoint.
  */
-static void insert_instance_check(nir_shader *shader)
+static nir_cursor insert_instance_check(nir_function_impl *entrypoint)
 {
-   /* Get original entrypoint. */
-   nir_function *orig_entrypoint = nir_shader_get_entrypoint(shader)->function;
-
-   /* Create a function for the instance check which will serve as the new
-    * entrypoint.
-    */
-   nir_function *inst_chk_func = nir_function_create(shader, INST_CHK_FUNC);
-
-   inst_chk_func->is_entrypoint = true;
-   orig_entrypoint->is_entrypoint = false;
-
-   nir_builder b = nir_builder_create(nir_function_impl_create(inst_chk_func));
-   b.cursor = nir_after_cf_list(&b.impl->body);
+   nir_builder b = nir_builder_at(nir_before_impl(entrypoint));
+   nir_cursor cursor;
 
    /* If the current instance index is greater than the total workgroup size,
     * we don't execute.
@@ -58,10 +46,11 @@ static void insert_instance_check(nir_shader *shader)
    nir_def *cond_inst_valid = nir_ilt(&b, flat_id, flat_size);
    nir_if *nif = nir_push_if(&b, cond_inst_valid);
    {
-      nir_call(&b, orig_entrypoint);
+      cursor = b.cursor;
    }
    nir_pop_if(&b, nif);
-   nir_jump(&b, nir_jump_return);
+
+   return cursor;
 }
 
 /**
@@ -77,22 +66,21 @@ bool pco_nir_compute_instance_check(nir_shader *shader)
    if (shader->info.internal)
       return false;
 
-   /* Check we haven't already done this. */
-   nir_foreach_function (function, shader) {
-      if (function->name && !strcmp(function->name, INST_CHK_FUNC))
-         return false;
-   }
+   nir_function_impl *entrypoint = nir_shader_get_entrypoint(shader);
 
-   insert_instance_check(shader);
+   if (nir_cf_list_is_empty_block(&entrypoint->body))
+      return false;
 
-   /* Re-inline. */
-   NIR_PASS(_, shader, nir_lower_variable_initializers, nir_var_function_temp);
-   NIR_PASS(_, shader, nir_lower_returns);
-   NIR_PASS(_, shader, nir_inline_functions);
-   NIR_PASS(_, shader, nir_copy_prop);
-   NIR_PASS(_, shader, nir_opt_deref);
-   nir_remove_non_entrypoints(shader);
-   NIR_PASS(_, shader, nir_lower_variable_initializers, ~0);
+   /* Extract the entire entrypoint. */
+   nir_cf_list cf_list;
+   nir_cf_extract(&cf_list,
+                  nir_before_impl(entrypoint),
+                  nir_after_impl(entrypoint));
+
+   nir_cursor cursor = insert_instance_check(entrypoint);
+
+   /* Re-insert the entrypoint inside the instance check. */
+   nir_cf_reinsert(&cf_list, cursor);
 
    return true;
 }