[Enhancement] Suport for image2DArray (GLSL) / Texture2DArray (Godot)

[jotaf98]

Hi, this is very useful! I was wondering if there was a way to also access multiple images in a batch, as image2DArray. I think that this maps to Godot's Texture2DArray, which we can build easily.

Justification: For parallel computation, which is the main use case of compute shaders, it would make sense to allow the "embarrassingly parallel" case of multiple images.

[jotaf98]

Here is the class implementing image2Darray, in case anyone is interested:

extends Uniform
class_name ImageArrayUniform
## [Uniform] corresponding to a texture array. Given to the shader as an image2DArray.

var texture: RID ## The [RID] of the corresponding texture array. Used internally.
var texture_size: Vector3i ## The resolution of the texture array (width, height, layers).
var image_format: Image.Format ## The [enum Image.Format] of the texture array.
var texture_format: RDTextureFormat ## The [RDTextureFormat] of the texture array.

signal async_image_retrieved(images: Array[Image])

## Returns a new ImageArrayUniform object using the given array of [param images].
static func create(images: Array[Image]) -> ImageArrayUniform:
	if images.is_empty():
		push_error("ImageArrayUniform.create: images array cannot be empty")
		return null
	
	var uniform := ImageArrayUniform.new()
	var first_image := images[0]
	var layer_size := first_image.get_size()
	
	# validate that all images have the same size and format
	uniform.image_format = first_image.get_format()
	for image in images:
		if image.get_size() != layer_size or image.get_format() != uniform.image_format:
			push_error("ImageArrayUniform.create: all images must have the same size and format")
			return null
	
	uniform.texture_size = Vector3i(layer_size.x, layer_size.y, images.size())
	
	# create texture format for 2D array
	uniform.texture_format = RDTextureFormat.new()
	uniform.texture_format.texture_type = RenderingDevice.TEXTURE_TYPE_2D_ARRAY
	uniform.texture_format.format = ImageFormatHelper.convert_image_format_to_data_format(uniform.image_format)
	uniform.texture_format.width = uniform.texture_size.x
	uniform.texture_format.height = uniform.texture_size.y
	uniform.texture_format.array_layers = uniform.texture_size.z
	uniform.texture_format.usage_bits = RenderingDevice.TEXTURE_USAGE_STORAGE_BIT | \
		RenderingDevice.TEXTURE_USAGE_SAMPLING_BIT | \
		RenderingDevice.TEXTURE_USAGE_CAN_UPDATE_BIT | \
		RenderingDevice.TEXTURE_USAGE_CAN_COPY_FROM_BIT
	
	# collect image data for all layers
	var image_data: Array[PackedByteArray] = []
	for image in images:
		image_data.append(image.get_data())
	
	uniform.texture = ComputeHelper.rd.texture_create(uniform.texture_format, ComputeHelper.view, image_data)
	return uniform

## ImageArrayUniform's custom implementation of [method Uniform.get_rd_uniform].
func get_rd_uniform(binding: int) -> RDUniform:
	var uniform := RDUniform.new()
	uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_IMAGE
	uniform.binding = binding
	uniform.add_id(texture)
	return uniform

