From 12f93a73779ea61c21d3433dcdccdebe437beb01 Mon Sep 17 00:00:00 2001 From: Mike Blumenkrantz Date: Sat, 5 Dec 2020 09:41:26 -0500 Subject: [PATCH] zink: hook up cs push constant for nir_intrinsic_load_work_dim Reviewed-by: Dave Airlie Part-of: --- src/gallium/drivers/zink/zink_compiler.c | 60 +++++++++++++++++++++++- src/gallium/drivers/zink/zink_draw.c | 5 ++ src/gallium/drivers/zink/zink_program.c | 9 ++++ src/gallium/drivers/zink/zink_program.h | 4 ++ 4 files changed, 77 insertions(+), 1 deletion(-) diff --git a/src/gallium/drivers/zink/zink_compiler.c b/src/gallium/drivers/zink/zink_compiler.c index 1655f7ebfc3..63b05138402 100644 --- a/src/gallium/drivers/zink/zink_compiler.c +++ b/src/gallium/drivers/zink/zink_compiler.c @@ -55,6 +55,26 @@ create_vs_pushconst(nir_shader *nir) vs_pushconst->data.location = INT_MAX; //doesn't really matter } +static void +create_cs_pushconst(nir_shader *nir) +{ + nir_variable *cs_pushconst; + /* create compatible layout for the ntv push constant loader */ + struct glsl_struct_field *fields = rzalloc_size(nir, 1 * sizeof(struct glsl_struct_field)); + fields[0].type = glsl_array_type(glsl_uint_type(), 1, 0); + fields[0].name = ralloc_asprintf(nir, "work_dim"); + fields[0].offset = 0; + cs_pushconst = nir_variable_create(nir, nir_var_mem_push_const, + glsl_struct_type(fields, 1, "struct", false), "cs_pushconst"); + cs_pushconst->data.location = INT_MAX; //doesn't really matter +} + +static bool +reads_work_dim(nir_shader *shader) +{ + return BITSET_TEST(shader->info.system_values_read, SYSTEM_VALUE_WORK_DIM); +} + static bool lower_discard_if_instr(nir_intrinsic_instr *instr, nir_builder *b) { @@ -139,6 +159,42 @@ lower_discard_if(nir_shader *shader) return progress; } +static bool +lower_work_dim_instr(nir_builder *b, nir_instr *in, void *data) +{ + if (in->type != nir_instr_type_intrinsic) + return false; + nir_intrinsic_instr *instr = nir_instr_as_intrinsic(in); + if (instr->intrinsic != nir_intrinsic_load_work_dim) + return false; + + if (instr->intrinsic == nir_intrinsic_load_work_dim) { + b->cursor = nir_after_instr(&instr->instr); + nir_intrinsic_instr *load = nir_intrinsic_instr_create(b->shader, nir_intrinsic_load_push_constant); + load->src[0] = nir_src_for_ssa(nir_imm_int(b, 0)); + nir_intrinsic_set_range(load, 3 * sizeof(uint32_t)); + load->num_components = 1; + nir_ssa_dest_init(&load->instr, &load->dest, 1, 32, "work_dim"); + nir_builder_instr_insert(b, &load->instr); + + nir_ssa_def_rewrite_uses(&instr->dest.ssa, &load->dest.ssa); + } + + return true; +} + +static bool +lower_work_dim(nir_shader *shader) +{ + if (shader->info.stage != MESA_SHADER_KERNEL) + return false; + + if (!reads_work_dim(shader)) + return false; + + return nir_shader_instructions_pass(shader, lower_work_dim_instr, nir_metadata_dominance, NULL); +} + static bool lower_64bit_vertex_attribs_instr(nir_builder *b, nir_instr *instr, void *data) { @@ -750,7 +806,8 @@ zink_shader_create(struct zink_screen *screen, struct nir_shader *nir, nir->info.stage == MESA_SHADER_TESS_EVAL) { NIR_PASS_V(nir, nir_lower_indirect_derefs, nir_var_shader_in | nir_var_shader_out, UINT_MAX); NIR_PASS_V(nir, nir_lower_io_arrays_to_elements_no_indirects, false); - } + } else if (nir->info.stage == MESA_SHADER_KERNEL) + create_cs_pushconst(nir); NIR_PASS_V(nir, nir_lower_uniforms_to_ubo, 16); if (nir->info.stage < MESA_SHADER_FRAGMENT) @@ -758,6 +815,7 @@ zink_shader_create(struct zink_screen *screen, struct nir_shader *nir, if (nir->info.stage == MESA_SHADER_GEOMETRY) NIR_PASS_V(nir, nir_lower_gs_intrinsics, nir_lower_gs_intrinsics_per_stream); NIR_PASS_V(nir, lower_basevertex); + NIR_PASS_V(nir, lower_work_dim); NIR_PASS_V(nir, nir_lower_regs_to_ssa); NIR_PASS_V(nir, lower_baseinstance); optimize_nir(nir); diff --git a/src/gallium/drivers/zink/zink_draw.c b/src/gallium/drivers/zink/zink_draw.c index c126d656c29..a7c9507e5ec 100644 --- a/src/gallium/drivers/zink/zink_draw.c +++ b/src/gallium/drivers/zink/zink_draw.c @@ -634,6 +634,11 @@ zink_launch_grid(struct pipe_context *pctx, const struct pipe_grid_info *info) vkCmdBindPipeline(batch->state->cmdbuf, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline); + if (BITSET_TEST(comp_program->shader->nir->info.system_values_read, SYSTEM_VALUE_WORK_DIM)) + vkCmdPushConstants(batch->state->cmdbuf, comp_program->base.layout, VK_SHADER_STAGE_COMPUTE_BIT, + offsetof(struct zink_cs_push_constant, work_dim), sizeof(uint32_t), + &info->work_dim); + if (info->indirect) { vkCmdDispatchIndirect(batch->state->cmdbuf, zink_resource(info->indirect)->obj->buffer, info->indirect_offset); zink_batch_reference_resource_rw(batch, zink_resource(info->indirect), false); diff --git a/src/gallium/drivers/zink/zink_program.c b/src/gallium/drivers/zink/zink_program.c index 36727c2d582..527d575de76 100644 --- a/src/gallium/drivers/zink/zink_program.c +++ b/src/gallium/drivers/zink/zink_program.c @@ -177,6 +177,15 @@ create_compute_pipeline_layout(VkDevice dev, struct zink_compute_program *comp) plci.pSetLayouts = layouts; plci.setLayoutCount = num_layouts; + VkPushConstantRange pcr = {}; + if (comp->shader->nir->info.stage == MESA_SHADER_KERNEL) { + pcr.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT; + pcr.offset = 0; + pcr.size = sizeof(struct zink_cs_push_constant); + plci.pushConstantRangeCount = 1; + plci.pPushConstantRanges = &pcr; + } + VkPipelineLayout layout; if (vkCreatePipelineLayout(dev, &plci, NULL, &layout) != VK_SUCCESS) { debug_printf("vkCreatePipelineLayout failed!\n"); diff --git a/src/gallium/drivers/zink/zink_program.h b/src/gallium/drivers/zink/zink_program.h index f0b7a39a357..168efbc424a 100644 --- a/src/gallium/drivers/zink/zink_program.h +++ b/src/gallium/drivers/zink/zink_program.h @@ -52,6 +52,10 @@ struct zink_gfx_push_constant { float default_outer_level[4]; }; +struct zink_cs_push_constant { + unsigned work_dim; +}; + /* a shader module is used for directly reusing a shader module between programs, * e.g., in the case where we're swapping out only one shader, * allowing us to skip going through shader keys