Metal: Texture write fix (#4952)

4 files changed, 60 insertions, 59 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index ed1c9fcb6..907822060 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -3498,6 +3498,14 @@ extension _Texture<T,Shape,isArray,0,sampleCount,$(access),isShadow, 0,format>
     {
         __intrinsic_asm "imageStore($0, $1, $V2)";
     }
+    
+    [require(metal, texture_sm_4_1)]
+    __intrinsic_op($(kIROp_ImageStore))
+    static void __metalImageStoreArray(This val, vector<uint, Shape.dimensions> location, T value, uint arrayIndex);
+
+    [require(metal, texture_sm_4_1)]
+    __intrinsic_op($(kIROp_ImageStore))
+    static void __metalImageStore(This val, vector<uint, Shape.dimensions+isArray> location, T value);
 
     __subscript(vector<uint, Shape.dimensions+isArray> location) -> T
     {
@@ -3567,55 +3575,14 @@ extension _Texture<T,Shape,isArray,0,sampleCount,$(access),isShadow, 0,format>
                         OpImageWrite $this $location __convertTexel(newValue);
                     };
                 case metal:
-                    switch (Shape.flavor)
+                    if (isArray != 0)
                     {
-                    case $(SLANG_TEXTURE_1D):
-                        // lod is not supported for 1D texture
-                        if (isArray == 1)
-                            // void write(Tv color, uint coord, uint array, uint lod = 0) const
-                            __intrinsic_asm "$0.write($2, uint(($1).x), uint(($1).y))";
-                        else
-                            // void write(Tv color, uint coord, uint lod = 0) const
-                            __intrinsic_asm "$0.write($2, uint($1))";
-                        break;
-                    case $(SLANG_TEXTURE_2D):
-                        if (isShadow == 1)
-                        {
-                            if (isArray == 1)
-                                // void write(Tv color, uint2 coord, uint array, uint lod = 0) const
-                                __intrinsic_asm "$0.write($2, vec<uint,2>(($1).xy), uint(($1).z))";
-                            else
-                                // void write(Tv color, uint2 coord, uint lod = 0) const
-                                __intrinsic_asm "$0.write($2, vec<uint,2>(($1).xy))";
-                        }
-                        else
-                        {
-                            if (isArray == 1)
-                                // void write(Tv color, uint2 coord, uint array, uint lod = 0) const
-                                __intrinsic_asm "$0.write($2, vec<uint,2>(($1).xy), uint(($1).z))";
-                            else
-                                // void write(Tv color, uint2 coord, uint lod = 0) const
-                                __intrinsic_asm "$0.write($2, vec<uint,2>(($1).xy))";
-                        }
-                        break;
-                    case $(SLANG_TEXTURE_3D):
-                        if (isShadow == 0 && isArray == 0)
-                            // void write(Tv color, uint3 coord, uint lod = 0) const
-                            __intrinsic_asm "$0.write($2, vec<uint,3>(($1).xyz))";
-                        break;
-                    case $(SLANG_TEXTURE_CUBE):
-                        static_assert(isArray == 0, "Unsupported 'Store' of 'texture cube array' for 'metal' target");
-                        if (isShadow == 1)
-                        {
-                            // void write(Tv color, uint2 coord, uint face, uint lod = 0) const
-                            __intrinsic_asm "$0.write($2, vec<uint,2>(($1).xy), uint(($1).z), uint(($1).w))";
-                        }
-                        else
-                        {
-                            // void write(Tv color, uint2 coord, uint face, uint lod = 0) const
-                            __intrinsic_asm "$0.write($2, vec<uint,2>(($1).xy), uint(($1).z), uint(($1).w))";
-                        }
-                        break;
+                        // last arg will be replaced with the split off array index
+                        __metalImageStoreArray(this, __vectorReshape<Shape.dimensions>(location), newValue, location[Shape.dimensions + isArray - 1]);
+                    }
+                    else
+                    {
+                        __metalImageStore(this, location, newValue);
                     }
             case wgsl:
                 static_assert(Shape.flavor == $(SLANG_TEXTURE_1D)
diff --git a/source/slang/slang-emit-metal.cpp b/source/slang/slang-emit-metal.cpp
index 56b55d57a..2d5a7d56b 100644
--- a/source/slang/slang-emit-metal.cpp
+++ b/source/slang/slang-emit-metal.cpp
@@ -644,7 +644,6 @@ bool MetalSourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inO
         }
         case kIROp_ImageStore:
         {
-            
             auto imageOp = as<IRImageStore>(inst);
             emitOperand(imageOp->getImage(), getInfo(EmitOp::General));
             m_writer->emit(".write(");
@@ -656,11 +655,6 @@ bool MetalSourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inO
                 m_writer->emit(",");
                 emitOperand(imageOp->getAuxCoord1(), getInfo(EmitOp::General));
             }
-            if(imageOp->hasAuxCoord2())
-            {
-                m_writer->emit(",");
-                emitOperand(imageOp->getAuxCoord2(), getInfo(EmitOp::General));
-            }
             m_writer->emit(")");
             return true;
         }
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index c2fd9da0a..eef9b86bc 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -2539,13 +2539,9 @@ struct IRImageStore : IRInst
     IRInst* getValue() { return getOperand(2); }
 
     /// If GLSL/SPIR-V, Sample coord
