/*
 * Copyright © 2017 Intel Corporation
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included
 * in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */

/**
 * @file iris_program_cache.c
 *
 * The in-memory program cache.  This is basically a hash table mapping
 * API-specified shaders and a state key to a compiled variant.  It also
 * takes care of uploading shader assembly into a BO for use on the GPU.
 */

#include <stdio.h>
#include <errno.h>
#include "pipe/p_defines.h"
#include "pipe/p_state.h"
#include "pipe/p_context.h"
#include "pipe/p_screen.h"
#include "util/u_atomic.h"
#include "util/u_upload_mgr.h"
#include "compiler/nir/nir.h"
#include "compiler/nir/nir_builder.h"
#include "intel/compiler/brw_compiler.h"
#include "intel/compiler/brw_nir.h"
#include "intel/compiler/elk/elk_compiler.h"
#include "intel/compiler/elk/elk_nir.h"
#include "iris_context.h"
#include "iris_resource.h"

struct keybox {
   uint16_t size;
   enum iris_program_cache_id cache_id;
   uint8_t data[0];
};

static struct keybox *
make_keybox(void *mem_ctx,
            enum iris_program_cache_id cache_id,
            const void *key,
            uint32_t key_size)
{
   struct keybox *keybox =
      ralloc_size(mem_ctx, sizeof(struct keybox) + key_size);

   keybox->cache_id = cache_id;
   keybox->size = key_size;
   memcpy(keybox->data, key, key_size);

   return keybox;
}

static uint32_t
keybox_hash(const void *void_key)
{
   const struct keybox *key = void_key;
   return _mesa_hash_data(&key->cache_id, key->size + sizeof(key->cache_id));
}

static bool
keybox_equals(const void *void_a, const void *void_b)
{
   const struct keybox *a = void_a, *b = void_b;
   if (a->size != b->size)
      return false;

   return memcmp(a->data, b->data, a->size) == 0;
}

struct iris_compiled_shader *
iris_find_cached_shader(struct iris_context *ice,
                        enum iris_program_cache_id cache_id,
                        uint32_t key_size,
                        const void *key)
{
   struct keybox *keybox = make_keybox(NULL, cache_id, key, key_size);
   struct hash_entry *entry =
      _mesa_hash_table_search(ice->shaders.cache, keybox);

   ralloc_free(keybox);

   return entry ? entry->data : NULL;
}

void
iris_delete_shader_variant(struct iris_compiled_shader *shader)
{
   pipe_resource_reference(&shader->assembly.res, NULL);
   util_queue_fence_destroy(&shader->ready);
   ralloc_free(shader);
}

struct iris_compiled_shader *
iris_create_shader_variant(const struct iris_screen *screen,
                           void *mem_ctx,
                           gl_shader_stage stage,
                           enum iris_program_cache_id cache_id,
                           uint32_t key_size,
                           const void *key)
{
#ifndef NDEBUG
   if (cache_id == IRIS_CACHE_BLORP) {
      /* Blorp shader must have a mem_ctx. */
      assert(mem_ctx != NULL);
   } else if (cache_id == IRIS_CACHE_TCS) {
      /* Pass-through tessellation control shaders (generated by the driver)
       * will have a mem_ctx, and other tessellation control shaders will not.
       */
   } else {
      /* Shaders that are neither blorp nor tessellation control must not have
       * a mem_ctx.
       */
      assert(mem_ctx == NULL);
   }
#endif

   struct iris_compiled_shader *shader =
      rzalloc_size(mem_ctx, sizeof(struct iris_compiled_shader) +
                   screen->vtbl.derived_program_state_size(cache_id));

   pipe_reference_init(&shader->ref, 1);
   util_queue_fence_init(&shader->ready);
   util_queue_fence_reset(&shader->ready);

   if (cache_id != IRIS_CACHE_BLORP) {
      assert(key_size <= sizeof(union iris_any_prog_key));
      memcpy(&shader->key, key, key_size);
   }

   shader->stage = stage;

   return shader;
}