## Updates the texture array to match [param images].
func update_images(images: Array[Image]) -> void:
	if images.is_empty():
		push_error("ImageArrayUniform.update_images: images array cannot be empty")
		return
	
	var first_image := images[0]
	var layer_size := first_image.get_size()
	var new_format := first_image.get_format()
	
	# validate that all images have the same size and format
	for image in images:
		if image.get_size() != layer_size or image.get_format() != new_format:
			push_error("ImageArrayUniform.update_images: all images must have the same size and format")
			return
	
	var new_size := Vector3i(layer_size.x, layer_size.y, images.size())
	
	# check if we can update in place or need to recreate
	if texture_size == new_size and image_format == new_format:
		# update each layer in place
		for layer in range(images.size()):
			ComputeHelper.rd.texture_update(texture, layer, images[layer].get_data())
	else:
		# recreate texture with new dimensions/format
		ComputeHelper.rd.free_rid(texture)
		image_format = new_format
		texture_size = new_size
		
		texture_format = RDTextureFormat.new()
		texture_format.texture_type = RenderingDevice.TEXTURE_TYPE_2D_ARRAY
		texture_format.format = ImageFormatHelper.convert_image_format_to_data_format(image_format)
		texture_format.width = texture_size.x
		texture_format.height = texture_size.y
		texture_format.array_layers = texture_size.z
		texture_format.usage_bits = RenderingDevice.TEXTURE_USAGE_STORAGE_BIT | \
			RenderingDevice.TEXTURE_USAGE_SAMPLING_BIT | \
			RenderingDevice.TEXTURE_USAGE_CAN_UPDATE_BIT | \
			RenderingDevice.TEXTURE_USAGE_CAN_COPY_FROM_BIT
		
		var image_data: Array[PackedByteArray] = []
		for image in images:
			image_data.append(image.get_data())
		
		texture = ComputeHelper.rd.texture_create(texture_format, ComputeHelper.view, image_data)
		rid_updated.emit(self)

## Updates a single layer of the texture array.
func update_layer(layer_index: int, image: Image) -> void:
	if layer_index < 0 or layer_index >= texture_size.z:
		push_error("ImageArrayUniform.update_layer: layer_index out of bounds")
		return
	
	if image.get_size() != Vector2i(texture_size.x, texture_size.y) or image.get_format() != image_format:
		push_error("ImageArrayUniform.update_layer: image size or format does not match texture array")
		return
	
	ComputeHelper.rd.texture_update(texture, layer_index, image.get_data())

## Returns an array of [Image] objects, one for each layer in the texture array. [b]Warning:[/b] Getting data from the GPU is very slow.
func get_images() -> Array[Image]:
	var images: Array[Image] = []
	for layer in range(texture_size.z):
		var image_data := ComputeHelper.rd.texture_get_data(texture, layer)
		var image := Image.create_from_data(texture_size.x, texture_size.y, false, image_format, image_data)
		images.append(image)
	return images

## Gets a single layer as an [Image]. [b]Warning:[/b] Getting data from the GPU is very slow.
func get_layer(layer_index: int) -> Image:
	if layer_index < 0 or layer_index >= texture_size.z:
		push_error("ImageArrayUniform.get_layer: layer_index out of bounds")
		return null
	
	var image_data := ComputeHelper.rd.texture_get_data(texture, layer_index)
	return Image.create_from_data(texture_size.x, texture_size.y, false, image_format, image_data)

## Gets all layers asynchronously. Returns a [Signal] with an Array[Image] that will be emitted when the images are retrieved. The signal remains the same each time you call this function, so feel free to cache it. [b]Note:[/b] The delay to when the signal is emitted corresponds to the amount of frames specified by [member ProjectSettings.rendering/rendering_device/vsync/frame_queue_size]. Also, this function does nothing before Godot 4.4.
func get_images_async() -> Signal:
	if ComputeHelper.version < 4:
		return async_image_retrieved
	
	# retrieve all layers asynchronously
	for layer in range(texture_size.z):
		ComputeHelper.rd.texture_get_data_async(texture, layer, async_image_callback.bind(layer))
	
	return async_image_retrieved

var _async_images: Array[Image] = []
var _async_layers_remaining: int = 0

## Used internally for [get_images_async].
func async_image_callback(image_data: PackedByteArray, layer_index: int) -> void:
	var image := Image.create_from_data(texture_size.x, texture_size.y, false, image_format, image_data)
	
	# initialize array on first callback
	if _async_layers_remaining == 0:
		_async_images.clear()
		_async_images.resize(texture_size.z)
		_async_layers_remaining = texture_size.z
	
	_async_images[layer_index] = image
	_async_layers_remaining -= 1
	
	# emit signal when all layers are retrieved
	if _async_layers_remaining == 0:
		async_image_retrieved.emit(_async_images)

func _notification(what: int) -> void:
	if what == NOTIFICATION_PREDELETE:
		ComputeHelper.rd.free_rid(texture)