-    /// If Metal, Array or Sample coord
+    /// Metal array/face index
     bool hasAuxCoord1() { return getOperandCount() > 3 && getOperand(3) != nullptr; }
     IRInst* getAuxCoord1() { return getOperand(3); }
-
-    /// If Metal, Sample coord
-    bool hasAuxCoord2() { return getOperandCount() > 4 && getOperand(4) != nullptr; }
-    IRInst* getAuxCoord2() { return getOperand(4); }
 };
 // Terminators
 
diff --git a/source/slang/slang-ir-metal-legalize.cpp b/source/slang/slang-ir-metal-legalize.cpp
index 00f70b00d..a92039210 100644
--- a/source/slang/slang-ir-metal-legalize.cpp
+++ b/source/slang/slang-ir-metal-legalize.cpp
@@ -1740,6 +1740,46 @@ namespace Slang
         }
     };
 
+    // metal textures only support writing 4-component values, even if the texture is only 1, 2, or 3-component
+    // in this case the other channels get ignored, but the signature still doesnt match
+    // so now we have to replace the value being written with a 4-component vector where 
+    // the new components get ignored, nice
+    void legalizeImageStoreValue(IRBuilder& builder, IRImageStore* imageStore)
+    {
+        builder.setInsertBefore(imageStore);
+        auto originalValue = imageStore->getValue();
+        auto valueBaseType = originalValue->getDataType();
+        IRType* elementType = nullptr;
+        List<IRInst*> components;
+        if(auto valueVectorType = as<IRVectorType>(valueBaseType))
+        {
+            if(auto originalElementCount = as<IRIntLit>(valueVectorType->getElementCount()))
+            {
+                if(originalElementCount->getValue() == 4)
+                {
+                    return;
+                }
+            }
+            elementType = valueVectorType->getElementType();
+            auto vectorValue = as<IRMakeVector>(originalValue);
+            for(UInt i = 0; i < vectorValue->getOperandCount(); i++)
+            {
+                components.add(vectorValue->getOperand(i));
+            }
+        }
+        else
+        {
+            elementType = valueBaseType;
+            components.add(originalValue);
+        }
+        for(UInt i = components.getCount(); i < 4; i++)
+        {
+            components.add(builder.getIntValue(builder.getIntType(), 0));
+        }
+        auto fourComponentVectorType = builder.getVectorType(elementType, 4);
+        imageStore->setOperand(2, builder.emitMakeVector(fourComponentVectorType, components));
+    }
+
     void legalizeFuncBody(IRFunc* func)
     {
         IRBuilder builder(func);
@@ -1796,6 +1836,10 @@ namespace Slang
                         arg->set(temp);
                     }
                 }
+                if(auto write = as<IRImageStore>(inst))
+                {
+                    legalizeImageStoreValue(builder, write);
+                }
             }
         }
     }