Fix crash with atomic op intrinsics (#2489)

tools/clang/lib/CodeGen/CGCall.cpp

@@ -2966,12 +2966,28 @@ void CodeGenFunction::EmitCallArg(CallArgList &args, const Expr *E,
     if (E->getObjectKind() == OK_VectorComponent) {
       if (const HLSLVectorElementExpr *VecElt = dyn_cast<HLSLVectorElementExpr>(E)) {
         LValue LV = EmitHLSLVectorElementExpr(VecElt);
-        llvm::Value *V = LV.getExtVectorAddr();
-        llvm::Constant *Elts = LV.getExtVectorElts();
-        // Only support scalar for atomic operations.
-        assert(Elts->getType()->getVectorNumElements() == 1);
-        llvm::Value *ch = Builder.CreateExtractElement(Elts, (uint64_t)0);
-        llvm::Value *Ptr = Builder.CreateGEP(V, {Builder.getInt32(0), ch});
+        llvm::Value *Ptr = nullptr;
+        if (LV.isSimple()) {
+          // Handle the special case when the vector component access
+          // is done on a scalar using .x or .r.
+          //
+          // Example 1:
+          // groupshared uint g;
+          // InterlockedAdd(g.x, 1);
+          //
+          // Example 2:
+          // RWBuffer<uint> buf;
+          // InterlockedAdd(buf[0].r, 1);
+          llvm::Value *V = LV.getAddress();
+          Ptr = Builder.CreateGEP(V, { Builder.getInt32(0) });
+        } else {
+          llvm::Value *V = LV.getExtVectorAddr();
+          llvm::Constant *Elts = LV.getExtVectorElts();
+          // Only support scalar for atomic operations.
+          assert(Elts->getType()->getVectorNumElements() == 1);
+          llvm::Value *ch = Builder.CreateExtractElement(Elts, (uint64_t)0);
+          Ptr = Builder.CreateGEP(V, { Builder.getInt32(0), ch });
+        }
         RValue RV = RValue::get(Ptr);
         return args.add(RV, type);
       } else {

tools/clang/test/HLSLFileCheck/hlsl/intrinsics/atomic/atomic.hlsl

@@ -17,6 +17,7 @@
 // CHECK: atomicCompareExchange
 // CHECK: AtomicAdd
 
+
 RWByteAddressBuffer rawBuf0 : register( u0 );
 
 struct Foo
@@ -25,6 +26,7 @@ struct Foo
   float3 b;
   uint   u;
   int2 c[4];
+  int d[4];
 };
 RWStructuredBuffer<Foo> structBuf1 : register( u1 );
 RWTexture2D<uint> rwTex2: register( u2 );
@@ -61,4 +63,20 @@ void main( uint GI : SV_GroupIndex, uint3 DTid : SV_DispatchThreadID )
     rawBuf0.InterlockedCompareStore( GI * 4, shareMem[GI], v );
     rawBuf0.InterlockedCompareExchange( GI * 4, shareMem[GI], 2, v );
     rawBuf0.InterlockedAdd( GI * 4, v );
+
+	// Special case: vector component access operation on scalars; Github issue# 2434
+	// CHECK: atomicBinOp
+	InterlockedAdd(structBuf1[DTid.z].u.x, 1);
+
+	// CHECK: atomicBinOp
+	InterlockedAdd(structBuf1[DTid.z].d[1].r, 1);
+
+	// CHECK: atomicBinOp
+	InterlockedAdd(structBuf1[DTid.z].d[1].r, 1);
+
+	// CHECK: atomicBinOp
+	InterlockedAdd(rwTex2[DTid.xy].x, 1);
+
+	// CHECK: atomicrmw
+	InterlockedAdd(shareMem[DTid.z].r, 1);
 }