aco: add infra for compiling TCS epilogs

Signed-off-by: Samuel Pitoiset <samuel.pitoiset@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/24417>

7 files changed, 75 insertions, 5 deletions

diff --git a/src/amd/compiler/aco_instruction_selection.cpp b/src/amd/compiler/aco_instruction_selection.cpp
index fef2746d363..95bf2c3dfeb 100644
--- a/src/amd/compiler/aco_instruction_selection.cpp
+++ b/src/amd/compiler/aco_instruction_selection.cpp
@@ -200,6 +200,7 @@ emit_bpermute(isel_context* ctx, Builder& bld, Temp index, Temp data)
       ctx->options->gfx_level >= GFX10 && ctx->options->gfx_level < GFX11 &&
       ctx->program->wave_size == 64 &&
       ((ctx->stage == fragment_fs && ctx->program->info.ps.has_epilog) ||
+       (ctx->stage == tess_control_hs && ctx->program->info.tcs.has_epilog) ||
        ctx->stage == raytracing_cs);
 
    if (ctx->options->gfx_level <= GFX7 || avoid_shared_vgprs) {
@@ -10855,6 +10856,12 @@ create_fs_jump_to_epilog(isel_context* ctx)
    ctx->block->instructions.emplace_back(std::move(jump));
 }
 
+static void
+create_tcs_jump_to_epilog(isel_context* ctx)
+{
+   /* TODO */
+}
+
 Pseudo_instruction*
 add_startpgm(struct isel_context* ctx)
 {
@@ -11295,6 +11302,8 @@ select_shader(isel_context& ctx, nir_shader* nir, const bool need_startpgm, cons
       /* FS epilogs always have at least one color/null export. */
       ctx.program->has_color_exports = true;
       ctx.block->kind |= block_kind_export_end;
+   } else if (ctx.stage == tess_control_hs && ctx.program->info.tcs.has_epilog) {
+      create_tcs_jump_to_epilog(&ctx);
    }
 
    if (endif_merged_wave_info) {
@@ -11359,7 +11368,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);
+      setup_isel_context(program, shader_count, shaders, config, options, info, args, false, false);
 
    if (ctx.stage == raytracing_cs)
       return select_program_rt(ctx, shader_count, shaders, args);
@@ -11949,7 +11958,8 @@ select_ps_epilog(Program* program, void* pinfo, ac_shader_config* config,
                  const struct ac_shader_args* args)
 {
    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);
+   isel_context ctx =
+      setup_isel_context(program, 0, NULL, config, options, info, args, true, false);
 
    ctx.block->fp_mode = program->next_fp_mode;
 
@@ -12011,4 +12021,33 @@ select_ps_epilog(Program* program, void* pinfo, ac_shader_config* config,
 
    cleanup_cfg(program);
 }
+
+void
+select_tcs_epilog(Program* program, void* pinfo, ac_shader_config* config,
+                  const struct aco_compiler_options* options, const struct aco_shader_info* info,
+                  const struct ac_shader_args* args)
+{
+   UNUSED 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);
+
+   ctx.block->fp_mode = program->next_fp_mode;
+
+   add_startpgm(&ctx);
+   append_logical_start(ctx.block);
+
+   Builder bld(ctx.program, ctx.block);
+
+   /* TODO */
+
+   program->config->float_mode = program->blocks[0].fp_mode.val;
+
+   append_logical_end(ctx.block);
+   ctx.block->kind |= block_kind_export_end;
+   bld.reset(ctx.block);
+   bld.sopp(aco_opcode::s_endpgm);
+
+   cleanup_cfg(program);
+}
+
 } // namespace aco
diff --git a/src/amd/compiler/aco_instruction_selection.h b/src/amd/compiler/aco_instruction_selection.h
index f1b1b04ce8c..fd886229435 100644
--- a/src/amd/compiler/aco_instruction_selection.h
+++ b/src/amd/compiler/aco_instruction_selection.h
@@ -123,7 +123,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);
+                                const struct ac_shader_args* args, bool is_ps_epilog,
+                                bool is_tcs_epilog);
 
 } // namespace aco
 