void
iris_upload_shader(struct iris_screen *screen,
                   struct iris_uncompiled_shader *ish,
                   struct iris_compiled_shader *shader,
                   struct hash_table *driver_shaders,
                   struct u_upload_mgr *uploader,
                   enum iris_program_cache_id cache_id,
                   uint32_t key_size,
                   const void *key,
                   const void *assembly)
{
   const struct intel_device_info *devinfo = screen->devinfo;

   u_upload_alloc(uploader, 0, shader->program_size, 64,
                  &shader->assembly.offset, &shader->assembly.res,
                  &shader->map);
   memcpy(shader->map, assembly, shader->program_size);

   struct iris_resource *res = (void *) shader->assembly.res;
   uint64_t shader_data_addr = res->bo->address +
                               shader->assembly.offset +
                               shader->const_data_offset;

   if (screen->brw) {
      struct brw_shader_reloc_value reloc_values[] = {
         {
            .id = BRW_SHADER_RELOC_CONST_DATA_ADDR_LOW,
            .value = shader_data_addr,
         },
         {
            .id = BRW_SHADER_RELOC_CONST_DATA_ADDR_HIGH,
            .value = shader_data_addr >> 32,
         },
      };
      brw_write_shader_relocs(&screen->brw->isa, shader->map,
                              shader->brw_prog_data, reloc_values,
                              ARRAY_SIZE(reloc_values));
   } else {
      struct elk_shader_reloc_value reloc_values[] = {
         {
            .id = BRW_SHADER_RELOC_CONST_DATA_ADDR_LOW,
            .value = shader_data_addr,
         },
         {
            .id = BRW_SHADER_RELOC_CONST_DATA_ADDR_HIGH,
            .value = shader_data_addr >> 32,
         },
      };
      elk_write_shader_relocs(&screen->elk->isa, shader->map,
                              shader->elk_prog_data, reloc_values,
                              ARRAY_SIZE(reloc_values));
   }

   /* Store the 3DSTATE shader packets and other derived state. */
   screen->vtbl.store_derived_program_state(devinfo, cache_id, shader);

   util_queue_fence_signal(&shader->ready);

   if (!ish) {
      struct keybox *keybox = make_keybox(shader, cache_id, key, key_size);
      _mesa_hash_table_insert(driver_shaders, keybox, shader);
   }
}

bool
iris_blorp_lookup_shader(struct blorp_batch *blorp_batch,
                         const void *key, uint32_t key_size,
                         uint32_t *kernel_out, void *prog_data_out)
{
   struct blorp_context *blorp = blorp_batch->blorp;
   struct iris_context *ice = blorp->driver_ctx;
   struct iris_batch *batch = blorp_batch->driver_batch;
   struct iris_screen *screen = batch->screen;
   struct iris_compiled_shader *shader =
      iris_find_cached_shader(ice, IRIS_CACHE_BLORP, key_size, key);

   if (!shader)
      return false;

   struct iris_bo *bo = iris_resource_bo(shader->assembly.res);
   *kernel_out =
      iris_bo_offset_from_base_address(bo) + shader->assembly.offset;
   *((void **) prog_data_out) = screen->brw ? (void *)shader->brw_prog_data
                                            : (void *)shader->elk_prog_data;

   iris_use_pinned_bo(batch, bo, false, IRIS_DOMAIN_NONE);

   return true;
}

bool
iris_blorp_upload_shader(struct blorp_batch *blorp_batch, uint32_t stage,
                         const void *key, uint32_t key_size,
                         const void *kernel, UNUSED uint32_t kernel_size,
                         const void *prog_data_templ,
                         UNUSED uint32_t prog_data_size,
                         uint32_t *kernel_out, void *prog_data_out)
{
   struct blorp_context *blorp = blorp_batch->blorp;
   struct iris_context *ice = blorp->driver_ctx;
   struct iris_batch *batch = blorp_batch->driver_batch;
   struct iris_screen *screen = batch->screen;

   struct iris_binding_table bt;
   memset(&bt, 0, sizeof(bt));

   struct iris_compiled_shader *shader =
      iris_create_shader_variant(screen, ice->shaders.cache, stage,
                                 IRIS_CACHE_BLORP, key_size, key);

   void *prog_data = ralloc_size(NULL, prog_data_size);
   memcpy(prog_data, prog_data_templ, prog_data_size);

   if (screen->brw) {
      iris_apply_brw_prog_data(shader, prog_data);
   } else {
      assert(screen->elk);
      iris_apply_elk_prog_data(shader, prog_data);
   }

   iris_finalize_program(shader, NULL, NULL, 0, 0, 0, &bt);

   iris_upload_shader(screen, NULL, shader, ice->shaders.cache,
                      ice->shaders.uploader_driver,
                      IRIS_CACHE_BLORP, key_size, key, kernel);

   struct iris_bo *bo = iris_resource_bo(shader->assembly.res);
   *kernel_out =
      iris_bo_offset_from_base_address(bo) + shader->assembly.offset;
   *((void **) prog_data_out) = screen->brw ? (void *)shader->brw_prog_data
                                            : (void *)shader->elk_prog_data;

   iris_use_pinned_bo(batch, bo, false, IRIS_DOMAIN_NONE);

   return true;
}

