diff --git a/src/amd/common/ac_shader_args.c b/src/amd/common/ac_shader_args.c
index d3816e1c0fc..2025cc75f82 100644
--- a/src/amd/common/ac_shader_args.c
+++ b/src/amd/common/ac_shader_args.c
@@ -52,3 +52,19 @@ void ac_add_arg(struct ac_shader_args *info, enum ac_arg_regfile regfile, unsign
 
    info->arg_count++;
 }
+
+void ac_add_return(struct ac_shader_args *info, enum ac_arg_regfile regfile)
+{
+   assert(info->return_count < AC_MAX_ARGS);
+
+   if (regfile == AC_ARG_SGPR) {
+      /* SGPRs must be inserted before VGPRs. */
+      assert(info->num_vgprs_returned == 0);
+      info->num_sgprs_returned++;;
+   } else {
+      assert(regfile == AC_ARG_VGPR);
+      info->num_vgprs_returned++;
+   }
+
+   info->return_count++;
+}
diff --git a/src/amd/common/ac_shader_args.h b/src/amd/common/ac_shader_args.h
index 80274374fe6..020e8b46ba8 100644
--- a/src/amd/common/ac_shader_args.h
+++ b/src/amd/common/ac_shader_args.h
@@ -67,6 +67,10 @@ struct ac_shader_args {
    uint16_t num_sgprs_used;
    uint16_t num_vgprs_used;
 
+   uint16_t return_count;
+   uint16_t num_sgprs_returned;
+   uint16_t num_vgprs_returned;
+
    struct ac_arg base_vertex;
    struct ac_arg start_instance;
    struct ac_arg draw_id;
@@ -108,5 +112,6 @@ struct ac_shader_args {
 
 void ac_add_arg(struct ac_shader_args *info, enum ac_arg_regfile regfile, unsigned registers,
                 enum ac_arg_type type, struct ac_arg *arg);
+void ac_add_return(struct ac_shader_args *info, enum ac_arg_regfile regfile);
 
 #endif