MSL: Emit spvTexelBufferCoord() on ImageWrite to a Buffer as well.

This is necessary to get the coordinates to give to the texture's
`write()` method.

reference/opt/shaders-msl/asm/comp/buffer-write.asm.comp

@@ -0,0 +1,23 @@
+#pragma clang diagnostic ignored "-Wmissing-prototypes"
+
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct cb
+{
+    float value;
+};
+
+// Returns 2D texture coords corresponding to 1D texel buffer coords
+uint2 spvTexelBufferCoord(uint tc)
+{
+    return uint2(tc % 4096, tc / 4096);
+}
+
+kernel void main0(constant cb& _6 [[buffer(7)]], texture2d<float, access::write> _buffer [[texture(0)]], uint3 gl_WorkGroupID [[threadgroup_position_in_grid]], uint gl_LocalInvocationIndex [[thread_index_in_threadgroup]])
+{
+    _buffer.write(_6.value, spvTexelBufferCoord(((32u * gl_WorkGroupID.x) + gl_LocalInvocationIndex)));
+}
+

reference/shaders-msl/asm/comp/buffer-write.asm.comp

@@ -0,0 +1,23 @@
+#pragma clang diagnostic ignored "-Wmissing-prototypes"
+
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct cb
+{
+    float value;
+};
+
+// Returns 2D texture coords corresponding to 1D texel buffer coords
+uint2 spvTexelBufferCoord(uint tc)
+{
+    return uint2(tc % 4096, tc / 4096);
+}
+
+kernel void main0(constant cb& _6 [[buffer(7)]], texture2d<float, access::write> _buffer [[texture(0)]], uint3 gl_WorkGroupID [[threadgroup_position_in_grid]], uint gl_LocalInvocationIndex [[thread_index_in_threadgroup]])
+{
+    _buffer.write(_6.value, spvTexelBufferCoord(((32u * gl_WorkGroupID.x) + gl_LocalInvocationIndex)));
+}
+

shaders-msl/asm/comp/buffer-write.asm.comp

@@ -0,0 +1,58 @@
+; SPIR-V
+; Version: 1.3
+; Generator: Khronos Glslang Reference Front End; 7
+; Bound: 63
+; Schema: 0
+               OpCapability Shader
+               OpCapability ImageBuffer
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %main "main" %group_id %group_index
+               OpExecutionMode %main LocalSize 32 1 1
+               OpSource HLSL 500
+               OpName %main "main"
+               OpName %cb "cb"
+               OpMemberName %cb 0 "value"
+               OpName %_ ""
+               OpName %buffer "buffer"
+               OpName %group_id "group_id"
+               OpName %group_index "group_index"
+               OpMemberDecorate %cb 0 Offset 0
+               OpDecorate %cb Block
+               OpDecorate %_ DescriptorSet 0
+               OpDecorate %_ Binding 7
+               OpDecorate %buffer DescriptorSet 0
+               OpDecorate %group_id BuiltIn WorkgroupId
+               OpDecorate %group_index BuiltIn LocalInvocationIndex
+       %void = OpTypeVoid
+          %3 = OpTypeFunction %void
+       %uint = OpTypeInt 32 0
+     %v3uint = OpTypeVector %uint 3
+    %uint_32 = OpConstant %uint 32
+      %float = OpTypeFloat 32
+         %cb = OpTypeStruct %float
+%_ptr_Uniform_cb = OpTypePointer Uniform %cb
+          %_ = OpVariable %_ptr_Uniform_cb Uniform
+        %int = OpTypeInt 32 1
+      %int_0 = OpConstant %int 0
+%_ptr_Uniform_float = OpTypePointer Uniform %float
+         %34 = OpTypeImage %float Buffer 0 0 0 2 R32f
+%_ptr_UniformConstant_34 = OpTypePointer UniformConstant %34
+     %buffer = OpVariable %_ptr_UniformConstant_34 UniformConstant
+%_ptr_Input_v3uint = OpTypePointer Input %v3uint
+   %group_id = OpVariable %_ptr_Input_v3uint Input
+%_ptr_Input_uint = OpTypePointer Input %uint
+%group_index = OpVariable %_ptr_Input_uint Input
+       %main = OpFunction %void None %3
+          %5 = OpLabel
+         %43 = OpLoad %v3uint %group_id
+         %47 = OpLoad %uint %group_index
+         %56 = OpCompositeExtract %uint %43 0
+         %57 = OpIMul %uint %uint_32 %56
+         %59 = OpIAdd %uint %57 %47
+         %60 = OpAccessChain %_ptr_Uniform_float %_ %int_0
+         %61 = OpLoad %float %60
+         %62 = OpLoad %34 %buffer
+               OpImageWrite %62 %59 %61
+               OpReturn
+               OpFunctionEnd

spirv_msl.cpp

@@ -4333,9 +4333,10 @@ CompilerMSL::SPVFuncImpl CompilerMSL::OpCodePreprocessor::get_spv_func_impl(Op o
 	}
 
 	case OpImageFetch:
+	case OpImageWrite:
 	{
 		// Retrieve the image type, and if it's a Buffer, emit a texel coordinate function
-		uint32_t tid = result_types[args[2]];
+		uint32_t tid = result_types[args[opcode == OpImageWrite ? 0 : 2]];
 		if (tid && compiler.get<SPIRType>(tid).image.dim == DimBuffer)
 			return SPVFuncImplTexelBufferCoords;