ac/llvm,radeonsi: lower ps color load in nir

Remove the color0/1 in ac_shader_abi which is used by
radeonsi only.

Reviewed-by: Marek Olšák <marek.olsak@amd.com>
Signed-off-by: Qiang Yu <yuq825@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/21683>

src/amd/llvm/ac_nir_to_llvm.c

@@ -3544,12 +3544,6 @@ static bool visit_intrinsic(struct ac_nir_context *ctx, nir_intrinsic_instr *ins
    case nir_intrinsic_is_helper_invocation:
       result = ac_build_is_helper_invocation(&ctx->ac);
       break;
-   case nir_intrinsic_load_color0:
-      result = ctx->abi->color0;
-      break;
-   case nir_intrinsic_load_color1:
-      result = ctx->abi->color1;
-      break;
    case nir_intrinsic_load_user_data_amd:
       assert(LLVMTypeOf(ctx->abi->user_data) == ctx->ac.v4i32);
       result = ctx->abi->user_data;

src/amd/llvm/ac_shader_abi.h

@@ -47,7 +47,6 @@ struct ac_shader_abi {
    LLVMValueRef vs_rel_patch_id;
    LLVMValueRef instance_id;
    LLVMValueRef persp_centroid, linear_centroid;
-   LLVMValueRef color0, color1;
    LLVMValueRef user_data;
 
    /* replaced registers when culling enabled */

src/gallium/drivers/radeonsi/si_nir_lower_abi.c

@@ -488,6 +488,29 @@ static bool lower_abi_instr(nir_builder *b, nir_instr *instr, struct lower_abi_s
       replacement = nir_ilt(b, prim_mask, nir_imm_int(b, 0));
       break;
    }
+   case nir_intrinsic_load_color0:
+   case nir_intrinsic_load_color1: {
+      uint32_t colors_read = sel->info.colors_read;
+
+      int start, offset;
+      if (intrin->intrinsic == nir_intrinsic_load_color0) {
+         start = 0;
+         offset = 0;
+      } else {
+         start = 4;
+         offset = util_bitcount(colors_read & 0xf);
+      }
+
+      nir_ssa_def *color[4];
+      for (int i = 0; i < 4; i++) {
+         color[i] = colors_read & BITFIELD_BIT(start + i) ?
+            ac_nir_load_arg_at_offset(b, &args->ac, args->color_start, offset++) :
+            nir_ssa_undef(b, 1, 32);
+      }
+
+      replacement = nir_vec(b, color, 4);
+      break;
+   }
    default:
       return false;
    }

src/gallium/drivers/radeonsi/si_shader.c

@@ -706,7 +706,7 @@ void si_init_shader_args(struct si_shader *shader, struct si_shader_args *args)
          unsigned num_color_elements = util_bitcount(shader->selector->info.colors_read);
 
          for (i = 0; i < num_color_elements; i++)
-            ac_add_arg(&args->ac, AC_ARG_VGPR, 1, AC_ARG_FLOAT, NULL);
+            ac_add_arg(&args->ac, AC_ARG_VGPR, 1, AC_ARG_FLOAT, i ? NULL : &args->color_start);
 
          num_prolog_vgprs += num_color_elements;
       }

src/gallium/drivers/radeonsi/si_shader_internal.h

@@ -82,6 +82,7 @@ struct si_shader_args {
    /* PS */
    struct ac_arg pos_fixed_pt;
    struct ac_arg alpha_reference;
+   struct ac_arg color_start;
    /* CS */
    struct ac_arg block_size;
    struct ac_arg cs_user_data;

src/gallium/drivers/radeonsi/si_shader_llvm.c

@@ -832,32 +832,6 @@ static bool si_llvm_translate_nir(struct si_shader_context *ctx, struct si_shade
       break;
 
    case MESA_SHADER_FRAGMENT: {
-      unsigned colors_read = ctx->shader->selector->info.colors_read;
-      LLVMValueRef main_fn = ctx->main_fn.value;
-
-      LLVMValueRef undef = LLVMGetUndef(ctx->ac.f32);
-
-      unsigned offset = SI_PARAM_POS_FIXED_PT + 1;
-
-      if (colors_read & 0x0f) {
-         unsigned mask = colors_read & 0x0f;
-         LLVMValueRef values[4];
-         values[0] = mask & 0x1 ? LLVMGetParam(main_fn, offset++) : undef;
-         values[1] = mask & 0x2 ? LLVMGetParam(main_fn, offset++) : undef;
-         values[2] = mask & 0x4 ? LLVMGetParam(main_fn, offset++) : undef;
-         values[3] = mask & 0x8 ? LLVMGetParam(main_fn, offset++) : undef;
-         ctx->abi.color0 = ac_to_integer(&ctx->ac, ac_build_gather_values(&ctx->ac, values, 4));
-      }
-      if (colors_read & 0xf0) {
-         unsigned mask = (colors_read & 0xf0) >> 4;
-         LLVMValueRef values[4];
-         values[0] = mask & 0x1 ? LLVMGetParam(main_fn, offset++) : undef;
-         values[1] = mask & 0x2 ? LLVMGetParam(main_fn, offset++) : undef;
-         values[2] = mask & 0x4 ? LLVMGetParam(main_fn, offset++) : undef;
-         values[3] = mask & 0x8 ? LLVMGetParam(main_fn, offset++) : undef;
-         ctx->abi.color1 = ac_to_integer(&ctx->ac, ac_build_gather_values(&ctx->ac, values, 4));
-      }
-
       ctx->abi.num_interp = si_get_ps_num_interp(shader);
 
       ctx->abi.kill_ps_if_inf_interp =