src/compiler/nir/nir_opt_call.c - platform/external/mesa3d - Git at Google

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

 #include "nir.h"
 #include "nir_builder.h"
 #include "nir_phi_builder.h"

 struct call_liveness_entry {
    struct list_head list;
    nir_call_instr *instr;
    const BITSET_WORD *live_set;
 };

 static bool
 can_remat_instr(nir_instr *instr)
 {
    switch (instr->type) {
    case nir_instr_type_alu:
    case nir_instr_type_load_const:
    case nir_instr_type_undef:
       return true;
    case nir_instr_type_intrinsic:
       switch (nir_instr_as_intrinsic(instr)->intrinsic) {
       case nir_intrinsic_load_ray_launch_id:
       case nir_intrinsic_load_ray_launch_size:
       case nir_intrinsic_vulkan_resource_index:
       case nir_intrinsic_vulkan_resource_reindex:
       case nir_intrinsic_load_vulkan_descriptor:
       case nir_intrinsic_load_push_constant:
       case nir_intrinsic_load_global_constant:
       case nir_intrinsic_load_smem_amd:
       case nir_intrinsic_load_scalar_arg_amd:
       case nir_intrinsic_load_vector_arg_amd:
          return true;
       default:
          return false;
       }
    default:
       return false;
    }
 }

 static void
 remat_ssa_def(nir_builder *b, nir_def *def, struct hash_table *remap_table,
               struct hash_table *phi_value_table,
               struct nir_phi_builder *phi_builder, BITSET_WORD *def_blocks)
 {
    memset(def_blocks, 0, BITSET_WORDS(b->impl->num_blocks) * sizeof(BITSET_WORD));
    BITSET_SET(def_blocks, def->parent_instr->block->index);
    BITSET_SET(def_blocks, nir_cursor_current_block(b->cursor)->index);
    struct nir_phi_builder_value *val =
       nir_phi_builder_add_value(phi_builder, def->num_components,
                                 def->bit_size, def_blocks);
    _mesa_hash_table_insert(phi_value_table, def, val);

    nir_instr *clone = nir_instr_clone_deep(b->shader, def->parent_instr,
                                            remap_table);
    nir_builder_instr_insert(b, clone);
    nir_def *new_def = nir_instr_def(clone);

    _mesa_hash_table_insert(remap_table, def, new_def);
    if (nir_cursor_current_block(b->cursor)->index !=
        def->parent_instr->block->index)
       nir_phi_builder_value_set_block_def(val, def->parent_instr->block, def);
    nir_phi_builder_value_set_block_def(val, nir_cursor_current_block(b->cursor),
                                        new_def);
 }

 struct remat_chain_check_data {
    struct hash_table *remap_table;
    unsigned chain_length;
 };

 static bool
 can_remat_chain(nir_src *src, void *data)
 {
    struct remat_chain_check_data *check_data = data;

    if (_mesa_hash_table_search(check_data->remap_table, src->ssa))
       return true;

    if (!can_remat_instr(src->ssa->parent_instr))
       return false;

    if (check_data->chain_length++ >= 16)
       return false;

    return nir_foreach_src(src->ssa->parent_instr, can_remat_chain, check_data);
 }

 struct remat_chain_data {
    nir_builder *b;
    struct hash_table *remap_table;
    struct hash_table *phi_value_table;
    struct nir_phi_builder *phi_builder;
    BITSET_WORD *def_blocks;
 };

 static bool
 do_remat_chain(nir_src *src, void *data)
 {
    struct remat_chain_data *remat_data = data;

    if (_mesa_hash_table_search(remat_data->remap_table, src->ssa))
       return true;

    nir_foreach_src(src->ssa->parent_instr, do_remat_chain, remat_data);

    remat_ssa_def(remat_data->b, src->ssa, remat_data->remap_table,
                  remat_data->phi_value_table, remat_data->phi_builder,
                  remat_data->def_blocks);
    return true;
 }

 static bool
 rewrite_instr_src_from_phi_builder(nir_src *src, void *data)
 {
    struct hash_table *phi_value_table = data;

    if (nir_src_is_const(*src)) {
       nir_builder b = nir_builder_at(nir_before_instr(nir_src_parent_instr(src)));
       nir_src_rewrite(src, nir_build_imm(&b, src->ssa->num_components,
                                          src->ssa->bit_size,
                                          nir_src_as_const_value(*src)));
       return true;
    }

    struct hash_entry *entry = _mesa_hash_table_search(phi_value_table, src->ssa);
    if (!entry)
       return true;

    nir_block *block = nir_src_parent_instr(src)->block;
    nir_def *new_def = nir_phi_builder_value_get_block_def(entry->data, block);

    bool can_rewrite = true;
    if (new_def->parent_instr->block == block && new_def->index != UINT32_MAX)
       can_rewrite =
          !nir_instr_is_before(nir_src_parent_instr(src), new_def->parent_instr);

    if (can_rewrite)
       nir_src_rewrite(src, new_def);
    return true;
 }

 static bool
 nir_minimize_call_live_states_impl(nir_function_impl *impl)
 {
    nir_metadata_require(impl, nir_metadata_block_index |
                                  nir_metadata_live_defs |
                                  nir_metadata_dominance);
    bool progress = false;
    void *mem_ctx = ralloc_context(NULL);

    struct list_head call_list;
    list_inithead(&call_list);
    unsigned num_defs = impl->ssa_alloc;

    nir_def **rematerializable =
       rzalloc_array_size(mem_ctx, sizeof(nir_def *), num_defs);

    nir_foreach_block(block, impl) {
       nir_foreach_instr(instr, block) {
          nir_def *def = nir_instr_def(instr);
          if (def &&
              can_remat_instr(instr)) {
             rematerializable[def->index] = def;
          }

          if (instr->type != nir_instr_type_call)
             continue;
          nir_call_instr *call = nir_instr_as_call(instr);
          if (!call->indirect_callee.ssa)
             continue;

          struct call_liveness_entry *entry =
             ralloc_size(mem_ctx, sizeof(struct call_liveness_entry));
          entry->instr = call;
          entry->live_set = nir_get_live_defs(nir_after_instr(instr), mem_ctx);
          list_addtail(&entry->list, &call_list);
       }
    }

    const unsigned block_words = BITSET_WORDS(impl->num_blocks);
    BITSET_WORD *def_blocks = ralloc_array(mem_ctx, BITSET_WORD, block_words);

    list_for_each_entry(struct call_liveness_entry, entry, &call_list, list) {
       unsigned i;

       nir_builder b = nir_builder_at(nir_after_instr(&entry->instr->instr));

       struct nir_phi_builder *builder = nir_phi_builder_create(impl);
       struct hash_table *phi_value_table =
          _mesa_pointer_hash_table_create(mem_ctx);
       struct hash_table *remap_table =
          _mesa_pointer_hash_table_create(mem_ctx);

       BITSET_FOREACH_SET(i, entry->live_set, num_defs) {
          if (!rematerializable[i] ||
              _mesa_hash_table_search(remap_table, rematerializable[i]))
             continue;

          assert(!_mesa_hash_table_search(phi_value_table, rematerializable[i]));

          struct remat_chain_check_data check_data = {
             .remap_table = remap_table,
             .chain_length = 1,
          };

          if (!nir_foreach_src(rematerializable[i]->parent_instr,
                               can_remat_chain, &check_data))
             continue;

          struct remat_chain_data remat_data = {
             .b = &b,
             .remap_table = remap_table,
             .phi_value_table = phi_value_table,
             .phi_builder = builder,
             .def_blocks = def_blocks,
          };

          nir_foreach_src(rematerializable[i]->parent_instr, do_remat_chain,
                          &remat_data);

          remat_ssa_def(&b, rematerializable[i], remap_table, phi_value_table,
                        builder, def_blocks);
          progress = true;
       }
       _mesa_hash_table_destroy(remap_table, NULL);

       nir_foreach_block(block, impl) {
          nir_foreach_instr(instr, block) {
             if (instr->type == nir_instr_type_phi)
                continue;

             nir_foreach_src(instr, rewrite_instr_src_from_phi_builder,
                             phi_value_table);
          }
       }

       nir_phi_builder_finish(builder);
       _mesa_hash_table_destroy(phi_value_table, NULL);
    }

    ralloc_free(mem_ctx);

    nir_progress(true, impl, nir_metadata_block_index | nir_metadata_dominance);
    return progress;
 }

 /* Tries to rematerialize as many live vars as possible after calls.
  * Note: nir_opt_cse will undo any rematerializations done by this pass,
  * so it shouldn't be run afterward.
  */
 bool
 nir_minimize_call_live_states(nir_shader *shader)
 {
    bool progress = false;

    nir_foreach_function_impl(impl, shader) {
       progress |= nir_minimize_call_live_states_impl(impl);
    }

    return progress;
 }
	/*
	* Copyright © 2024 Valve Corporation
	* SPDX-License-Identifier: MIT
	*/

	#include "nir.h"
	#include "nir_builder.h"
	#include "nir_phi_builder.h"

	struct call_liveness_entry {
	struct list_head list;
	nir_call_instr *instr;
	const BITSET_WORD *live_set;
	};

	static bool
	can_remat_instr(nir_instr *instr)
	{
	switch (instr->type) {
	case nir_instr_type_alu:
	case nir_instr_type_load_const:
	case nir_instr_type_undef:
	return true;
	case nir_instr_type_intrinsic:
	switch (nir_instr_as_intrinsic(instr)->intrinsic) {
	case nir_intrinsic_load_ray_launch_id:
	case nir_intrinsic_load_ray_launch_size:
	case nir_intrinsic_vulkan_resource_index:
	case nir_intrinsic_vulkan_resource_reindex:
	case nir_intrinsic_load_vulkan_descriptor:
	case nir_intrinsic_load_push_constant:
	case nir_intrinsic_load_global_constant:
	case nir_intrinsic_load_smem_amd:
	case nir_intrinsic_load_scalar_arg_amd:
	case nir_intrinsic_load_vector_arg_amd:
	return true;
	default:
	return false;
	}
	default:
	return false;
	}
	}

	static void
	remat_ssa_def(nir_builder b, nir_def def, struct hash_table *remap_table,
	struct hash_table *phi_value_table,
	struct nir_phi_builder phi_builder, BITSET_WORD def_blocks)
	{
	memset(def_blocks, 0, BITSET_WORDS(b->impl->num_blocks) * sizeof(BITSET_WORD));
	BITSET_SET(def_blocks, def->parent_instr->block->index);
	BITSET_SET(def_blocks, nir_cursor_current_block(b->cursor)->index);
	struct nir_phi_builder_value *val =
	nir_phi_builder_add_value(phi_builder, def->num_components,
	def->bit_size, def_blocks);
	_mesa_hash_table_insert(phi_value_table, def, val);

	nir_instr *clone = nir_instr_clone_deep(b->shader, def->parent_instr,
	remap_table);
	nir_builder_instr_insert(b, clone);
	nir_def *new_def = nir_instr_def(clone);

	_mesa_hash_table_insert(remap_table, def, new_def);
	if (nir_cursor_current_block(b->cursor)->index !=
	def->parent_instr->block->index)
	nir_phi_builder_value_set_block_def(val, def->parent_instr->block, def);
	nir_phi_builder_value_set_block_def(val, nir_cursor_current_block(b->cursor),
	new_def);
	}

	struct remat_chain_check_data {
	struct hash_table *remap_table;
	unsigned chain_length;
	};

	static bool
	can_remat_chain(nir_src src, void data)
	{
	struct remat_chain_check_data *check_data = data;

	if (_mesa_hash_table_search(check_data->remap_table, src->ssa))
	return true;

	if (!can_remat_instr(src->ssa->parent_instr))
	return false;

	if (check_data->chain_length++ >= 16)
	return false;

	return nir_foreach_src(src->ssa->parent_instr, can_remat_chain, check_data);
	}

	struct remat_chain_data {
	nir_builder *b;
	struct hash_table *remap_table;
	struct hash_table *phi_value_table;
	struct nir_phi_builder *phi_builder;
	BITSET_WORD *def_blocks;
	};

	static bool
	do_remat_chain(nir_src src, void data)
	{
	struct remat_chain_data *remat_data = data;

	if (_mesa_hash_table_search(remat_data->remap_table, src->ssa))
	return true;

	nir_foreach_src(src->ssa->parent_instr, do_remat_chain, remat_data);

	remat_ssa_def(remat_data->b, src->ssa, remat_data->remap_table,
	remat_data->phi_value_table, remat_data->phi_builder,
	remat_data->def_blocks);
	return true;
	}

	static bool
	rewrite_instr_src_from_phi_builder(nir_src src, void data)
	{
	struct hash_table *phi_value_table = data;

	if (nir_src_is_const(*src)) {
	nir_builder b = nir_builder_at(nir_before_instr(nir_src_parent_instr(src)));
	nir_src_rewrite(src, nir_build_imm(&b, src->ssa->num_components,
	src->ssa->bit_size,
	nir_src_as_const_value(*src)));
	return true;
	}

	struct hash_entry *entry = _mesa_hash_table_search(phi_value_table, src->ssa);
	if (!entry)
	return true;

	nir_block *block = nir_src_parent_instr(src)->block;
	nir_def *new_def = nir_phi_builder_value_get_block_def(entry->data, block);

	bool can_rewrite = true;
	if (new_def->parent_instr->block == block && new_def->index != UINT32_MAX)
	can_rewrite =
	!nir_instr_is_before(nir_src_parent_instr(src), new_def->parent_instr);

	if (can_rewrite)
	nir_src_rewrite(src, new_def);
	return true;
	}

	static bool
	nir_minimize_call_live_states_impl(nir_function_impl *impl)
	{
	nir_metadata_require(impl, nir_metadata_block_index \|
	nir_metadata_live_defs \|
	nir_metadata_dominance);
	bool progress = false;
	void *mem_ctx = ralloc_context(NULL);

	struct list_head call_list;
	list_inithead(&call_list);
	unsigned num_defs = impl->ssa_alloc;

	nir_def **rematerializable =
	rzalloc_array_size(mem_ctx, sizeof(nir_def *), num_defs);

	nir_foreach_block(block, impl) {
	nir_foreach_instr(instr, block) {
	nir_def *def = nir_instr_def(instr);
	if (def &&
	can_remat_instr(instr)) {
	rematerializable[def->index] = def;
	}

	if (instr->type != nir_instr_type_call)
	continue;
	nir_call_instr *call = nir_instr_as_call(instr);
	if (!call->indirect_callee.ssa)
	continue;

	struct call_liveness_entry *entry =
	ralloc_size(mem_ctx, sizeof(struct call_liveness_entry));
	entry->instr = call;
	entry->live_set = nir_get_live_defs(nir_after_instr(instr), mem_ctx);
	list_addtail(&entry->list, &call_list);
	}
	}

	const unsigned block_words = BITSET_WORDS(impl->num_blocks);
	BITSET_WORD *def_blocks = ralloc_array(mem_ctx, BITSET_WORD, block_words);

	list_for_each_entry(struct call_liveness_entry, entry, &call_list, list) {
	unsigned i;

	nir_builder b = nir_builder_at(nir_after_instr(&entry->instr->instr));

	struct nir_phi_builder *builder = nir_phi_builder_create(impl);
	struct hash_table *phi_value_table =
	_mesa_pointer_hash_table_create(mem_ctx);
	struct hash_table *remap_table =
	_mesa_pointer_hash_table_create(mem_ctx);

	BITSET_FOREACH_SET(i, entry->live_set, num_defs) {
	if (!rematerializable[i] \|\|
	_mesa_hash_table_search(remap_table, rematerializable[i]))
	continue;

	assert(!_mesa_hash_table_search(phi_value_table, rematerializable[i]));

	struct remat_chain_check_data check_data = {
	.remap_table = remap_table,
	.chain_length = 1,
	};

	if (!nir_foreach_src(rematerializable[i]->parent_instr,
	can_remat_chain, &check_data))
	continue;

	struct remat_chain_data remat_data = {
	.b = &b,
	.remap_table = remap_table,
	.phi_value_table = phi_value_table,
	.phi_builder = builder,
	.def_blocks = def_blocks,
	};

	nir_foreach_src(rematerializable[i]->parent_instr, do_remat_chain,
	&remat_data);

	remat_ssa_def(&b, rematerializable[i], remap_table, phi_value_table,
	builder, def_blocks);
	progress = true;
	}
	_mesa_hash_table_destroy(remap_table, NULL);

	nir_foreach_block(block, impl) {
	nir_foreach_instr(instr, block) {
	if (instr->type == nir_instr_type_phi)
	continue;

	nir_foreach_src(instr, rewrite_instr_src_from_phi_builder,
	phi_value_table);
	}
	}

	nir_phi_builder_finish(builder);
	_mesa_hash_table_destroy(phi_value_table, NULL);
	}

	ralloc_free(mem_ctx);

	nir_progress(true, impl, nir_metadata_block_index \| nir_metadata_dominance);
	return progress;
	}

	/* Tries to rematerialize as many live vars as possible after calls.
	* Note: nir_opt_cse will undo any rematerializations done by this pass,
	* so it shouldn't be run afterward.
	*/
	bool
	nir_minimize_call_live_states(nir_shader *shader)
	{
	bool progress = false;

	nir_foreach_function_impl(impl, shader) {
	progress \|= nir_minimize_call_live_states_impl(impl);
	}

	return progress;
	}