void
iris_init_program_cache(struct iris_context *ice)
{
   ice->shaders.cache =
      _mesa_hash_table_create(ice, keybox_hash, keybox_equals);

   ice->shaders.uploader_driver =
      u_upload_create(&ice->ctx, 64 * 1024,
                      PIPE_BIND_CUSTOM, PIPE_USAGE_IMMUTABLE,
                      IRIS_RESOURCE_FLAG_SHADER_MEMZONE |
                      IRIS_RESOURCE_FLAG_DEVICE_MEM);
   ice->shaders.uploader_unsync =
      u_upload_create(&ice->ctx, 64 * 1024,
                      PIPE_BIND_CUSTOM, PIPE_USAGE_IMMUTABLE,
                      IRIS_RESOURCE_FLAG_SHADER_MEMZONE |
                      IRIS_RESOURCE_FLAG_DEVICE_MEM);
}

void
iris_destroy_program_cache(struct iris_context *ice)
{
   for (int i = 0; i < MESA_SHADER_STAGES; i++) {
      iris_shader_variant_reference(&ice->shaders.prog[i], NULL);
   }
   iris_shader_variant_reference(&ice->shaders.last_vue_shader, NULL);

   hash_table_foreach(ice->shaders.cache, entry) {
      struct iris_compiled_shader *shader = entry->data;
      iris_delete_shader_variant(shader);
   }

   u_upload_destroy(ice->shaders.uploader_driver);
   u_upload_destroy(ice->shaders.uploader_unsync);

   ralloc_free(ice->shaders.cache);
}

static void
link_libintel_shaders(nir_shader *nir, const nir_shader *libintel)
{
   nir_link_shader_functions(nir, libintel);
   NIR_PASS_V(nir, nir_inline_functions);
   NIR_PASS_V(nir, nir_remove_non_entrypoints);
   NIR_PASS_V(nir, nir_lower_vars_to_explicit_types, nir_var_function_temp,
              glsl_get_cl_type_size_align);
   NIR_PASS_V(nir, nir_opt_deref);
   NIR_PASS_V(nir, nir_lower_vars_to_ssa);
   NIR_PASS_V(nir, nir_lower_explicit_io,
              nir_var_shader_temp | nir_var_function_temp | nir_var_mem_shared |
                 nir_var_mem_global,
              nir_address_format_62bit_generic);
}

