MSL: Emit correct address space when casting during OpStore.

When storing to local variable (eg. OpCopyLogical), the
default device address space used during casts is illegal.
Determine correct address space based on variable type.

3 files changed, 123 insertions, 4 deletions

diff --git a/reference/shaders-msl-no-opt/asm/comp/copy-logical-offset-and-array-stride-diffs.spv14.asm.comp b/reference/shaders-msl-no-opt/asm/comp/copy-logical-offset-and-array-stride-diffs.spv14.asm.comp
new file mode 100644
index 00000000..54087ddc
--- /dev/null
+++ b/reference/shaders-msl-no-opt/asm/comp/copy-logical-offset-and-array-stride-diffs.spv14.asm.comp
@@ -0,0 +1,54 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct _8
+{
+    char _m0_pad[4];
+    uint _m0;
+};
+
+struct _9
+{
+    char _m0_pad[8];
+    uint _m0;
+};
+
+struct _4
+{
+    uint _m0;
+    uint4 _m1[2];
+    uint _m2;
+    char _m3_pad[12];
+    _8 _m3;
+    float4 _m4;
+    float3 _m5;
+    float2 _m6;
+};
+
+struct _5
+{
+    uint _m0;
+    uint _m1[2];
+    uint _m2;
+    _9 _m3;
+    float4 _m4;
+    float3 _m5;
+    float2 _m6;
+};
+
+kernel void main0(device _5& _2 [[buffer(0)]], device _4& _3 [[buffer(1)]])
+{
+    _4 _23;
+    _23._m0 = _2._m0;
+    (thread uint&)_23._m1[0] = _2._m1[0];
+    (thread uint&)_23._m1[1] = _2._m1[1];
+    _23._m2 = _2._m2;
+    _23._m3._m0 = _2._m3._m0;
+    _23._m4 = _2._m4;
+    _23._m5 = _2._m5;
+    _23._m6 = _2._m6;
+    _3 = _23;
+}
+
diff --git a/shaders-msl-no-opt/asm/comp/copy-logical-offset-and-array-stride-diffs.spv14.asm.comp b/shaders-msl-no-opt/asm/comp/copy-logical-offset-and-array-stride-diffs.spv14.asm.comp
new file mode 100644
index 00000000..026bd113
--- /dev/null
+++ b/shaders-msl-no-opt/asm/comp/copy-logical-offset-and-array-stride-diffs.spv14.asm.comp
@@ -0,0 +1,60 @@
+; SPIR-V
+; Version: 1.4
+; Generator: Khronos SPIR-V Tools Assembler; 0
+; Bound: 24
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %2 "main" %3 %4
+               OpExecutionMode %2 LocalSize 1 1 1
+               OpDecorate %5 Block
+               OpMemberDecorate %5 0 Offset 0
+               OpMemberDecorate %5 1 Offset 16
+               OpMemberDecorate %5 2 Offset 48
+               OpMemberDecorate %5 3 Offset 64
+               OpMemberDecorate %5 4 Offset 80
+               OpMemberDecorate %5 5 Offset 96
+               OpMemberDecorate %5 6 Offset 112
+               OpDecorate %6 Block
+               OpMemberDecorate %6 0 Offset 0
+               OpMemberDecorate %6 1 Offset 4
+               OpMemberDecorate %6 2 Offset 12
+               OpMemberDecorate %6 3 Offset 16
+               OpMemberDecorate %6 4 Offset 32
+               OpMemberDecorate %6 5 Offset 48
+               OpMemberDecorate %6 6 Offset 64
+               OpDecorate %3 DescriptorSet 0
+               OpDecorate %3 Binding 0
+               OpDecorate %4 DescriptorSet 0
+               OpDecorate %4 Binding 1
+               OpDecorate %7 ArrayStride 4
+               OpDecorate %8 ArrayStride 16
+               OpMemberDecorate %9 0 Offset 4
+               OpMemberDecorate %10 0 Offset 8
+         %11 = OpTypeVoid
+         %12 = OpTypeFloat 32
+         %13 = OpTypeVector %12 2
+         %14 = OpTypeVector %12 3
+         %15 = OpTypeVector %12 4
+         %16 = OpTypeMatrix %15 4
+         %17 = OpTypeInt 32 0
+         %18 = OpConstant %17 2
+          %7 = OpTypeArray %17 %18
+          %8 = OpTypeArray %17 %18
+          %9 = OpTypeStruct %17
+         %10 = OpTypeStruct %17
+          %5 = OpTypeStruct %17 %8 %17 %9 %15 %14 %13
+         %19 = OpTypePointer StorageBuffer %5
+          %6 = OpTypeStruct %17 %7 %17 %10 %15 %14 %13
+         %20 = OpTypePointer StorageBuffer %6
+          %3 = OpVariable %20 StorageBuffer
+          %4 = OpVariable %19 StorageBuffer
+         %21 = OpTypeFunction %11
+          %2 = OpFunction %11 None %21
+          %1 = OpLabel
+         %22 = OpLoad %6 %3
+         %23 = OpCopyLogical %5 %22
+               OpStore %4 %23
+               OpReturn
+               OpFunctionEnd
+
diff --git a/spirv_msl.cpp b/spirv_msl.cpp
index d81a8189..f870444e 100644
--- a/spirv_msl.cpp
+++ b/spirv_msl.cpp
@@ -4740,6 +4740,11 @@ void CompilerMSL::emit_store_statement(uint32_t lhs_expression, uint32_t rhs_exp
 
 		auto &physical_type = get<SPIRType>(physical_type_id);
 
+		string cast_addr_space = "thread";
+		auto *p_var_lhs = maybe_get_backing_variable(lhs_expression);
+		if (p_var_lhs)
+			cast_addr_space = get_type_address_space(get<SPIRType>(p_var_lhs->basetype), lhs_expression);
+
 		if (is_matrix(type))
 		{
 			const char *packed_pfx = lhs_packed_type ? "packed_" : "";
@@ -4767,7 +4772,7 @@ void CompilerMSL::emit_store_statement(uint32_t lhs_expression, uint32_t rhs_exp
 				write_type.columns = 1;
 
 				if (physical_type.columns != type.columns)
-					cast_expr = join("(device ", packed_pfx, type_to_glsl(write_type), "&)");
+					cast_expr = join("(", cast_addr_space, " ", packed_pfx, type_to_glsl(write_type), "&)");
 
 				if (rhs_transpose)
 				{
@@ -4809,7 +4814,7 @@ void CompilerMSL::emit_store_statement(uint32_t lhs_expression, uint32_t rhs_exp
 				write_type.columns = 1;
 
 				if (physical_type.vecsize != type.vecsize)
-					cast_expr = join("(device ", packed_pfx, type_to_glsl(write_type), "&)");
+					cast_expr = join("(", cast_addr_space, " ", packed_pfx, type_to_glsl(write_type), "&)");
 
 				if (rhs_transpose)
 				{
@@ -4864,7 +4869,7 @@ void CompilerMSL::emit_store_statement(uint32_t lhs_expression, uint32_t rhs_exp
 				auto column_index = lhs_expr.find_last_of('[');
 				if (column_index != string::npos)
 				{
-					statement("((device ", type_to_glsl(write_type), "*)&",
+					statement("((", cast_addr_space, " ", type_to_glsl(write_type), "*)&",
 					          lhs_expr.insert(column_index, join('[', c, ']', ")")), " = ",
 					          to_extract_component_expression(rhs_expression, c), ";");
 				}
@@ -4887,7 +4892,7 @@ void CompilerMSL::emit_store_statement(uint32_t lhs_expression, uint32_t rhs_exp
 
 			// Unpack the expression so we can store to it with a float or float2.
 			// It's still an l-value, so it's fine. Most other unpacking of expressions turn them into r-values instead.
-			lhs = join("(device ", type_to_glsl(type), "&)", enclose_expression(lhs));
+			lhs = join("(", cast_addr_space, " ", type_to_glsl(type), "&)", enclose_expression(lhs));
 			if (!optimize_read_modify_write(expression_type(rhs_expression), lhs, rhs))
 				statement(lhs, " = ", rhs, ";");
 		}