pco: update virtual register support for bools and nir reg translation

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/.clang-format

@@ -228,6 +228,7 @@ ForEachMacros: [
   'nir_foreach_block',
   'nir_foreach_block_safe',
   'nir_foreach_block_unstructured',
+  'nir_foreach_reg_decl',
   'nir_foreach_entrypoint',
   'nir_foreach_entrypoint_safe',
   'nir_foreach_function',

src/imagination/pco/pco_bool.c

@@ -27,7 +27,7 @@ static bool lower_bools(pco_func *func)
 
    /* Update to 32-bit. */
    pco_foreach_instr_in_func (instr, func) {
-      pco_foreach_instr_dest_ssa (pdest, instr) {
+      pco_foreach_instr_dest_vreg_ssa (pdest, instr) {
          if (pco_ref_get_bits(*pdest) != 1)
             continue;
 
@@ -36,7 +36,7 @@ static bool lower_bools(pco_func *func)
          progress = true;
       }
 
-      pco_foreach_instr_src_ssa (psrc, instr) {
+      pco_foreach_instr_src_vreg_ssa (psrc, instr) {
          if (pco_ref_get_bits(*psrc) != 1)
             continue;
 

src/imagination/pco/pco_internal.h

@@ -2184,24 +2184,26 @@ static inline pco_ref pco_ref_new_vreg(pco_func *func)
 }
 
 /**
- * \brief Builds and returns a scalar hardware register reference.
+ * \brief Remaps an SSA register to an existing virtual register.
  *
- * \param[in] index Register index.
- * \param[in] reg_class Register class.
- * \return Hardware register reference.
+ * \param[in,out] func The function.
+ * \param[in] ref Base SSA reference.
+ * \param[in] bits New bit width, or 0 if unchanged.
+ * \return Virtual register reference.
  */
-static inline pco_ref pco_ref_hwreg(unsigned index,
-                                    enum pco_reg_class reg_class)
+static inline pco_ref
+pco_ref_ssa_vreg(ASSERTED pco_func *func, pco_ref ref, unsigned bits)
 {
-   assert(index < 256);
-   assert(reg_class != PCO_REG_CLASS_VIRT);
+   assert(pco_ref_is_ssa(ref));
+   assert(ref.val < func->next_vreg);
 
-   return (pco_ref){
-      .val = index,
-      .bits = PCO_BITS_32,
-      .type = PCO_REF_TYPE_REG,
-      .reg_class = reg_class,
-   };
+   ref.type = PCO_REF_TYPE_REG;
+   ref.reg_class = PCO_REG_CLASS_VIRT;
+
+   if (bits)
+      ref.bits = pco_bits(bits);
+
+   return ref;
 }
 
 /**
@@ -2227,6 +2229,19 @@ pco_ref_hwreg_vec(unsigned index, enum pco_reg_class reg_class, unsigned chans)
    };
 }
 
+/**
+ * \brief Builds and returns a scalar hardware register reference.
+ *
+ * \param[in] index Register index.
+ * \param[in] reg_class Register class.
+ * \return Hardware register reference.
+ */
+static inline pco_ref pco_ref_hwreg(unsigned index,
+                                    enum pco_reg_class reg_class)
+{
+   return pco_ref_hwreg_vec(index, reg_class, 1);
+}
+
 /**
  * \brief Builds 32x1[2] hardware register address component references.
  *

src/imagination/pco/pco_opt.c

@@ -258,6 +258,9 @@ static inline bool try_back_prop_instr(struct pco_use *uses, pco_instr *instr)
 
    pco_ref *pdest_from = &use->instr->dest[0];
 
+   if (pco_ref_is_vreg(*pdest_from))
+      return false;
+
    assert(pco_ref_get_bits(*pdest_from) == pco_ref_get_bits(*pdest_to));
    assert(pco_ref_get_chans(*pdest_from) == pco_ref_get_chans(*pdest_to));
    assert(!pco_ref_has_mods_set(*pdest_from) &&

src/imagination/pco/pco_trans_nir.c

@@ -909,6 +909,69 @@ static pco_instr *lower_alphatst(trans_ctx *tctx,
                    pco_zero);
 }
 
+static inline unsigned lookup_reg_bits(nir_intrinsic_instr *intr)
+{
+   nir_def *reg;
+   switch (intr->intrinsic) {
+   case nir_intrinsic_load_reg:
+      reg = intr->src[0].ssa;
+      break;
+
+   case nir_intrinsic_store_reg:
+      reg = intr->src[1].ssa;
+      break;
+
+   default:
+      UNREACHABLE("");
+   }
+
+   nir_intrinsic_instr *reg_decl = nir_reg_get_decl(reg);
+   return nir_intrinsic_bit_size(reg_decl);
+}
+
+static pco_instr *trans_reg_intr(trans_ctx *tctx,
+                                 nir_intrinsic_instr *intr,
+                                 pco_ref dest,
+                                 pco_ref src0,
+                                 pco_ref src1)
+{
+   pco_func *func = tctx->func;
+
+   /* Special case; consume and reserve. */
+   if (intr->intrinsic == nir_intrinsic_decl_reg) {
+      assert(nir_intrinsic_num_components(intr) == 1);
+      assert(!nir_intrinsic_num_array_elems(intr));
+
+      /* nir_trivialize_registers puts decl_regs into the start block. */
+      assert(func->next_vreg == dest.val);
+      ++func->next_vreg;
+
+      return NULL;
+   }
+
+   unsigned reg_bits = lookup_reg_bits(intr);
+   switch (intr->intrinsic) {
+   case nir_intrinsic_load_reg:
+      assert(!nir_intrinsic_base(intr));
+      assert(!nir_intrinsic_legacy_fabs(intr));
+      assert(!nir_intrinsic_legacy_fneg(intr));
+
+      return pco_mov(&tctx->b, dest, pco_ref_ssa_vreg(func, src0, reg_bits));
+
+   case nir_intrinsic_store_reg:
+      assert(!nir_intrinsic_base(intr));
+      assert(nir_intrinsic_write_mask(intr) == 1);
+      assert(!nir_intrinsic_legacy_fsat(intr));
+
+      return pco_mov(&tctx->b, pco_ref_ssa_vreg(func, src1, reg_bits), src0);
+
+   default:
+      break;
+   }
+
+   UNREACHABLE("");
+}
+
 /**
  * \brief Translates a NIR intrinsic instruction into PCO.
  *
@@ -1025,6 +1088,12 @@ static pco_instr *trans_intr(trans_ctx *tctx, nir_intrinsic_instr *intr)
       instr = lower_alphatst(tctx, dest, src[0], src[1], src[2]);
       break;
 
+   case nir_intrinsic_decl_reg:
+   case nir_intrinsic_load_reg:
+   case nir_intrinsic_store_reg:
+      instr = trans_reg_intr(tctx, intr, dest, src[0], src[1]);
+      break;
+
    default:
       printf("Unsupported intrinsic: \"");
       nir_print_instr(&intr->instr, stdout);