Allow for dynamic indexing of texture arrays.

[Makogan]

I am not entirely sure whether to flag this as a bug or as missing functionality. Consider this HLSL example which I have tested by compiling it to spir-v and running it:

float4 main(
    [[vk::location(0)]] float3 normal : NORMAL0,
    [[vk::location(1)]] float3 uv : UV0,
    [[vk::location(2)]] float3 position : POSITION0,
    [[vk::location(3)]] float3 camera_position : CAMPOS0
) : SV_TARGET
{
    uint index = int(round(uv.z));
    float4 albedo = textures[index].Sample(textures_sampler[index], uv.xy);

    return (Lambert(
        position,
        normal,
        albedo.xyz,
        vec3(5, 1000, 5),
        normalize(vec3(1, 1, 1))) * 2.0 + 0.1);
}

I was trying to port this code onto rust and compiling it through rustgpu, as in this example:

#![cfg_attr(target_arch = "spirv", no_std, feature(lang_items))]
#![allow(internal_features)]

extern crate bytemuck;
extern crate spirv_std;

pub use spirv_std::glam::*;
use spirv_std::image::*;
use spirv_std::num_traits::Float;
use spirv_std::spirv;

type Texture2D = Image!(2D, type=f32, sampled);

// Fragment shader
#[spirv(fragment)]
pub fn main_fs(
    in_normal: Vec3,
    in_uv: Vec3,
    in_position: Vec3,
    in_camera_position: Vec3,
    #[spirv(descriptor_set = 1, binding = 1)] image: &[SampledImage<Texture2D>; 69],
    output: &mut Vec4,
)
{
    let index = in_uv.z.round() as usize;
    let albedo = image[index].sample(in_uv.xy());

    *output = albedo;
}

// Vertex shader
pub struct CamUBO
{
    model: Mat4,
    view: Mat4,
    proj: Mat4,
}

#[spirv(vertex)]
pub fn main_vs(
    in_position: Vec3,
    in_normal: Vec3,
    in_uv: Vec3,
    #[spirv(uniform, descriptor_set = 1, binding = 0)] camera_ubo: &CamUBO,

    #[spirv(position, invariant)] screen_pos: &mut Vec4,
    out_normal: &mut Vec3,
    out_uv: &mut Vec3,
    out_position: &mut Vec3,
    out_camera_position: &mut Vec3,
)
{
    *screen_pos = camera_ubo.proj
        * camera_ubo.view
        * camera_ubo.model
        * Vec4::new(in_position.x, in_position.y, in_position.z, 1.0);

    *out_position = screen_pos.xyz();
    *out_normal = (camera_ubo.model * in_normal.extend(0.0)).xyz();
    *out_uv = in_uv;
    *out_camera_position = -camera_ubo.view.col(3).xyz();
}

This creates a long list of errors, but the most pertinent I think is this:

error: cannot offset a pointer to an arbitrary element

And indeed, setting the index to a constant does compile and run fine. Given that the HLSL code is able to compile down into spirv, it should be possible to generate the correct assembly to index into the texture array. I am unsure if there might be an existing work around however.

Example & Steps To Reproduce

Run spirv-builder and compile the above snippet.

• Rustc version: rustc 1.77.0-nightly (75c68cfd2 2024-01-07)
• OS: Ubuntu 22.04
• GPU: GTX 1070
• SPIR-V: SPIRV-Tools v2024.1 unknown hash, 2024-03-12T15:23:26+00:00

[Makogan]

I am not well versed in the internals of rustgpu nor with spirv. But perhaps it would be possible to add an attribute macro to guide the compiler to compile slices of image types into SampledImageArrayDynamicIndexing, it seems that spirv builder already defines enumerator types that match this case. So i assume it;s just a matter of specifying how and when to turn certain slices into the corresponding dynamic indexing type?

https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_capability

[Patryk27]

This feels like RuntimeArray, i.e. &RuntimeArray<SampledImage<Texture2D>>.

[Makogan]

This feels like RuntimeArray, i.e. &RuntimeArray<SampledImage<Texture2D>>.

This is the correct answer, thank you and sorry I missed it.

A different alternative if one uses the old way of doing things, with fixed array sizes is to use unchecked indexing, like this:

// Fragment shader
#[spirv(fragment)]
pub fn main_fs(
    in_normal: Vec3,
    in_uv: Vec3,
    in_position: Vec3,
    #[spirv(descriptor_set = 1, binding = 1)] image: &[SampledTexture2D; 69],
    // use this for boundless arrays
    // #[spirv(descriptor_set = 1, binding = 1)] image: &RuntimeArray<SampledTexture2D>,
    output: &mut Vec4,
)
{
    // let albedo = unsafe { image.index(index).sample(in_uv.xy()) };
    let index = in_uv.z.round() as usize;
    let albedo = unsafe { image.index_unchecked(index).sample(in_uv.xy()) };

    *output = lambert(
        in_position,
        in_normal,
        albedo.xyz(),
        Vec3::new(5., 1000., 5.),
        Vec3::new(1., 1., 1.).normalize() * 1.5 + Vec3::splat(0.5),
    );
}