radeonsi: make si_declare_compute_memory() more generic and call for nir

Reviewed-by: Marek Olšák <marek.olsak@amd.com>

src/gallium/drivers/radeonsi/si_shader.c

@@ -2231,16 +2231,13 @@ void si_load_system_value(struct si_shader_context *ctx,
 	ctx->system_values[index] = value;
 }
 
-void si_declare_compute_memory(struct si_shader_context *ctx,
-			       const struct tgsi_full_declaration *decl)
+void si_declare_compute_memory(struct si_shader_context *ctx)
 {
 	struct si_shader_selector *sel = ctx->shader->selector;
 
 	LLVMTypeRef i8p = LLVMPointerType(ctx->i8, AC_LOCAL_ADDR_SPACE);
 	LLVMValueRef var;
 
-	assert(decl->Declaration.MemType == TGSI_MEMORY_TYPE_SHARED);
-	assert(decl->Range.First == decl->Range.Last);
 	assert(!ctx->ac.lds);
 
 	var = LLVMAddGlobalInAddressSpace(ctx->ac.module,
@@ -2252,6 +2249,15 @@ void si_declare_compute_memory(struct si_shader_context *ctx,
 	ctx->ac.lds = LLVMBuildBitCast(ctx->ac.builder, var, i8p, "");
 }
 
+void si_tgsi_declare_compute_memory(struct si_shader_context *ctx,
+				    const struct tgsi_full_declaration *decl)
+{
+	assert(decl->Declaration.MemType == TGSI_MEMORY_TYPE_SHARED);
+	assert(decl->Range.First == decl->Range.Last);
+
+	si_declare_compute_memory(ctx);
+}
+
 static LLVMValueRef load_const_buffer_desc(struct si_shader_context *ctx, int i)
 {
 	LLVMValueRef list_ptr = LLVMGetParam(ctx->main_fn,

src/gallium/drivers/radeonsi/si_shader_internal.h

@@ -320,7 +320,8 @@ LLVMValueRef si_load_image_desc(struct si_shader_context *ctx,
 void si_load_system_value(struct si_shader_context *ctx,
 			  unsigned index,
 			  const struct tgsi_full_declaration *decl);
-void si_declare_compute_memory(struct si_shader_context *ctx,
+void si_declare_compute_memory(struct si_shader_context *ctx);
+void si_tgsi_declare_compute_memory(struct si_shader_context *ctx,
 				    const struct tgsi_full_declaration *decl);
 
 void si_llvm_load_input_vs(

src/gallium/drivers/radeonsi/si_shader_nir.c

@@ -857,6 +857,10 @@ bool si_nir_build_llvm(struct si_shader_context *ctx, struct nir_shader *nir)
 	ctx->num_samplers = util_last_bit(info->samplers_declared);
 	ctx->num_images = util_last_bit(info->images_declared);
 
+	if (ctx->shader->selector->local_size) {
+		assert(nir->info.stage == MESA_SHADER_COMPUTE);
+		si_declare_compute_memory(ctx);
+	}
 	ac_nir_translate(&ctx->ac, &ctx->abi, nir, NULL);
 
 	return true;

src/gallium/drivers/radeonsi/si_shader_tgsi_setup.c

@@ -851,7 +851,7 @@ static void emit_declaration(struct lp_build_tgsi_context *bld_base,
 	}
 
 	case TGSI_FILE_MEMORY:
-		si_declare_compute_memory(ctx, decl);
+		si_tgsi_declare_compute_memory(ctx, decl);
 		break;
 
 	default: