How to modify a texture inside a compute shader

Archive
1
:information_source: Attention Topic was automatically imported from the old Question2Answer platform.
:bust_in_silhouette: Asked By creative-mind

Hi,

I have a compute shader that takes a texture with sampler as input and outputs a float buffer array (see code bellow). However, I would like to modify the texture inside the compute shader and get it as the output instead of the float buffer array. Is there a way how to read the texture afterwards, on the CPU side?

Here’s the CPU side:

@tool
extends Node
 
@export var generate := false:
	set(_value):
		generate = false
		generate_normals(load("res://texture.png"))


func _ready():
	generate_normals(load("res://texture.png"))
	
func generate_normals(texture: Texture2D):
	# Load GLSL shader
	var shader_file := load("res://minimal.glsl")
	var shader_spirv: RDShaderSPIRV = shader_file.get_spirv()
	
	# Create a local rendering device.
	var rd := RenderingServer.create_local_rendering_device()
	var shader := rd.shader_create_from_spirv(shader_spirv)
	
	# Create input texture uniform
	var img = texture.get_image()
	var img_pba = img.get_data()
	var width = texture.get_width()
	var height = texture.get_height()
	
	var fmt = RDTextureFormat.new()
	fmt.width = width
	fmt.height = height
	fmt.usage_bits = RenderingDevice.TEXTURE_USAGE_CAN_UPDATE_BIT | RenderingDevice.TEXTURE_USAGE_SAMPLING_BIT
	fmt.format = RenderingDevice.DATA_FORMAT_R8G8B8A8_SRGB
	
	var read_data = PackedByteArray(img_pba)
	
	var v_tex = rd.texture_create(fmt, RDTextureView.new(), [img_pba])
	var samp_state = RDSamplerState.new()
	samp_state.unnormalized_uvw = true
	var samp = rd.sampler_create(samp_state)
	
	var tex_uniform = RDUniform.new()
	tex_uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_SAMPLER_WITH_TEXTURE
	tex_uniform.binding = 0
	tex_uniform.add_id(samp)
	tex_uniform.add_id(v_tex)
	
	# Create storage buffer for the input array.

	# Initialise the byte array the shader will write to.
	var write_data = PackedByteArray()
	write_data.resize(read_data.size())

	# Create storage buffer for the output array.
	var write_buffer = rd.storage_buffer_create(write_data.size(), write_data)
	var write_uniform := RDUniform.new()
	write_uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_STORAGE_BUFFER
	write_uniform.binding = 1
	write_uniform.add_id(write_buffer)

	# Create buffer for sending grid size data to array.
	var size_data_bytes := PackedByteArray(PackedInt32Array([width, height]).to_byte_array())
	var size_data_buffer = rd.storage_buffer_create(8, size_data_bytes)
	var size_data_uniform := RDUniform.new()
	size_data_uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_STORAGE_BUFFER
	size_data_uniform.binding = 2
	size_data_uniform.add_id(size_data_buffer)

	# Create a set for the uniforms.
	var uniform_set = rd.uniform_set_create([tex_uniform, write_uniform, size_data_uniform], shader, 0)

	# Create a compute pipeline
	var pipeline := rd.compute_pipeline_create(shader)
	var compute_list := rd.compute_list_begin()
	rd.compute_list_bind_compute_pipeline(compute_list, pipeline)
	rd.compute_list_bind_uniform_set(compute_list, uniform_set, 0)
	rd.compute_list_dispatch(compute_list, width / 8, height / 8, 1)
	rd.compute_list_end()
	
	# Submit to GPU and wait for sync
	rd.submit()
	rd.sync()
	
	# Read back the data from the buffers
	var output_bytes := rd.buffer_get_data(write_buffer)
	var output := output_bytes.to_float32_array()
	print(output)

And here’s the shader code:

#[compute]
#version 450

// Invocations in the (x, y, z) dimension
layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;

// Input Heightmap texture
layout(set = 0, binding = 0) uniform sampler2D heightmap;

// Output Normal texture
layout(set = 0, binding = 1, std430) restrict buffer FloatBufferOut {
  float data[];
}
buffer_out;

// Texture size info
layout(set = 0, binding = 2, std430) restrict buffer SizeDataBuffer {
  int width;
  int height;
}
size_data;


int mod(int x, int y) {
	return (x - y * (x/y));
}

void main() {
	int idx = int(gl_GlobalInvocationID.y * size_data.width + gl_GlobalInvocationID.x);
	vec2 pos = vec2(mod(int(gl_GlobalInvocationID.x), size_data.width), 
					mod(int(gl_GlobalInvocationID.y), size_data.height));
	
	vec4 color = texture(heightmap, pos);
    buffer_out.data[idx] = float(color[0]);
}