Implement pruning unused depth reference variables

README.md

@@ -13,12 +13,13 @@ This project aims to transform common but unsupported SPIRV shaders into a form
 
 At the moment, the following transformations are supported:
 
-| Feature                   | `spirv-val` | `naga` | `tint` |
-| ------------------------- | ----------- | ------ | ------ |
-| Combined Image Samplers   | ✅          | ✅     | ✅     |
-| Mixed Depth / Comparison  | ✅          | ⚠️\*   | ❌     |
-| isnan / isinf Patching    | ✅          | ✅     | ✅     |
-| Storage Cube Patching     | ✅          | ✅     | ✅     |
+| Feature                           | `spirv-val` | `naga` | `tint` |
+| --------------------------------- | ----------- | ------ | ------ |
+| Combined Image Samplers           | ✅          | ✅     | ✅     |
+| Mixed Depth / Comparison          | ✅          | ⚠️\*   | ❌     |
+| isnan / isinf Patching            | ✅          | ✅     | ✅     |
+| Storage Cube Patching             | ✅          | ✅     | ✅     |
+| Unused Image Sampler Pruning      | ✅          | ✅     | ✅     |
 
 > \* Simple cases are OK.
 > With some [special patches](https://github.com/davnotdev/wgpu/tree/trunk-naga-patches), `naga` can process these.
@@ -159,6 +160,27 @@ void main() {
 - You *can* nest `imageCube` usages in functions, but this will not translate to WGSL
 - `imageCubeArray` is not supported
 
+## Unused Image Sampler Pruning
+
+In ubershader configurations, it is common to have uniforms that are referenced in some variants but not others.
+If you were to declare a `uniform texture2D depth_buffer;` that is unused, the WGSL translation is always `texture_2d`.
+However, in shaders that DO reference `depth_buffer`, you likely get a `texture_depth_2d`.
+This may cause conflicts in your bind group layouts or errors when creating bind groups.
+This patch prunes samplers and textures that are unused.
+
+### Tests
+
+| Test                           | `spirv-val` | Naga   | Tint |
+| ------------------------------ | ----------- | ------ | ---- |
+| `pruneunuseddref.frag`         | ✅          | ✅     | ✅   |
+| `pruneunuseddref_nested.frag`  | ✅          | ✅     | ✅   |
+| `pruneunuseddref_storage.frag` | ✅          | ✅     | ✅   |
+
+### Additional Notes
+
+- Does not prune combined image samplers
+- Does not prune storage textures since they are separate category of texture
+
 ## Library Usage
 
 Add this to your `Cargo.toml`:

ffi/bin/spv_webgpu_transform.c

@@ -1,89 +1,94 @@
+#include "../spirv_webgpu_transform.h"
+#include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
-#include <stdint.h>
-#include "../spirv_webgpu_transform.h"
 
 #define BAD_FILE_PATH "./bad.spv"
 
 void print_set_binding(TransformCorrectionMap map, uint32_t set, uint32_t binding);
 
 int main() {
-    // 1. Read the SPIRV file
-    FILE* file = fopen(BAD_FILE_PATH, "rb");
-    fseek(file, 0, SEEK_END);
-    int spirv_bytes = ftell(file);
-    fseek(file, 0, SEEK_SET);
-    uint8_t* spirv = (uint8_t*)malloc(spirv_bytes);
-    fread(spirv, 1, spirv_bytes, file);
-    fclose(file);
-
-    // 2. Run the transformations
-    TransformCorrectionMap correction_map = SPIRV_WEBGPU_TRANSFORM_CORRECTION_MAP_NULL;
-
-    uint32_t* comb_out_spv;
-    uint32_t comb_out_count;
-    spirv_webgpu_transform_combimgsampsplitter_alloc((uint32_t*)spirv, spirv_bytes / 4, &comb_out_spv, &comb_out_count, &correction_map);
-
-    uint32_t* dref_out_spv;
-    uint32_t dref_out_count;
-    spirv_webgpu_transform_drefsplitter_alloc(comb_out_spv, comb_out_count, &dref_out_spv, &dref_out_count, &correction_map);
-
-    uint32_t* isnanisinf_out_spv;
-    uint32_t isnanisinf_out_count;
-    spirv_webgpu_transform_isnanisinfpatch_alloc(dref_out_spv, dref_out_count, &isnanisinf_out_spv, &isnanisinf_out_count);
-
-    uint32_t* storagecube_out_spv;
-    uint32_t storagecube_out_count;
-    spirv_webgpu_transform_storagecubepatch_alloc(isnanisinf_out_spv, isnanisinf_out_count, &storagecube_out_spv, &storagecube_out_count, &correction_map);
-
-    // 3. Observe the patched variables
-    print_set_binding(correction_map, 0, 0);
-    print_set_binding(correction_map, 0, 1);
-    print_set_binding(correction_map, 1, 0);
-    print_set_binding(correction_map, 2, 0);
-
-    // Fluke values should return None
-    print_set_binding(correction_map, 1, 1);
-    print_set_binding(correction_map, 3, 0);
-
-    // 4. Free memory
-    spirv_webgpu_transform_storagecubepatch_free(storagecube_out_spv);
-    spirv_webgpu_transform_isnanisinfpatch_free(isnanisinf_out_spv);
-    spirv_webgpu_transform_drefsplitter_free(dref_out_spv);
-    spirv_webgpu_transform_combimgsampsplitter_free(comb_out_spv);
-    spirv_webgpu_transform_correction_map_free(correction_map);
-    free(spirv);
+	// 1. Read the SPIRV file
+	FILE *file = fopen(BAD_FILE_PATH, "rb");
+	fseek(file, 0, SEEK_END);
+	int spirv_bytes = ftell(file);
+	fseek(file, 0, SEEK_SET);
+	uint8_t *spirv = (uint8_t *)malloc(spirv_bytes);
+	fread(spirv, 1, spirv_bytes, file);
+	fclose(file);
+
+	// 2. Run the transformations
+	TransformCorrectionMap correction_map = SPIRV_WEBGPU_TRANSFORM_CORRECTION_MAP_NULL;
+
+	uint32_t *comb_out_spv;
+	uint32_t comb_out_count;
+	spirv_webgpu_transform_combimgsampsplitter_alloc((uint32_t *)spirv, spirv_bytes / 4, &comb_out_spv, &comb_out_count, &correction_map);
+
+	uint32_t *dref_out_spv;
+	uint32_t dref_out_count;
+	spirv_webgpu_transform_drefsplitter_alloc(comb_out_spv, comb_out_count, &dref_out_spv, &dref_out_count, &correction_map);
+
+	uint32_t *isnanisinf_out_spv;
+	uint32_t isnanisinf_out_count;
+	spirv_webgpu_transform_isnanisinfpatch_alloc(dref_out_spv, dref_out_count, &isnanisinf_out_spv, &isnanisinf_out_count);
+
+	uint32_t *storagecube_out_spv;
+	uint32_t storagecube_out_count;
+	spirv_webgpu_transform_storagecubepatch_alloc(isnanisinf_out_spv, isnanisinf_out_count, &storagecube_out_spv, &storagecube_out_count, &correction_map);
+
+	uint32_t *pruneunuseddref_out_spv;
+	uint32_t pruneunuseddref_out_count;
+	spirv_webgpu_transform_pruneunuseddref_alloc(storagecube_out_spv, storagecube_out_count, &pruneunuseddref_out_spv, &pruneunuseddref_out_count);
+
+	// 3. Observe the patched variables
+	print_set_binding(correction_map, 0, 0);
+	print_set_binding(correction_map, 0, 1);
+	print_set_binding(correction_map, 1, 0);
+	print_set_binding(correction_map, 2, 0);
+
+	// Fluke values should return None
+	print_set_binding(correction_map, 1, 1);
+	print_set_binding(correction_map, 3, 0);
+
+	// 4. Free memory
+	spirv_webgpu_transform_pruneunuseddref_free(pruneunuseddref_out_spv);
+	spirv_webgpu_transform_storagecubepatch_free(storagecube_out_spv);
+	spirv_webgpu_transform_isnanisinfpatch_free(isnanisinf_out_spv);
+	spirv_webgpu_transform_drefsplitter_free(dref_out_spv);
+	spirv_webgpu_transform_combimgsampsplitter_free(comb_out_spv);
+	spirv_webgpu_transform_correction_map_free(correction_map);
+	free(spirv);
 }
 
 void print_set_binding(TransformCorrectionMap map, uint32_t set, uint32_t binding) {
-    uint16_t* corrections;
-    uint32_t correction_count;
-    TransformCorrectionStatus status = spirv_webgpu_transform_correction_map_index(map, set, binding, &corrections, &correction_count);
-
-    printf("For set %d, binding %d:\n", set, binding);
-
-    if (status == SPIRV_WEBGPU_TRANSFORM_CORRECTION_STATUS_NONE)  {
-        printf("\tNone\n");
-    } else {
-        printf("\tSome\n");
-    }
-
-    printf("\t");
-    for (int i = 0; i < correction_count; i++) {
-        switch (corrections[i]) {
-            case SPIRV_WEBGPU_TRANSFORM_CORRECTION_TYPE_SPLIT_COMBINED:
-                printf("SPLIT_COMBINED ");
-                break;
-            case SPIRV_WEBGPU_TRANSFORM_CORRECTION_TYPE_SPLIT_DREF_REGULAR:
-                printf("SPLIT_DREF_REGULAR ");
-                break;
-            case SPIRV_WEBGPU_TRANSFORM_CORRECTION_TYPE_SPLIT_DREF_COMPARISON:
-                printf("SPLIT_DREF_COMPARISON ");
-                break;
-            case SPIRV_WEBGPU_TRANSFORM_CORRECTION_TYPE_CONVERT_STORAGE_CUBE:
-                printf("CONVERT_STORAGE_CUBE ");
-                break;
-        }
-    }
-    printf("\n");
+	uint16_t *corrections;
+	uint32_t correction_count;
+	TransformCorrectionStatus status = spirv_webgpu_transform_correction_map_index(map, set, binding, &corrections, &correction_count);
+
+	printf("For set %d, binding %d:\n", set, binding);
+
+	if (status == SPIRV_WEBGPU_TRANSFORM_CORRECTION_STATUS_NONE) {
+		printf("\tNone\n");
+	} else {
+		printf("\tSome\n");
+	}
+
+	printf("\t");
+	for (int i = 0; i < correction_count; i++) {
+		switch (corrections[i]) {
+			case SPIRV_WEBGPU_TRANSFORM_CORRECTION_TYPE_SPLIT_COMBINED:
+				printf("SPLIT_COMBINED ");
+				break;
+			case SPIRV_WEBGPU_TRANSFORM_CORRECTION_TYPE_SPLIT_DREF_REGULAR:
+				printf("SPLIT_DREF_REGULAR ");
+				break;
+			case SPIRV_WEBGPU_TRANSFORM_CORRECTION_TYPE_SPLIT_DREF_COMPARISON:
+				printf("SPLIT_DREF_COMPARISON ");
+				break;
+			case SPIRV_WEBGPU_TRANSFORM_CORRECTION_TYPE_CONVERT_STORAGE_CUBE:
+				printf("CONVERT_STORAGE_CUBE ");
+				break;
+		}
+	}
+	printf("\n");
 }

ffi/spirv_webgpu_transform.h

@@ -19,6 +19,8 @@ void spirv_webgpu_transform_isnanisinfpatch_alloc(uint32_t *in_spv, uint32_t in_
 void spirv_webgpu_transform_isnanisinfpatch_free(uint32_t *out_spv);
 void spirv_webgpu_transform_storagecubepatch_alloc(uint32_t *in_spv, uint32_t in_count, uint32_t **out_spv, uint32_t *out_count, TransformCorrectionMap *correction_map);
 void spirv_webgpu_transform_storagecubepatch_free(uint32_t *out_spv);
+void spirv_webgpu_transform_pruneunuseddref_alloc(uint32_t *int_spv, uint32_t in_count, uint32_t **out_spv, uint32_t *out_count);
+void spirv_webgpu_transform_pruneunuseddref_free(uint32_t *out_spv);
 
 void spirv_webgpu_transform_mirrorpatch_alloc(
 		uint32_t *in_left_spv, uint32_t in_left_count, TransformCorrectionMap *left_corrections,

ffi/src/lib.rs

@@ -3,7 +3,7 @@
 use core::{ffi, ptr, slice};
 use spirv_webgpu_transform::{
     CorrectionMap, combimgsampsplitter, drefsplitter, isnanisinfpatch, mirrorpatch,
-    storagecubepatch,
+    pruneunuseddref, storagecubepatch,
 };
 
 type TransformCorrectionMap = *mut ffi::c_void;
@@ -147,6 +147,32 @@ pub unsafe extern "C" fn spirv_webgpu_transform_storagecubepatch_free(out_spv: *
     unsafe { drop(Box::from_raw(out_spv)) }
 }
 
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn spirv_webgpu_transform_pruneunuseddref_alloc(
+    in_spv: *const u32,
+    in_count: u32,
+    out_spv: *mut *const u32,
+    out_count: *mut u32,
+) {
+    let in_spv = unsafe { slice::from_raw_parts(in_spv, in_count as usize) };
+    match pruneunuseddref(in_spv) {
+        Ok(spv) => unsafe {
+            *out_count = spv.len() as u32;
+            let leaked = Box::leak(spv.into_boxed_slice());
+            *out_spv = leaked.as_ptr();
+        },
+        Err(_) => unsafe {
+            *out_spv = ptr::null();
+            *out_count = 0;
+        },
+    }
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn spirv_webgpu_transform_pruneunuseddref_free(out_spv: *mut u32) {
+    unsafe { drop(Box::from_raw(out_spv)) }
+}
+
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn spirv_webgpu_transform_mirrorpatch_alloc(
     in_left_spv: *const u32,

src/bin/spv_webgpu_transform.rs

@@ -5,7 +5,7 @@ fn main() {
 
     if args.len() != 4 {
         eprintln!(
-            "Usage: spv_webgpu_transform <combimg|dref|isnanisinf|storagecube> <input.spv> <output.spv>"
+            "Usage: spv_webgpu_transform <combimg|dref|isnanisinf|storagecube|pruneunuseddref> <input.spv> <output.spv>"
         );
         process::exit(1);
     }
@@ -30,6 +30,7 @@ fn main() {
         "storagecube" => {
             spirv_webgpu_transform::storagecubepatch(&spv, &mut out_correction_map).unwrap()
         }
+        "pruneunuseddref" => spirv_webgpu_transform::pruneunuseddref(&spv).unwrap(),
         mode => {
             eprintln!("unknown mode {:?}", mode);
             process::exit(1)

src/lib.rs

@@ -31,6 +31,7 @@ use std::collections::{HashMap, HashSet};
 mod correction;
 mod isnanisinfpatch;
 mod mirrorpatch;
+mod pruneunuseddref;
 mod splitcombined;
 mod splitdref;
 mod spv;
@@ -46,6 +47,7 @@ use util::*;
 pub use correction::*;
 pub use isnanisinfpatch::*;
 pub use mirrorpatch::*;
+pub use pruneunuseddref::*;
 pub use splitcombined::*;
 pub use splitdref::*;
 pub use storagecubepatch::*;