diff --git a/src/amd/compiler/aco_instruction_selection.cpp b/src/amd/compiler/aco_instruction_selection.cpp
index 2e4f7af7dee..eef9f28df67 100644
--- a/src/amd/compiler/aco_instruction_selection.cpp
+++ b/src/amd/compiler/aco_instruction_selection.cpp
@@ -11680,7 +11680,7 @@ select_program(Program* program, unsigned shader_count, struct nir_shader* const
                const struct aco_shader_info* info, const struct ac_shader_args* args)
 {
    isel_context ctx =
-      setup_isel_context(program, shader_count, shaders, config, options, info, args, false, false);
+      setup_isel_context(program, shader_count, shaders, config, options, info, args);
 
    if (ctx.stage == raytracing_cs)
       return select_program_rt(ctx, shader_count, shaders, args);
@@ -12285,7 +12285,7 @@ select_ps_epilog(Program* program, void* pinfo, ac_shader_config* config,
 {
    const struct aco_ps_epilog_info* einfo = (const struct aco_ps_epilog_info*)pinfo;
    isel_context ctx =
-      setup_isel_context(program, 0, NULL, config, options, info, args, true, false);
+      setup_isel_context(program, 0, NULL, config, options, info, args, SWStage::FS);
 
    ctx.block->fp_mode = program->next_fp_mode;
 
@@ -12355,7 +12355,7 @@ select_tcs_epilog(Program* program, void* pinfo, ac_shader_config* config,
 {
    const struct aco_tcs_epilog_info* einfo = (const struct aco_tcs_epilog_info*)pinfo;
    isel_context ctx =
-      setup_isel_context(program, 0, NULL, config, options, info, args, false, true);
+      setup_isel_context(program, 0, NULL, config, options, info, args, SWStage::TCS);
 
    ctx.block->fp_mode = program->next_fp_mode;
 
diff --git a/src/amd/compiler/aco_instruction_selection.h b/src/amd/compiler/aco_instruction_selection.h
index 43bb317257d..9ce9de95224 100644
--- a/src/amd/compiler/aco_instruction_selection.h
+++ b/src/amd/compiler/aco_instruction_selection.h
@@ -118,8 +118,8 @@ isel_context setup_isel_context(Program* program, unsigned shader_count,
                                 struct nir_shader* const* shaders, ac_shader_config* config,
                                 const struct aco_compiler_options* options,
                                 const struct aco_shader_info* info,
-                                const struct ac_shader_args* args, bool is_ps_epilog,
-                                bool is_tcs_epilog);
+                                const struct ac_shader_args* args,
+                                SWStage sw_stage = SWStage::None);
 
 } // namespace aco
 
diff --git a/src/amd/compiler/aco_instruction_selection_setup.cpp b/src/amd/compiler/aco_instruction_selection_setup.cpp
index a5959323a47..27b318451d7 100644
--- a/src/amd/compiler/aco_instruction_selection_setup.cpp
+++ b/src/amd/compiler/aco_instruction_selection_setup.cpp
@@ -649,9 +649,8 @@ isel_context
 setup_isel_context(Program* program, unsigned shader_count, struct nir_shader* const* shaders,
                    ac_shader_config* config, const struct aco_compiler_options* options,
                    const struct aco_shader_info* info, const struct ac_shader_args* args,
-                   bool is_ps_epilog, bool is_tcs_epilog)
+                   SWStage sw_stage)
 {
-   SWStage sw_stage = SWStage::None;
    for (unsigned i = 0; i < shader_count; i++) {
       switch (shaders[i]->info.stage) {
       case MESA_SHADER_VERTEX: sw_stage = sw_stage | SWStage::VS; break;
@@ -673,16 +672,6 @@ setup_isel_context(Program* program, unsigned shader_count, struct nir_shader* c
       }
    }
 
-   if (is_ps_epilog) {
-      assert(shader_count == 0 && !shaders);
-      sw_stage = SWStage::FS;
-   }
-
-   if (is_tcs_epilog) {
-      assert(shader_count == 0 && !shaders);
-      sw_stage = SWStage::TCS;
-   }
-
    init_program(program, Stage{info->hw_stage, sw_stage}, info, options->gfx_level, options->family,
                 options->wgp_mode, config);