Transfer offset decorations when legalizing laid-out structs

Struct legalization removing fields not representable in memory should transfer offset decorations in case the struct has already had offsets calculated. Closes shader-slang#5264.
aleino-nv · Nov 8, 2024 · c8b88c3 · c8b88c3
1 parent ad72ee6
commit c8b88c3
Show file tree

Hide file tree

Showing 7 changed files with 28 additions and 10 deletions.
diff --git a/source/slang/slang-ir.h b/source/slang/slang-ir.h
@@ -593,6 +593,18 @@ struct IRInst
     IRDecoration* getLastDecoration();
     IRInstList<IRDecoration> getDecorations();
 
+    // Step to the next child of a certain type, T
+    template<typename T>
+    T* getNextChildOfType(IRInst* child = nullptr)
+    {
+        for (IRInst* cc = child ? child->getNextInst() : getFirstDecorationOrChild(); cc != nullptr; cc = cc->getNextInst())
+        {
+            if (T* cct = as<T>(cc))
+                return cct;
+        }
+        return nullptr;
+    }
+
     // Look up a decoration in the list of decorations
     IRDecoration* findDecorationImpl(IROp op);
     template<typename T>
@@ -780,8 +792,19 @@ struct IRInst
     void removeOperand(Index index);
 
     /// Transfer any decorations of this instruction to the `target` instruction.
+    // TODO: rewrite as transferDecorationsTo(target) -> transferChildrenOfTypeTo<IRDecoration>(target)
     void transferDecorationsTo(IRInst* target);
 
+    template <typename T>
+    void transferChildrenOfTypeTo(IRInst* target)
+    {
+        for (T* cc = getNextChildOfType<T>(); cc != nullptr; cc = getNextChildOfType<T>(cc))
+        {
+            cc->removeFromParent();
+            cc->insertAtStart(target);
+        }
+    }
+
     /// Does this instruction have any uses?
     bool hasUses() const { return firstUse != nullptr; }
 

diff --git a/source/slang/slang-legalize-types.cpp b/source/slang/slang-legalize-types.cpp
@@ -487,8 +487,11 @@ struct TupleTypeBuilder
                     fieldType = context->getBuilder()->getVoidType();
 
                 // TODO: shallow clone of modifiers, etc.
-
-                builder->createStructField(ordinaryStructType, ee.fieldKey, fieldType);
+                IRStructField* originalField = findStructField(originalStructType, ee.fieldKey);
+                IRStructField* newField = builder->createStructField(ordinaryStructType, ee.fieldKey, fieldType);
+                // In case the original struct had offsets calculated, transfer those as well.
+                // The original offsets should still be valid, since we only skip fields of types that aren't representable in memory.
+                originalField->transferChildrenOfTypeTo<IROffsetDecoration>(newField);
             }
 
             ordinaryType = LegalType::simple((IRType*)ordinaryStructType);

diff --git a/tests/autodiff/existential-1.slang b/tests/autodiff/existential-1.slang
@@ -2,7 +2,6 @@
 
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
-//DISABLE_TEST(compute):COMPARE_COMPUTE_EX:-wgpu
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;

diff --git a/tests/autodiff/existential-2.slang b/tests/autodiff/existential-2.slang
@@ -2,7 +2,6 @@
 
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
-//DISABLE_TEST(compute):COMPARE_COMPUTE_EX:-wgpu
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;

diff --git a/tests/autodiff/material/diff-bwd-falcor-material-system.slang b/tests/autodiff/material/diff-bwd-falcor-material-system.slang
@@ -1,6 +1,5 @@
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
-//DISABLE_TEST(compute):COMPARE_COMPUTE_EX:-wgpu
 
 // outputBuffer is defined in IBSDF.slang
 //TEST_INPUT:ubuffer(data=[0 0 0 0 0 0], stride=4):out,name=outputBuffer

diff --git a/tests/autodiff/material2/diff-bwd-falcor-material-system.slang b/tests/autodiff/material2/diff-bwd-falcor-material-system.slang
@@ -1,6 +1,5 @@
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
-//DISABLE_TEST(compute):COMPARE_COMPUTE_EX:-wgpu
 
 // outputBuffer is defined in IBSDF.slang
 //TEST_INPUT:ubuffer(data=[0 0 0 0 0 0 0 0 0], stride=4):out,name=outputBuffer

diff --git a/tests/expected-failure-github.txt b/tests/expected-failure-github.txt
@@ -6,10 +6,6 @@ tests/language-feature/saturated-cooperation/fuse.slang (vk)
 tests/bugs/byte-address-buffer-interlocked-add-f32.slang (vk)
 tests/serialization/obfuscated-serialized-module-test.slang.2 syn (mtl)
 tests/autodiff/custom-intrinsic.slang.2 syn (wgpu)
-tests/autodiff/existential-1.slang.2 syn (wgpu)
-tests/autodiff/existential-2.slang.2 syn (wgpu)
-tests/autodiff/material/diff-bwd-falcor-material-system.slang.2 syn (wgpu)
-tests/autodiff/material2/diff-bwd-falcor-material-system.slang.2 syn (wgpu)
 tests/bugs/atomic-coerce.slang.3 syn (wgpu)
 tests/bugs/buffer-swizzle-store.slang.3 syn (wgpu)
 tests/bugs/dxbc-double-problem.slang.6 syn (wgpu)