void
iris_ensure_indirect_generation_shader(struct iris_batch *batch)
{
   struct iris_context *ice = batch->ice;
   if (ice->draw.generation.shader)
      return;

   struct iris_screen *screen = (struct iris_screen *)ice->ctx.screen;
   const struct {
      char name[40];
   } key = {
      .name = "iris-generation-shader",
   };
   ice->draw.generation.shader =
      iris_find_cached_shader(ice, IRIS_CACHE_BLORP, sizeof(key), &key);
   if (ice->draw.generation.shader != NULL)
      return;

   const nir_shader_compiler_options *nir_options =
      screen->brw ? screen->brw->nir_options[MESA_SHADER_COMPUTE]
                  : screen->elk->nir_options[MESA_SHADER_COMPUTE];

   nir_builder b = nir_builder_init_simple_shader(MESA_SHADER_FRAGMENT,
                                                  nir_options,
                                                  "iris-indirect-generate");

   uint32_t uniform_size =
      screen->vtbl.call_generation_shader(screen, &b);

   nir_shader *nir = b.shader;

   void *mem_ctx = ralloc_context(NULL);
   link_libintel_shaders(nir, screen->vtbl.load_shader_lib(screen, mem_ctx));

   NIR_PASS_V(nir, nir_lower_vars_to_ssa);
   NIR_PASS_V(nir, nir_opt_cse);
   NIR_PASS_V(nir, nir_opt_gcm, true);
   NIR_PASS_V(nir, nir_opt_peephole_select, 1, false, false);

   NIR_PASS_V(nir, nir_lower_variable_initializers, ~0);

   NIR_PASS_V(nir, nir_split_var_copies);
   NIR_PASS_V(nir, nir_split_per_member_structs);

   if (screen->brw) {
      struct brw_nir_compiler_opts opts = {};
      brw_preprocess_nir(screen->brw, nir, &opts);
   } else {
      assert(screen->elk);
      struct elk_nir_compiler_opts opts = {};
      elk_preprocess_nir(screen->elk, nir, &opts);
   }

   NIR_PASS_V(nir, nir_propagate_invariant, false);

   NIR_PASS_V(nir, nir_lower_input_attachments,
              &(nir_input_attachment_options) {
                 .use_fragcoord_sysval = true,
                 .use_layer_id_sysval = true,
              });

   /* Reset sizes before gathering information */
   nir->global_mem_size = 0;
   nir->scratch_size = 0;
   nir->info.shared_size = 0;
   nir_shader_gather_info(nir, nir_shader_get_entrypoint(nir));

   NIR_PASS_V(nir, nir_copy_prop);
   NIR_PASS_V(nir, nir_opt_constant_folding);
   NIR_PASS_V(nir, nir_opt_dce);

   /* Do vectorizing here. For some reason when trying to do it in the back
    * this just isn't working.
    */
   nir_load_store_vectorize_options options = {
      .modes = nir_var_mem_ubo | nir_var_mem_ssbo | nir_var_mem_global,
      .callback = brw_nir_should_vectorize_mem,
      .robust_modes = (nir_variable_mode)0,
   };
   NIR_PASS_V(nir, nir_opt_load_store_vectorize, &options);

   nir->num_uniforms = uniform_size;

   struct iris_compiled_shader *shader =
      iris_create_shader_variant(screen, ice->shaders.cache,
                                 MESA_SHADER_FRAGMENT,
                                 IRIS_CACHE_BLORP,
                                 sizeof(key), &key);

   const unsigned *program;
   if (screen->brw) {
      union brw_any_prog_key prog_key;
      memset(&prog_key, 0, sizeof(prog_key));

      struct brw_wm_prog_data *prog_data = ralloc_size(NULL, sizeof(*prog_data));
      memset(prog_data, 0, sizeof(*prog_data));
      prog_data->base.nr_params = nir->num_uniforms / 4;

      brw_nir_analyze_ubo_ranges(screen->brw, nir, prog_data->base.ubo_ranges);

      struct brw_compile_stats stats[3];
      struct brw_compile_fs_params params = {
         .base = {
            .nir = nir,
            .log_data = &ice->dbg,
            .debug_flag = DEBUG_WM,
            .stats = stats,
            .mem_ctx = mem_ctx,
         },
         .key = &prog_key.wm,
         .prog_data = prog_data,
      };
      program = brw_compile_fs(screen->brw, &params);
      assert(program);
      iris_apply_brw_prog_data(shader, &prog_data->base);
   } else {
      union elk_any_prog_key prog_key;
      memset(&prog_key, 0, sizeof(prog_key));

      struct elk_wm_prog_data *prog_data = ralloc_size(NULL, sizeof(*prog_data));
      memset(prog_data, 0, sizeof(*prog_data));
      prog_data->base.nr_params = nir->num_uniforms / 4;

      elk_nir_analyze_ubo_ranges(screen->elk, nir, prog_data->base.ubo_ranges);

      struct elk_compile_stats stats[3];
      struct elk_compile_fs_params params = {
         .base = {
            .nir = nir,
            .log_data = &ice->dbg,
            .debug_flag = DEBUG_WM,
            .stats = stats,
            .mem_ctx = mem_ctx,
         },
         .key = &prog_key.wm,
         .prog_data = prog_data,
      };
      program = elk_compile_fs(screen->elk, &params);
      assert(program);
      iris_apply_elk_prog_data(shader, &prog_data->base);
   }

   struct iris_binding_table bt;
   memset(&bt, 0, sizeof(bt));

   iris_finalize_program(shader, NULL, NULL, 0, 0, 0, &bt);

   iris_upload_shader(screen, NULL, shader, ice->shaders.cache,
                      ice->shaders.uploader_driver,
                      IRIS_CACHE_BLORP, sizeof(key), &key, program);

   ralloc_free(mem_ctx);

   struct iris_bo *bo = iris_resource_bo(shader->assembly.res);
   iris_use_pinned_bo(batch, bo, false, IRIS_DOMAIN_NONE);

   ice->draw.generation.shader = shader;
}