[spirv] Fix vector swizzling on RWBuffer/RWTexture elements (#979)

For such vector swizzling, we cannot get the pointer and use OpStore
like normal vectors. Instead we should use OpImageWrite.

tools/clang/lib/SPIRV/SPIRVEmitter.cpp

@@ -4426,11 +4426,14 @@ SPIRVEmitter::tryToAssignToVectorElements(const Expr *lhs,
   if (!lhsExpr)
     return 0;
 
-  if (!isVectorShuffle(lhs)) {
-    // No vector shuffle needed to be generated for this assignment.
-    // Should fall back to the normal handling of assignment.
-    return 0;
-  }
+  // Special case for <scalar-value>.x, which will have an AST of
+  // HLSLVectorElementExpr whose base is an ImplicitCastExpr
+  // (CK_HLSLVectorSplat). We just need to assign to <scalar-value>
+  // for such case.
+  if (const auto *baseCast = dyn_cast<CastExpr>(lhsExpr->getBase()))
+    if (baseCast->getCastKind() == CastKind::CK_HLSLVectorSplat &&
+        hlsl::GetHLSLVecSize(baseCast->getType()) == 1)
+      return processAssignment(baseCast->getSubExpr(), rhs, false);
 
   const Expr *base = nullptr;
   hlsl::VectorMemberAccessPositions accessor;
@@ -4438,12 +4441,50 @@ SPIRVEmitter::tryToAssignToVectorElements(const Expr *lhs,
 
   const QualType baseType = base->getType();
   assert(hlsl::IsHLSLVecType(baseType));
-  const auto baseSizse = hlsl::GetHLSLVecSize(baseType);
+  const uint32_t baseTypeId = typeTranslator.translateType(baseType);
+  const auto baseSize = hlsl::GetHLSLVecSize(baseType);
+  const auto accessorSize = accessor.Count;
+  // Whether selecting the whole original vector
+  bool isSelectOrigin = accessorSize == baseSize;
+
+  // Assigning to one component
+  if (accessorSize == 1) {
+    if (isBufferTextureIndexing(dyn_cast_or_null<CXXOperatorCallExpr>(base))) {
+      // Assigning to one component of a RWBuffer/RWTexture element
+      // We need to use OpImageWrite here.
+      // Compose the new vector value first
+      const uint32_t oldVec = doExpr(base);
+      const uint32_t newVec = theBuilder.createCompositeInsert(
+          baseTypeId, oldVec, {accessor.Swz0}, rhs);
+      const auto result = tryToAssignToRWBufferRWTexture(base, newVec);
+      assert(result); // Definitely RWBuffer/RWTexture assignment
+      return rhs;     // TODO: incorrect for compound assignments
+    } else {
+      // Assigning to one normal vector component. Nothing special, just fall
+      // back to the normal CodeGen path.
+      return 0;
+    }
+  }
+
+  if (isSelectOrigin) {
+    for (uint32_t i = 0; i < accessorSize; ++i) {
+      uint32_t position;
+      accessor.GetPosition(i, &position);
+      if (position != i)
+        isSelectOrigin = false;
+    }
+  }
+
+  // Assigning to the original vector
+  if (isSelectOrigin) {
+    // Ignore this HLSLVectorElementExpr and dispatch to base
+    return processAssignment(base, rhs, false);
+  }
 
   llvm::SmallVector<uint32_t, 4> selectors;
-  selectors.resize(baseSizse);
+  selectors.resize(baseSize);
   // Assume we are selecting all original elements first.
-  for (uint32_t i = 0; i < baseSizse; ++i) {
+  for (uint32_t i = 0; i < baseSize; ++i) {
     selectors[i] = i;
   }
 
@@ -4453,16 +4494,17 @@ SPIRVEmitter::tryToAssignToVectorElements(const Expr *lhs,
   for (uint32_t i = 0; i < accessor.Count; ++i) {
     uint32_t position;
     accessor.GetPosition(i, &position);
-    selectors[position] = baseSizse + i;
+    selectors[position] = baseSize + i;
   }
 
-  const uint32_t baseTypeId = typeTranslator.translateType(baseType);
-  const uint32_t vec1 = doExpr(base);
-  const uint32_t vec1Val = theBuilder.createLoad(baseTypeId, vec1);
+  const auto vec1 = doExpr(base);
+  const uint32_t vec1Val =
+      vec1.isRValue() ? vec1 : theBuilder.createLoad(baseTypeId, vec1);
   const uint32_t shuffle =
       theBuilder.createVectorShuffle(baseTypeId, vec1Val, rhs, selectors);
 
-  theBuilder.createStore(vec1, shuffle);
+  if (!tryToAssignToRWBufferRWTexture(base, shuffle))
+    theBuilder.createStore(vec1, shuffle);
 
   // TODO: OK, this return value is incorrect for compound assignments, for
   // which cases we should return lvalues. Should at least emit errors if

tools/clang/lib/SPIRV/SPIRVEmitter.h

@@ -220,10 +220,6 @@ private:
   /// Tries to emit instructions for assigning to the given vector element
   /// accessing expression. Returns 0 if the trial fails and no instructions
   /// are generated.
-  ///
-  /// This method handles the cases that we are writing to neither one element
-  /// or all elements in their original order. For other cases, 0 will be
-  /// returned and the normal assignment process should be used.
   SpirvEvalInfo tryToAssignToVectorElements(const Expr *lhs,
                                             const SpirvEvalInfo &rhs);
 

tools/clang/test/CodeGenSPIRV/binary-op.assign.vector.hlsl

@@ -0,0 +1,53 @@
+// Run: %dxc -T vs_6_0 -E main
+
+RWBuffer<uint>     MyRWBuffer;
+RWTexture2D<int>   MyRWTexture;
+RWBuffer<uint4>    MyRWBuffer4;
+RWTexture3D<uint3> MyRWTexture3;
+
+// CHECK:      [[c25:%\d+]] = OpConstantComposite %v3uint %uint_25 %uint_25 %uint_25
+// CHECK:      [[c26:%\d+]] = OpConstantComposite %v4uint %uint_26 %uint_26 %uint_26 %uint_26
+// CHECK:      [[c28:%\d+]] = OpConstantComposite %v2uint %uint_28 %uint_28
+
+void main() {
+    // <scalar-value>.x
+// CHECK:      [[buf:%\d+]] = OpLoad %type_buffer_image %MyRWBuffer
+// CHECK-NEXT:                OpImageWrite [[buf]] %uint_1 %uint_15
+    MyRWBuffer[1].x = 15;
+
+    // <scalar-value>.x
+// CHECK:      [[tex:%\d+]] = OpLoad %type_2d_image %MyRWTexture
+// CHECK-NEXT:                OpImageWrite [[tex]] {{%\d+}} %int_16
+    MyRWTexture[uint2(2, 3)].x = 16;
+
+    // Out-of-order swizzling
+// CHECK:      [[buf:%\d+]] = OpLoad %type_buffer_image_0 %MyRWBuffer4
+// CHECK-NEXT: [[old:%\d+]] = OpImageRead %v4uint [[buf]] %uint_4 None
+// CHECK-NEXT: [[new:%\d+]] = OpVectorShuffle %v4uint [[old]] [[c25]] 6 1 4 5
+// CHECK-NEXT: [[buf:%\d+]] = OpLoad %type_buffer_image_0 %MyRWBuffer4
+// CHECK-NEXT:                OpImageWrite [[buf]] %uint_4 [[new]]
+    MyRWBuffer4[4].zwx = 25;
+
+    // Swizzling resulting in the original vector
+// CHECK:      [[buf:%\d+]] = OpLoad %type_buffer_image_0 %MyRWBuffer4
+// CHECK-NEXT:                OpImageWrite [[buf]] %uint_4 [[c26]]
+    MyRWBuffer4[4].xyzw = 26;
+
+    // Selecting one element
+// CHECK:      [[tex:%\d+]] = OpLoad %type_3d_image %MyRWTexture3
+// CHECK-NEXT: [[old:%\d+]] = OpImageRead %v4uint [[tex]] {{%\d+}} None
+// CHECK-NEXT:  [[v3:%\d+]] = OpVectorShuffle %v3uint [[old]] [[old]] 0 1 2
+// CHECK-NEXT: [[new:%\d+]] = OpCompositeInsert %v3uint %uint_27 [[v3]] 1
+// CHECK-NEXT: [[tex:%\d+]] = OpLoad %type_3d_image %MyRWTexture3
+// CHECK-NEXT:                OpImageWrite [[tex]] {{%\d+}} [[new]]
+    MyRWTexture3[uint3(5, 6, 7)].y = 27;
+
+    // In-order swizzling
+// CHECK:      [[tex:%\d+]] = OpLoad %type_3d_image %MyRWTexture3
+// CHECK-NEXT: [[old:%\d+]] = OpImageRead %v4uint [[tex]] {{%\d+}} None
+// CHECK-NEXT:  [[v3:%\d+]] = OpVectorShuffle %v3uint [[old]] [[old]] 0 1 2
+// CHECK-NEXT: [[new:%\d+]] = OpVectorShuffle %v3uint [[v3]] [[c28]] 3 4 2
+// CHECK-NEXT: [[tex:%\d+]] = OpLoad %type_3d_image %MyRWTexture3
+// CHECK-NEXT:                OpImageWrite [[tex]] {{%\d+}} [[new]]
+    MyRWTexture3[uint3(8, 9, 10)].xy = 28;
+}

tools/clang/unittests/SPIRV/CodeGenSPIRVTest.cpp

@@ -156,6 +156,9 @@ TEST_F(FileTest, UnaryOpLogicalNot) {
 
 // For assignments
 TEST_F(FileTest, BinaryOpAssign) { runFileTest("binary-op.assign.hlsl"); }
+TEST_F(FileTest, BinaryOpAssignVector) {
+  runFileTest("binary-op.assign.vector.hlsl");
+}
 TEST_F(FileTest, BinaryOpAssignComposite) {
   runFileTest("binary-op.assign.composite.hlsl");
 }