diff --git a/src/amd/compiler/aco_instruction_selection_setup.cpp b/src/amd/compiler/aco_instruction_selection_setup.cpp
index ea4810112ab..cc2089c9787 100644
--- a/src/amd/compiler/aco_instruction_selection_setup.cpp
+++ b/src/amd/compiler/aco_instruction_selection_setup.cpp
@@ -651,7 +651,7 @@ 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_ps_epilog, bool is_tcs_epilog)
 {
    SWStage sw_stage = SWStage::None;
    for (unsigned i = 0; i < shader_count; i++) {
@@ -680,6 +680,11 @@ setup_isel_context(Program* program, unsigned shader_count, struct nir_shader* c
       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);
 
@@ -696,7 +701,7 @@ setup_isel_context(Program* program, unsigned shader_count, struct nir_shader* c
    ASSERTED bool mesh_shading = ctx.stage.has(SWStage::TS) || ctx.stage.has(SWStage::MS);
    assert(!mesh_shading || ctx.program->gfx_level >= GFX10_3);
 
-   if (ctx.stage == tess_control_hs)
+   if (ctx.stage == tess_control_hs && !is_tcs_epilog)
       setup_tcs_info(&ctx, shaders[0], NULL);
    else if (ctx.stage == vertex_tess_control_hs)
       setup_tcs_info(&ctx, shaders[1], shaders[0]);
diff --git a/src/amd/compiler/aco_interface.cpp b/src/amd/compiler/aco_interface.cpp
index 40995cc9fe4..fccad69a5c0 100644
--- a/src/amd/compiler/aco_interface.cpp
+++ b/src/amd/compiler/aco_interface.cpp
@@ -386,3 +386,13 @@ aco_compile_ps_epilog(const struct aco_compiler_options* options,
    aco_compile_shader_part(options, info, args, aco::select_ps_epilog, (void*)pinfo, build_epilog,
                            binary);
 }
+
+void
+aco_compile_tcs_epilog(const struct aco_compiler_options* options,
+                       const struct aco_shader_info* info, const struct aco_tcs_epilog_info* pinfo,
+                       const struct ac_shader_args* args, aco_shader_part_callback* build_epilog,
+                       void** binary)
+{
+   aco_compile_shader_part(options, info, args, aco::select_tcs_epilog, (void*)pinfo, build_epilog,
+                           binary);
+}
diff --git a/src/amd/compiler/aco_interface.h b/src/amd/compiler/aco_interface.h
index e28ff73191b..44c0fad3ef8 100644
--- a/src/amd/compiler/aco_interface.h
+++ b/src/amd/compiler/aco_interface.h
@@ -77,6 +77,12 @@ void aco_compile_ps_epilog(const struct aco_compiler_options* options,
                            const struct ac_shader_args* args,
                            aco_shader_part_callback* build_epilog, void** binary);
 
+void aco_compile_tcs_epilog(const struct aco_compiler_options* options,
+                            const struct aco_shader_info* info,
+                            const struct aco_tcs_epilog_info* epilog_info,
+                            const struct ac_shader_args* args,
+                            aco_shader_part_callback* build_epilog, void** binary);
+
 uint64_t aco_get_codegen_flags();
 
 #ifdef __cplusplus
diff --git a/src/amd/compiler/aco_ir.h b/src/amd/compiler/aco_ir.h
index 775e82d70c9..b1f0280437e 100644
--- a/src/amd/compiler/aco_ir.h
+++ b/src/amd/compiler/aco_ir.h
@@ -2232,6 +2232,10 @@ void select_ps_epilog(Program* program, void* pinfo, ac_shader_config* config,
                       const struct aco_compiler_options* options,
                       const struct aco_shader_info* info, const struct ac_shader_args* args);
 
+void select_tcs_epilog(Program* program, void* pinfo, ac_shader_config* config,
+                       const struct aco_compiler_options* options,
+                       const struct aco_shader_info* info, const struct ac_shader_args* args);
+
 void lower_phis(Program* program);
 void calc_min_waves(Program* program);
 void update_vgpr_sgpr_demand(Program* program, const RegisterDemand new_demand);
diff --git a/src/amd/compiler/aco_shader_info.h b/src/amd/compiler/aco_shader_info.h
index c234f9e635c..b057cb356f4 100644
--- a/src/amd/compiler/aco_shader_info.h
+++ b/src/amd/compiler/aco_shader_info.h
@@ -77,6 +77,11 @@ struct aco_ps_epilog_info {
    bool mrt0_is_dual_src;
 };
 
+struct aco_tcs_epilog_info {
+   enum tess_primitive_mode primitive_mode;
+   bool tes_reads_tessfactors;
+};
+
 struct aco_shader_info {
    enum ac_hw_stage hw_stage;
    uint8_t wave_size;