Fix DispatchMesh with groupshared payload.

lib/Transforms/Scalar/ScalarReplAggregatesHLSL.cpp

@@ -2588,6 +2588,13 @@ void SROA_Helper::RewriteCall(CallInst *CI) {
         RewriteCallArg(CI, HLOperandIndex::kCallShaderPayloadOpIdx,
                        /*bIn*/ true, /*bOut*/ true);
       } break;
+      case IntrinsicOp::IOP_DispatchMesh: {
+        if (OldVal ==
+            CI->getArgOperand(HLOperandIndex::kDispatchMeshOpPayload)) {
+          RewriteCallArg(CI, HLOperandIndex::kDispatchMeshOpPayload,
+                         /*bIn*/ true, /*bOut*/ false);
+        }
+      } break;
       case IntrinsicOp::MOP_TraceRayInline: {
         if (OldVal ==
             CI->getArgOperand(HLOperandIndex::kTraceRayInlineRayDescOpIdx)) {

tools/clang/lib/CodeGen/CGCall.cpp

@@ -3379,6 +3379,13 @@ RValue CodeGenFunction::EmitCall(const CGFunctionInfo &CallInfo,
           EmitAggregateCopy(AI, Addr, I->Ty, RV.isVolatileQualified());
         } else {
           // Skip the extra memcpy call.
+          // HLSL Change Starts
+          // Generate AddrSpaceCast for shared memory.
+          if (RVAddrSpace != ArgAddrSpace) {
+            Addr = Builder.CreateAddrSpaceCast(
+                Addr, IRFuncTy->getParamType(FirstIRArg));
+          }
+          // HLSL Change Ends
           IRCallArgs[FirstIRArg] = Addr;
         }
       }

tools/clang/test/HLSLFileCheck/shader_targets/mesh/as-groupshared-payload.hlsl

@@ -0,0 +1,17 @@
+// RUN: %dxc -E amplification -T as_6_5 %s | FileCheck %s
+
+// CHECK: define void @amplification
+
+struct MeshPayload
+{
+  uint data[4];
+};
+
+groupshared MeshPayload pld;
+
+[numthreads(4,1,1)]
+void amplification(uint gtid : SV_GroupIndex)
+{
+  pld.data[gtid] = gtid;
+  DispatchMesh(1,1,1,pld);
+}