Metal: Texture write fix

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)

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;
         }

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
 

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);
+                }
             }
         }
     }

tests/metal/texture-write.slang

@@ -0,0 +1,57 @@
+//TEST:SIMPLE(filecheck=METAL): -stage compute -entry computeMain -target metal -DEMIT_SOURCE
+//TEST:SIMPLE(filecheck=METALLIB): -stage compute -entry computeMain -target metallib
+// for some reason, metal textures dont have an overload for less-than-four component
+// writes, they need to be converted to 4-components in a legalize step, as the other components
+// get discarded
+struct TextureWrite
+{
+    //TEST_INPUT: RWTexture2D(size=4, content = zero):name t2D_f32
+    RWTexture2D<float> tex1;
+    //TEST_INPUT: RWTexture2D(size=4, content = zero):name t2D_f32v2
+    RWTexture2D<float2> tex2;
+    //TEST_INPUT: RWTexture2D(size=4, content = zero):name t2D_f32v3
+    RWTexture2D<float3> tex3;
+    //TEST_INPUT: RWTexture2D(size=4, content = zero):name t2D_f32v4
+    RWTexture2D<float4> tex4;
+
+    //TEST_INPUT: RWTexture2DArray(size=4, content = zero):name t2D_f32
+    RWTexture2DArray<float> tex1Array;
+    //TEST_INPUT: RWTexture2DArray(size=4, content = zero):name t2D_f32v2
+    RWTexture2DArray<float2> tex2Array;
+    //TEST_INPUT: RWTexture2DArray(size=4, content = zero):name t2D_f32v3
+    RWTexture2DArray<float3> tex3Array;
+    //TEST_INPUT: RWTexture2DArray(size=4, content = zero):name t2D_f32v4
+    RWTexture2DArray<float4> tex4Array;
+}
+ParameterBlock<TextureWrite> pWrites;
+
+[numthreads(1, 1, 1)]
+void computeMain()
+{
+    // TODO: check for the type of first parameter to be a 4-component vector
+    // METALLIB: call {{.*}}.write_texture_2d.v4f32(
+    // METAL: .write(
+    pWrites.tex1[uint2(1, 1)] = 1;
+    // METALLIB: call {{.*}}.write_texture_2d.v4f32(
+    // METAL: .write(
+    pWrites.tex2[uint2(2, 2)] = float2(1, 2);
+    // METALLIB: call {{.*}}.write_texture_2d.v4f32(
+    // METAL: .write(
+    pWrites.tex3[uint2(3, 3)] = float3(1, 2, 3);
+    // METALLIB: call {{.*}}.write_texture_2d.v4f32(
+    // METAL: .write(
+    pWrites.tex4[uint2(4, 4)] = float4(1, 2, 3, 4);
+
+    // METALLIB: call {{.*}}.write_texture_2d_array.v4f32(
+    // METAL: .write(
+    pWrites.tex1Array[uint3(1, 1, 1)] = 1;
+    // METALLIB: call {{.*}}.write_texture_2d_array.v4f32(
+    // METAL: .write(
+    pWrites.tex2Array[uint3(2, 2, 2)] = float2(1, 2);
+    // METALLIB: call {{.*}}.write_texture_2d_array.v4f32(
+    // METAL: .write(
+    pWrites.tex3Array[uint3(3, 3, 3)] = float3(1, 2, 3);
+    // METALLIB: call {{.*}}.write_texture_2d_array.v4f32(
+    // METAL: .write(
+    pWrites.tex4Array[uint3(4, 4, 4)] = float4(1, 2, 3, 4);
+}