Metal remove void field (#6725)

* Reapply "Eliminate empty struct on metal target (#6603)" (#6711)

This reverts commit bc9dc6557fc0cc3a4c0c2ff27e636940e361cf5d.

* Remove argument in make_struct call corresponding to void field

This is a follow-up of #6543, where we leave the VoidType field as it in
make_struct call during legalization pass.

So during cleaning_void IR pass, when we remove "VoidType" from struct,
we will have to also clean up the argument corresponding to the
"VoidType" field.

6 files changed, 87 insertions, 1 deletions

diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 841f44a80..8d7577b52 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -1278,6 +1278,15 @@ Result linkAndOptimizeIR(
         //
         legalizeResourceTypes(targetProgram, irModule, sink);
 
+        // We also need to legalize empty types for Metal targets.
+        switch (target)
+        {
+        case CodeGenTarget::Metal:
+        case CodeGenTarget::MetalLib:
+        case CodeGenTarget::MetalLibAssembly:
+            legalizeEmptyTypes(targetProgram, irModule, sink);
+            break;
+        }
         //  Debugging output of legalization
 #if 0
         dumpIRIfEnabled(codeGenContext, irModule, "LEGALIZED");
diff --git a/source/slang/slang-ir-cleanup-void.cpp b/source/slang/slang-ir-cleanup-void.cpp
index 3a776fc4f..84532d0ec 100644
--- a/source/slang/slang-ir-cleanup-void.cpp
+++ b/source/slang/slang-ir-cleanup-void.cpp
@@ -122,6 +122,8 @@ struct CleanUpVoidContext
         case kIROp_StructType:
             {
                 List<IRInst*> toRemove;
+                UInt fieldCount = 0;
+                ShortList<UInt> voidFieldIndex;
                 for (auto child : inst->getChildren())
                 {
                     if (auto field = as<IRStructField>(child))
@@ -129,11 +131,33 @@ struct CleanUpVoidContext
                         if (field->getFieldType()->getOp() == kIROp_VoidType)
                         {
                             toRemove.add(field);
+                            voidFieldIndex.add(fieldCount);
                         }
                     }
+                    fieldCount++;
                 }
                 for (auto ii : toRemove)
                     ii->removeAndDeallocate();
+
+                // Once we remove the void fields in the struct, we also need update the make_struct
+                // call sites to remove the arguments corresponding to the void fields.
+                if (inst->hasUses() && voidFieldIndex.getCount())
+                {
+                    UInt currentFieldCount = fieldCount - toRemove.getCount();
+                    for (auto use = inst->firstUse; use; use = use->nextUse)
+                    {
+                        if (auto makeStructInst = as<IRMakeStruct>(use->user))
+                        {
+                            if (makeStructInst->getOperandCount() != currentFieldCount)
+                            {
+                                for (Int i = 0; i < voidFieldIndex.getCount(); i++)
+                                {
+                                    makeStructInst->removeOperand(voidFieldIndex[i]);
+                                }
+                            }
+                        }
+                    }
+                }
             }
             break;
         default:
diff --git a/source/slang/slang-ir-legalize-types.cpp b/source/slang/slang-ir-legalize-types.cpp
index 197cd93c5..837193abc 100644
--- a/source/slang/slang-ir-legalize-types.cpp
+++ b/source/slang/slang-ir-legalize-types.cpp
@@ -4123,6 +4123,11 @@ struct IREmptyTypeLegalizationContext : IRTypeLegalizationContext
 
     bool isSimpleType(IRType* type) override
     {
+        if (isMetalTarget(targetProgram->getTargetReq()))
+        {
+            return false;
+        }
+
         // If type is used as public interface, then treat it as simple.
         for (auto decor : type->getDecorations())
         {
@@ -4145,6 +4150,11 @@ struct IREmptyTypeLegalizationContext : IRTypeLegalizationContext
     {
         return LegalType();
     }
+
+    virtual bool shouldLegalizeParameterBlockElementType() override
+    {
+        return isMetalTarget(targetProgram->getTargetReq());
+    }
 };
 
 // The main entry points that are used when transforming IR code
diff --git a/source/slang/slang-legalize-types.cpp b/source/slang/slang-legalize-types.cpp
index e26475522..cc2c12d42 100644
--- a/source/slang/slang-legalize-types.cpp
+++ b/source/slang/slang-legalize-types.cpp
@@ -1211,11 +1211,15 @@ LegalType legalizeTypeImpl(TypeLegalizationContext* context, IRType* type)
         LegalType legalElementType;
 
         if (isMetalTarget(context->targetProgram->getTargetReq()) &&
-            as<IRParameterBlockType>(uniformBufferType))
+            as<IRParameterBlockType>(uniformBufferType) &&
+            !context->shouldLegalizeParameterBlockElementType())
         {
             // On Metal, we do not need to legalize the element type of
             // a parameter block because we can translate it directly into
             // an argument buffer.
+            //
+            // But we do need empty type legalized for Metal, because Metal doesn't
+            // allow empty struct in argument buffer.
             legalElementType = LegalType::simple(originalElementType);
         }
         else
diff --git a/source/slang/slang-legalize-types.h b/source/slang/slang-legalize-types.h
index ae76cbd39..0c705c4fa 100644
--- a/source/slang/slang-legalize-types.h
+++ b/source/slang/slang-legalize-types.h
@@ -659,6 +659,12 @@ struct IRTypeLegalizationContext
         IROp op,
         LegalType legalElementType,
         IRInst* layoutOperand) = 0;
+
+    /// Customization point to decide whether a parameter block type should be legalized.
+    ///
+    /// This function is called in `legalizeTypeImpl` to decide whether a parameter block
+    /// type should be legalized. Not all legalization passes need to legalize parameter block.
+    virtual bool shouldLegalizeParameterBlockElementType() { return false; }
 };
 
 // This typedef exists to support pre-existing code from when
diff --git a/tests/metal/empty-struct-remove.slang b/tests/metal/empty-struct-remove.slang
new file mode 100644
index 000000000..5ac1ccb69
--- /dev/null
+++ b/tests/metal/empty-struct-remove.slang
@@ -0,0 +1,33 @@
+
+//TEST:SIMPLE(filecheck=LIB):-target metallib -entry computeMain -stage compute
+//TEST:SIMPLE(filecheck=METAL):-target metal -entry computeMain -stage compute
+
+// METAL-NOT: struct emptyStruct
+struct emptyStruct
+{
+    void set(RWStructuredBuffer<int> buffer, int value) {buffer[0] = value;}
+}
+
+
+struct MyStruct
+{
+    RWStructuredBuffer<int> buffer;
+    int value;
+    void set()
+    {
+        e.set(buffer, value);
+    }
+    emptyStruct e;
+}
+
+ParameterBlock<MyStruct> param;
+
+// LIB: @computeMain
+[shader("compute")]
+[numthreads(1, 1, 1)]
+void computeMain(
+    uint tid: SV_DispatchThreadID,
+)
+{
+    param.set();
+}