pvr, pco: temporary initial scratch memory support

Signed-off-by: Simon Perretta <simon.perretta@imgtec.com>
Acked-by: Erik Faye-Lund <erik.faye-lund@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/36412>

src/imagination/pco/pco_data.h

@@ -199,6 +199,9 @@ typedef struct _pco_common_data {
    unsigned spilled_temps;
    pco_range spill_info; /* addr_lo, addr_hi, block_size */
 
+   unsigned scratch;
+   pco_range scratch_info; /* addr_lo, addr_hi, block_size */
+
    unsigned coeffs; /** Number of allocated coefficient registers. */
    unsigned shareds; /** Number of allocated shared registers. */
 

src/imagination/pco/pco_nir.c

@@ -350,6 +350,8 @@ static void gather_common_data(nir_shader *nir, pco_data *data)
                               gather_common_data_pass,
                               nir_metadata_all,
                               data);
+
+   data->common.scratch = nir->scratch_size;
 }
 
 /**
@@ -536,6 +538,17 @@ void pco_preprocess_nir(pco_ctx *ctx, nir_shader *nir)
 
    NIR_PASS(_, nir, nir_lower_vars_to_ssa);
 
+   if (!nir->info.internal) {
+      /* TODO: test with different size_threshold values. */
+      NIR_PASS(_,
+               nir,
+               nir_lower_vars_to_scratch,
+               nir_var_function_temp,
+               8,
+               glsl_get_natural_size_align_bytes,
+               glsl_get_word_size_align_bytes);
+   }
+
    NIR_PASS(_,
             nir,
             nir_lower_indirect_derefs,

src/imagination/pco/pco_trans_nir.c

@@ -1195,6 +1195,85 @@ static pco_instr *trans_global_atomic_buffer(trans_ctx *tctx,
    UNREACHABLE("");
 }
 
+static pco_instr *trans_scratch(trans_ctx *tctx,
+                                pco_ref dest,
+                                pco_ref offset_src,
+                                pco_ref data_src)
+{
+   const pco_common_data *common = &tctx->shader->data.common;
+
+   assert(common->scratch_info.count > 0);
+   unsigned base_addr_idx = common->scratch_info.start;
+   unsigned block_size_idx = common->scratch_info.start + 2;
+
+   pco_ref base_addr[2];
+   pco_ref_hwreg_addr_comps(base_addr_idx, PCO_REG_CLASS_SHARED, base_addr);
+
+   pco_ref block_size = pco_ref_hwreg(block_size_idx, PCO_REG_CLASS_SHARED);
+   pco_ref local_addr_inst_num =
+      pco_ref_hwreg(PCO_SR_LOCAL_ADDR_INST_NUM, PCO_REG_CLASS_SPEC);
+
+   pco_ref inst_base_addr[2];
+   pco_ref_new_ssa_addr_comps(tctx->func, inst_base_addr);
+
+   pco_imadd64(&tctx->b,
+               inst_base_addr[0],
+               inst_base_addr[1],
+               block_size,
+               local_addr_inst_num,
+               base_addr[0],
+               base_addr[1],
+               pco_ref_null());
+
+   pco_ref addr_data_comps[3];
+   pco_ref_new_ssa_addr_comps(tctx->func, addr_data_comps);
+
+   pco_add64_32(&tctx->b,
+                addr_data_comps[0],
+                addr_data_comps[1],
+                inst_base_addr[0],
+                inst_base_addr[1],
+                offset_src,
+                pco_ref_null(),
+                .s = true);
+
+   bool is_load = pco_ref_is_null(data_src);
+   if (is_load) {
+      ASSERTED unsigned bits = pco_ref_get_bits(dest);
+
+      /* TODO: 8/16-bit support via masking. */
+      assert(bits == 32);
+
+      pco_ref addr = pco_ref_new_ssa_addr(tctx->func);
+      pco_vec(&tctx->b, addr, 2, addr_data_comps);
+
+      unsigned chans = pco_ref_get_chans(dest);
+      return pco_ld(&tctx->b,
+                    dest,
+                    pco_ref_drc(PCO_DRC_0),
+                    pco_ref_imm8(chans),
+                    addr);
+   }
+
+   unsigned chans = pco_ref_get_chans(data_src);
+   pco_ref addr_data = pco_ref_new_ssa_addr_data(tctx->func, chans);
+
+   addr_data_comps[2] = data_src;
+   pco_vec(&tctx->b, addr_data, ARRAY_SIZE(addr_data_comps), addr_data_comps);
+
+   pco_ref data_comp = pco_ref_new_ssa(tctx->func,
+                                       pco_ref_get_bits(data_src),
+                                       pco_ref_get_chans(data_src));
+   pco_comp(&tctx->b, data_comp, addr_data, pco_ref_val16(2));
+
+   return pco_st32(&tctx->b,
+                   data_comp,
+                   pco_ref_drc(PCO_DRC_0),
+                   pco_ref_imm8(chans),
+                   addr_data,
+                   pco_ref_null());
+}
+
 static inline enum pco_reg_class sys_val_to_reg_class(gl_system_value sys_val,
                                                       mesa_shader_stage stage)
 {
@@ -1627,6 +1706,14 @@ static pco_instr *trans_intr(trans_ctx *tctx, nir_intrinsic_instr *intr)
       instr = trans_global_atomic_buffer(tctx, intr, dest, src[0]);
       break;
 
+   case nir_intrinsic_load_scratch:
+      instr = trans_scratch(tctx, dest, src[0], src[1]);
+      break;
+
+   case nir_intrinsic_store_scratch:
+      instr = trans_scratch(tctx, dest, src[1], src[0]);
+      break;
+
    /* Vertex sysvals. */
    case nir_intrinsic_load_vertex_id:
    case nir_intrinsic_load_instance_id:

src/imagination/vulkan/pds/pvr_pds.h

@@ -907,6 +907,7 @@ struct pvr_pds_descriptor_set {
 #define PVR_BUFFER_TYPE_FS_META (10)
 #define PVR_BUFFER_TYPE_TILE_BUFFERS (11)
 #define PVR_BUFFER_TYPE_SPILL_INFO (12)
+#define PVR_BUFFER_TYPE_SCRATCH_INFO (13)
 #define PVR_BUFFER_TYPE_INVALID (~0)
 
 struct pvr_pds_buffer {

src/imagination/vulkan/pds/pvr_pipeline_pds.c

@@ -1582,7 +1582,8 @@ void pvr_pds_generate_descriptor_upload_program(
       case PVR_BUFFER_TYPE_FRONT_FACE_OP:
       case PVR_BUFFER_TYPE_FS_META:
       case PVR_BUFFER_TYPE_TILE_BUFFERS:
-      case PVR_BUFFER_TYPE_SPILL_INFO: {
+      case PVR_BUFFER_TYPE_SPILL_INFO:
+      case PVR_BUFFER_TYPE_SCRATCH_INFO: {
          struct pvr_const_map_entry_special_buffer *special_buffer_entry;
 
          special_buffer_entry =

src/imagination/vulkan/pvr_cmd_buffer.c

@@ -4020,6 +4020,43 @@ static VkResult pvr_setup_descriptor_mappings(
             break;
          }
 
+         case PVR_BUFFER_TYPE_SCRATCH_INFO: {
+            assert(data->common.scratch);
+            unsigned scratch_block_size = data->common.scratch;
+
+            /* TODO: 2048 is to account for each instance... do this
+             * programmatically!
+             */
+            struct pvr_suballoc_bo *scratch_buffer_bo;
+            result = pvr_cmd_buffer_upload_general(cmd_buffer,
+                                                   NULL,
+                                                   scratch_block_size * 2048,
+                                                   &scratch_buffer_bo);
+
+            if (result != VK_SUCCESS)
+               return result;
+
+            uint32_t scratch_info[3] = {
+               [0] = scratch_buffer_bo->dev_addr.addr & 0xffffffff,
+               [1] = scratch_buffer_bo->dev_addr.addr >> 32,
+               [2] = scratch_block_size,
+            };
+
+            struct pvr_suballoc_bo *scratch_info_bo;
+            result = pvr_cmd_buffer_upload_general(cmd_buffer,
+                                                   scratch_info,
+                                                   sizeof(scratch_info),
+                                                   &scratch_info_bo);
+
+            if (result != VK_SUCCESS)
+               return result;
+
+            PVR_WRITE(qword_buffer,
+                      scratch_info_bo->dev_addr.addr,
+                      special_buff_entry->const_offset,
+                      pds_info->data_size_in_dwords);
+            break;
+         }
          default:
             UNREACHABLE("Unsupported special buffer type.");
          }

src/imagination/vulkan/pvr_pipeline.c

@@ -623,6 +623,14 @@ static VkResult pvr_pds_descriptor_program_create_and_upload(
       };
    }
 
+   if (data->common.scratch_info.count > 0) {
+      program.buffers[program.buffer_count++] = (struct pvr_pds_buffer){
+         .type = PVR_BUFFER_TYPE_SCRATCH_INFO,
+         .size_in_dwords = data->common.scratch_info.count,
+         .destination = data->common.scratch_info.start,
+      };
+   }
+
    if (stage == MESA_SHADER_FRAGMENT &&
        data->common.sys_vals[SYSTEM_VALUE_FRONT_FACE].count > 0) {
       program.buffers[program.buffer_count++] = (struct pvr_pds_buffer){
@@ -2424,6 +2432,15 @@ static void pvr_setup_descriptors(pco_data *data,
 
       data->common.shareds += 3;
    }
+
+   if (data->common.scratch) {
+      data->common.scratch_info = (pco_range){
+         .start = data->common.shareds,
+         .count = 3,
+      };
+
+      data->common.shareds += 3;
+   }
 }
 
 static void