Mali SPIR-V driver crash: textureLoad via fn param + wgpu Restrict bounds check

Cargo.toml

[package]
name = "mali-repro"
version = "0.1.0"
edition = "2021"

[[bin]]
name = "mali-repro"
path = "main.rs"

[dependencies]
wgpu = { version = "30", default-features = false, features = ["wgsl", "vulkan"] }
futures = "0.3"

main.rs

// Repro for a Mali driver SIGSEGV triggered by wgpu's default (Restrict) image bounds-check
// policy. Builds a render pipeline whose fragment shader calls textureLoad() on a texture
// received as a function parameter.
//
// - default shader module: naga emits OpImageQueryLevels for the bounds check, which segfaults
//   the Mali r54p3 SPIR-V compiler when the image reaches it through an OpFunctionParameter.
// - `--unchecked`: same WGSL via create_shader_module_trusted with bounds checks disabled;
//   no OpImageQueryLevels is emitted, pipeline creation succeeds.
//
// Run on the device with the affected driver, e.g. over `adb shell` on Android:
//   cargo build --release --target aarch64-linux-android
//   adb push target/aarch64-linux-android/release/mali-repro /data/local/tmp/
//   adb shell /data/local/tmp/mali-repro              # SIGSEGV in libGLES_mali.so
//   adb shell /data/local/tmp/mali-repro --unchecked   # ok

const SHADER: &str = r#"
@group(0) @binding(0) var t: texture_2d<f32>;

fn load(tex: texture_2d<f32>) -> vec4<f32> {
    return textureLoad(tex, vec2<i32>(0, 0), 0);
}

@vertex
fn vs_main(@builtin(vertex_index) i: u32) -> @builtin(position) vec4<f32> {
    let x = f32(i32(i & 1u) * 4 - 1);
    let y = f32(i32(i >> 1u) * 4 - 1);
    return vec4<f32>(x, y, 0.0, 1.0);
}

@fragment
fn fs_main() -> @location(0) vec4<f32> {
    return load(t);
}
"#;

fn main() {
    let unchecked = std::env::args().nth(1).as_deref() == Some("--unchecked");

    let mut desc = wgpu::InstanceDescriptor::new_without_display_handle_from_env();
    desc.backends = wgpu::Backends::VULKAN; // the driver bug is Vulkan-only; GL is unaffected
    let instance = wgpu::Instance::new(desc);
    let adapter = futures::executor::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
        power_preference: wgpu::PowerPreference::HighPerformance,
        force_fallback_adapter: false,
        compatible_surface: None,
        apply_limit_buckets: false,
    }))
    .expect("no adapter");
    let info = adapter.get_info();
    println!("adapter: {} {:?} driver {} {}", info.name, info.backend, info.driver, info.driver_info);

    let (device, _queue) = futures::executor::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
        label: None,
        required_features: wgpu::Features::empty(),
        required_limits: wgpu::Limits::downlevel_webgl2_defaults().using_resolution(adapter.limits()),
        memory_hints: Default::default(),
        trace: Default::default(),
        experimental_features: Default::default(),
    }))
    .expect("no device");

    let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        label: None,
        entries: &[wgpu::BindGroupLayoutEntry {
            binding: 0,
            visibility: wgpu::ShaderStages::FRAGMENT,
            ty: wgpu::BindingType::Texture {
                sample_type: wgpu::TextureSampleType::Float { filterable: true },
                view_dimension: wgpu::TextureViewDimension::D2,
                multisampled: false,
            },
            count: None,
        }],
    });
    let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: None,
        bind_group_layouts: &[Some(&bgl)],
        immediate_size: 0,
    });

    let source = wgpu::ShaderSource::Wgsl(SHADER.into());
    let module = if unchecked {
        // SAFETY: the shader indexes in bounds; this only disables naga's inserted checks.
        unsafe {
            device.create_shader_module_trusted(
                wgpu::ShaderModuleDescriptor { label: None, source },
                wgpu::ShaderRuntimeChecks::unchecked(),
            )
        }
    } else {
        device.create_shader_module(wgpu::ShaderModuleDescriptor { label: None, source })
    };

    println!("creating pipeline (unchecked={unchecked}) ...");
    let _pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
        label: None,
        layout: Some(&layout),
        vertex: wgpu::VertexState {
            module: &module,
            entry_point: Some("vs_main"),
            compilation_options: Default::default(),
            buffers: &[],
        },
        fragment: Some(wgpu::FragmentState {
            module: &module,
            entry_point: Some("fs_main"),
            compilation_options: Default::default(),
            targets: &[Some(wgpu::ColorTargetState {
                format: wgpu::TextureFormat::Rgba8Unorm,
                blend: None,
                write_mask: wgpu::ColorWrites::ALL,
            })],
        }),
        primitive: Default::default(),
        depth_stencil: None,
        multisample: Default::default(),
        multiview_mask: None,
        cache: None,
    });
    device.poll(wgpu::PollType::Wait { submission_index: None, timeout: None }).ok();
    println!("pipeline created ok");
}