/*
 * Copyright © 2020 Valve Corporation
 * SPDX-License-Identifier: MIT
 */

#include "tu_private.h"
#include "nir_builder.h"

/* Some a6xx variants cannot support a non-contiguous multiview mask. Instead,
 * inside the shader something like this needs to be inserted:
 *
 * gl_Position = ((1ull << gl_ViewIndex) & view_mask) ? gl_Position : vec4(0.);
 *
 * Scan backwards until we find the gl_Position write (there should only be
 * one).
 */
static bool
lower_multiview_mask(nir_shader *nir, uint32_t *mask)
{
   nir_function_impl *impl = nir_shader_get_entrypoint(nir);

   if (util_is_power_of_two_or_zero(*mask + 1)) {
      nir_metadata_preserve(impl, nir_metadata_all);
      return false;
   }

   nir_builder b;
   nir_builder_init(&b, impl);

   uint32_t old_mask = *mask;
   *mask = BIT(util_logbase2(old_mask) + 1) - 1;

   nir_foreach_block_reverse(block, impl) {
      nir_foreach_instr_reverse(instr, block) {
         if (instr->type != nir_instr_type_intrinsic)
            continue;

         nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
         if (intrin->intrinsic != nir_intrinsic_store_deref)
            continue;

         nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
         if (deref->mode != nir_var_shader_out)
            continue;

         nir_variable *var = nir_deref_instr_get_variable(deref);
         if (var->data.location != VARYING_SLOT_POS)
            continue;

         assert(intrin->src[1].is_ssa);
         nir_ssa_def *orig_src = intrin->src[1].ssa;
         b.cursor = nir_before_instr(instr);

         /* ((1ull << gl_ViewIndex) & mask) != 0 */
         nir_ssa_def *cmp =
            nir_i2b(&b, nir_iand(&b, nir_imm_int(&b, old_mask),
                                  nir_ishl(&b, nir_imm_int(&b, 1),
                                           nir_load_view_index(&b))));

         nir_ssa_def *src = nir_bcsel(&b, cmp, orig_src, nir_imm_float(&b, 0.));
         nir_instr_rewrite_src(instr, &intrin->src[1], nir_src_for_ssa(src));

         nir_metadata_preserve(impl, nir_metadata_block_index |
                                     nir_metadata_dominance);
         return true;
      }
   }

   nir_metadata_preserve(impl, nir_metadata_all);
   return false;
}

bool
tu_nir_lower_multiview(nir_shader *nir, uint32_t mask, bool *multi_pos_output,
                       struct tu_device *dev)
{
   *multi_pos_output = false;

   bool progress = false;

   if (!dev->physical_device->supports_multiview_mask)
      NIR_PASS(progress, nir, lower_multiview_mask, &mask);

   unsigned num_views = util_logbase2(mask) + 1;

   /* Speculatively assign output locations so that we know num_outputs. We
    * will assign output locations for real after this pass.
    */
   unsigned num_outputs;
   nir_assign_io_var_locations(nir, nir_var_shader_out, &num_outputs, MESA_SHADER_VERTEX);

   /* In addition to the generic checks done by NIR, check that we don't
    * overflow VPC with the extra copies of gl_Position.
    */
   if (likely(!(dev->physical_device->instance->debug_flags & TU_DEBUG_NOMULTIPOS)) &&
       num_outputs + (num_views - 1) <= 32 && nir_can_lower_multiview(nir)) {
      *multi_pos_output = true;

      /* It appears that the multiview mask is ignored when multi-position
       * output is enabled, so we have to write 0 to inactive views ourselves.
       */
      NIR_PASS(progress, nir, lower_multiview_mask, &mask);

      NIR_PASS_V(nir, nir_lower_multiview, mask);
      progress = true;
   }

   return progress;
}