ac/nir: load buffer descriptors via ac_shader_abi::load_ssbo

2 files changed, 30 insertions, 8 deletions

diff --git a/src/amd/common/ac_nir_to_llvm.c b/src/amd/common/ac_nir_to_llvm.c
index 329fd4e3f5..6382a9d54c 100644
--- a/src/amd/common/ac_nir_to_llvm.c
+++ b/src/amd/common/ac_nir_to_llvm.c
@@ -2265,10 +2265,8 @@ static void visit_store_ssbo(struct nir_to_llvm_context *ctx,
 	LLVMValueRef base_data, base_offset;
 	LLVMValueRef params[6];
 
-	if (ctx->stage == MESA_SHADER_FRAGMENT)
-		ctx->shader_info->fs.writes_memory = true;
-
-	params[1] = get_src(ctx->nir, instr->src[1]);
+	params[1] = ctx->abi.load_ssbo(&ctx->abi,
+				       get_src(ctx->nir, instr->src[1]), true);
 	params[2] = LLVMConstInt(ctx->i32, 0, false); /* vindex */
 	params[4] = ctx->i1false;  /* glc */
 	params[5] = ctx->i1false;  /* slc */
@@ -2345,14 +2343,14 @@ static LLVMValueRef visit_atomic_ssbo(struct nir_to_llvm_context *ctx,
 	const char *name;
 	LLVMValueRef params[6];
 	int arg_count = 0;
-	if (ctx->stage == MESA_SHADER_FRAGMENT)
-		ctx->shader_info->fs.writes_memory = true;
 
 	if (instr->intrinsic == nir_intrinsic_ssbo_atomic_comp_swap) {
 		params[arg_count++] = llvm_extract_elem(&ctx->ac, get_src(ctx->nir, instr->src[3]), 0);
 	}
 	params[arg_count++] = llvm_extract_elem(&ctx->ac, get_src(ctx->nir, instr->src[2]), 0);
-	params[arg_count++] = get_src(ctx->nir, instr->src[0]);
+	params[arg_count++] = ctx->abi.load_ssbo(&ctx->abi,
+						 get_src(ctx->nir, instr->src[0]),
+						 true);
 	params[arg_count++] = LLVMConstInt(ctx->i32, 0, false); /* vindex */
 	params[arg_count++] = get_src(ctx->nir, instr->src[1]);      /* voffset */
 	params[arg_count++] = ctx->i1false;  /* slc */
@@ -2426,7 +2424,9 @@ static LLVMValueRef visit_load_buffer(struct nir_to_llvm_context *ctx,
 			unreachable("unhandled number of components");
 
 		LLVMValueRef params[] = {
-			get_src(ctx->nir, instr->src[0]),
+			ctx->abi.load_ssbo(&ctx->abi,
+					   get_src(ctx->nir, instr->src[0]),
+					   false),
 			LLVMConstInt(ctx->i32, 0, false),
 			offset,
 			ctx->i1false,
@@ -4112,6 +4112,17 @@ static void visit_intrinsic(struct ac_nir_context *ctx,
 	}
 }
 
+static LLVMValueRef radv_load_ssbo(struct ac_shader_abi *abi,
+				   LLVMValueRef buffer, bool write)
+{
+	struct nir_to_llvm_context *ctx = nir_to_llvm_context_from_abi(abi);
+
+	if (write && ctx->stage == MESA_SHADER_FRAGMENT)
+		ctx->shader_info->fs.writes_memory = true;
+
+	return buffer;
+}
+
 static LLVMValueRef radv_get_sampler_desc(struct ac_shader_abi *abi,
 					  unsigned descriptor_set,
 					  unsigned base_index,
@@ -6172,6 +6183,7 @@ LLVMModuleRef ac_translate_nir_to_llvm(LLVMTargetMachineRef tm,
 	ctx.abi.chip_class = options->chip_class;
 	ctx.abi.inputs = &ctx.inputs[0];
 	ctx.abi.emit_outputs = handle_shader_outputs_post;
+	ctx.abi.load_ssbo = radv_load_ssbo;
 	ctx.abi.load_sampler_desc = radv_get_sampler_desc;
 
 	nir_foreach_variable(variable, &nir->outputs)
diff --git a/src/amd/common/ac_shader_abi.h b/src/amd/common/ac_shader_abi.h
index fe804cf5c8..efbd7de460 100644
--- a/src/amd/common/ac_shader_abi.h
+++ b/src/amd/common/ac_shader_abi.h
@@ -59,6 +59,16 @@ struct ac_shader_abi {
 	LLVMValueRef (*load_ubo)(struct ac_shader_abi *abi, LLVMValueRef index);
 
 	/**
+	 * Load the descriptor for the given buffer.
+	 *
+	 * \param buffer the buffer as presented in NIR: this is the descriptor
+	 *               in Vulkan, and the buffer index in OpenGL/Gallium
+	 * \param write whether buffer contents will be written
+	 */
+	LLVMValueRef (*load_ssbo)(struct ac_shader_abi *abi,
+				  LLVMValueRef buffer, bool write);
+
+	/**
 	 * Load a descriptor associated to a sampler.
 	 *
 	 * \param descriptor_set the descriptor set index (only for Vulkan)