[spirv] Fix ConstantBuffer/TextureBuffer assignment (#929)

Assigning a ConstantBuffer<T> to a T is indeed a copy.

tools/clang/lib/SPIRV/SPIRVEmitter.cpp

@@ -199,21 +199,6 @@ const DeclContext *isConstantTextureBufferDeclRef(const Expr *expr) {
   return nullptr;
 }
 
-/// Returns true if the given expr is loading a ConstantBuffer/TextureBuffer
-/// as a whole.
-bool isConstantTextureBufferLoad(const Expr *expr) {
-  if (!expr)
-    return false;
-
-  // The expression should be a LValueToRValue implict cast from a DeclRefExpr
-  // referencing a ConstantBuffer or TextureBuffer variable.
-  if (const auto *castExpr = dyn_cast<ImplicitCastExpr>(expr))
-    if (castExpr->getCastKind() == CK_LValueToRValue)
-      return isConstantTextureBufferDeclRef(castExpr);
-
-  return false;
-}
-
 /// Returns true if
 /// * the given expr is an DeclRefExpr referencing a kind of structured or byte
 /// buffer and it is non-alias one, or
@@ -646,8 +631,8 @@ SpirvEvalInfo SPIRVEmitter::doExpr(const Expr *expr) {
   return 0;
 }
 
-SpirvEvalInfo SPIRVEmitter::loadIfGLValue(const Expr *expr) {
-  auto info = doExpr(expr);
+SpirvEvalInfo SPIRVEmitter::loadIfGLValue(const Expr *expr,
+                                          SpirvEvalInfo info) {
 
   if (!info.isRValue()) {
     // Check whether we are trying to load an externally visible structured/byte
@@ -981,7 +966,7 @@ void SPIRVEmitter::doVarDecl(const VarDecl *decl) {
       if (const auto constId = tryToEvaluateAsConst(init))
         theBuilder.createStore(varId, constId);
       else
-        storeValue(varId, loadIfGLValue(init), decl->getType(), init);
+        storeValue(varId, loadIfGLValue(init), decl->getType());
 
       // Update counter variable associatd with local variables
       tryToAssignCounterVar(decl, init);
@@ -1439,7 +1424,7 @@ void SPIRVEmitter::doReturnStmt(const ReturnStmt *stmt) {
       // class and then return.
       const uint32_t valType = typeTranslator.translateType(retType);
       const uint32_t tempVar = theBuilder.addFnVar(valType, "temp.var.ret");
-      storeValue(tempVar, retInfo, retType, retVal);
+      storeValue(tempVar, retInfo, retType);
 
       theBuilder.createReturnValue(theBuilder.createLoad(valType, tempVar));
     } else {
@@ -1559,8 +1544,7 @@ SpirvEvalInfo SPIRVEmitter::doBinaryOperator(const BinaryOperator *expr) {
       tryToAssignCounterVar(dstDecl, expr->getRHS());
 
     return processAssignment(expr->getLHS(), loadIfGLValue(expr->getRHS()),
-                             /*isCompoundAssignment=*/false, /*lhsPtr=*/0,
-                             expr->getRHS());
+                             /*isCompoundAssignment=*/false);
   }
 
   // Try to optimize floatMxN * float and floatN * float case
@@ -1605,8 +1589,8 @@ SpirvEvalInfo SPIRVEmitter::processCall(const CallExpr *callExpr) {
   const auto numParams = callee->getNumParams();
   bool isNonStaticMemberCall = false;
 
-  llvm::SmallVector<uint32_t, 4> params; // Temporary variables
-  llvm::SmallVector<uint32_t, 4> args;   // Evaluated arguments
+  llvm::SmallVector<uint32_t, 4> params;    // Temporary variables
+  llvm::SmallVector<SpirvEvalInfo, 4> args; // Evaluated arguments
 
   if (const auto *memberCall = dyn_cast<CXXMemberCallExpr>(callExpr)) {
     isNonStaticMemberCall =
@@ -1658,16 +1642,9 @@ SpirvEvalInfo SPIRVEmitter::processCall(const CallExpr *callExpr) {
     // Update counter variable associated with function parameters
     tryToAssignCounterVar(param, arg);
 
-    if (canActAsOutParmVar(param)) {
-      // The current parameter is marked as out/inout. The argument then is
-      // essentially passed in by reference. We need to load the value
-      // explicitly here since the AST won't inject LValueToRValue implicit
-      // cast for this case.
-      const uint32_t value = theBuilder.createLoad(varType, args.back());
-      theBuilder.createStore(tempVarId, value);
-    } else {
-      theBuilder.createStore(tempVarId, args.back());
-    }
+    const auto rhsVal = loadIfGLValue(arg, args.back());
+    // Initialize the temporary variables using the contents of the arguments
+    storeValue(tempVarId, rhsVal, param->getType());
   }
 
   // Push the callee into the work queue if it is not there.
@@ -1689,7 +1666,7 @@ SpirvEvalInfo SPIRVEmitter::processCall(const CallExpr *callExpr) {
       const uint32_t index = i + isNonStaticMemberCall;
       const uint32_t typeId = typeTranslator.translateType(param->getType());
       const uint32_t value = theBuilder.createLoad(typeId, params[index]);
-      theBuilder.createStore(args[index], value);
+      storeValue(args[index], value, param->getType());
     }
   }
 
@@ -2690,13 +2667,6 @@ SPIRVEmitter::processACSBufferAppendConsume(const CXXMemberCallExpr *expr) {
 
   if (isAppend) {
     // Write out the value
-    // Note: we don't want to supply the rhsExpr here to give storeValue() hints
-    // about loading a ConstantBuffer/TextureBuffer. We want it to decompose
-    // and write each component out even if the rhs is indeed a ConstantBuffer
-    // or TextureBuffer. That is because if the source code is .Append()ing some
-    // ConstantBuffer or TextureBuffer, the intent is almost surely write out
-    // the whole struct instead of creating a temporary copy and eliminate
-    // later.
     storeValue(bufferInfo, doExpr(expr->getArg(0)), bufferElemTy);
     return 0;
   } else {
@@ -3741,8 +3711,7 @@ spv::Op SPIRVEmitter::translateOp(BinaryOperator::Opcode op, QualType type) {
 SpirvEvalInfo SPIRVEmitter::processAssignment(const Expr *lhs,
                                               const SpirvEvalInfo &rhs,
                                               const bool isCompoundAssignment,
-                                              SpirvEvalInfo lhsPtr,
-                                              const Expr *rhsExpr) {
+                                              SpirvEvalInfo lhsPtr) {
   // Assigning to vector swizzling should be handled differently.
   if (SpirvEvalInfo result = tryToAssignToVectorElements(lhs, rhs))
     return result;
@@ -3760,7 +3729,7 @@ SpirvEvalInfo SPIRVEmitter::processAssignment(const Expr *lhs,
   if (!lhsPtr)
     lhsPtr = doExpr(lhs);
 
-  storeValue(lhsPtr, rhs, lhs->getType(), rhsExpr);
+  storeValue(lhsPtr, rhs, lhs->getType());
 
   // Plain assignment returns a rvalue, while compound assignment returns
   // lvalue.
@@ -3769,7 +3738,7 @@ SpirvEvalInfo SPIRVEmitter::processAssignment(const Expr *lhs,
 
 void SPIRVEmitter::storeValue(const SpirvEvalInfo &lhsPtr,
                               const SpirvEvalInfo &rhsVal,
-                              const QualType lhsValType, const Expr *rhsExpr) {
+                              const QualType lhsValType) {
   if (typeTranslator.isScalarType(lhsValType) ||
       typeTranslator.isVectorType(lhsValType) ||
       typeTranslator.isMxNMatrix(lhsValType)) {
@@ -3779,26 +3748,35 @@ void SPIRVEmitter::storeValue(const SpirvEvalInfo &lhsPtr,
     // from rhs and directly store into lhs and avoid decomposing rhs.
     // TODO: is this optimization always correct?
     theBuilder.createStore(lhsPtr, rhsVal);
-  } else if (hlsl::IsHLSLResourceType(lhsValType) ||
-             isConstantTextureBufferLoad(rhsExpr)) {
+  } else if (TypeTranslator::isOpaqueType(lhsValType)) {
     // Resource types are represented using RecordType in the AST.
     // Handle them before the general RecordType.
     //
-    // HLSL allows to put resource types in structs, or assign to variables
-    // of resource types. These can all result in illegal SPIR-V for Vulkan.
-    // We just translate here literally and let SPIRV-Tools opt to do the
-    // legalization work.
+    // HLSL allows to put resource types that translating into SPIR-V opaque
+    // types in structs, or assign to variables of resource types. These can all
+    // result in illegal SPIR-V for Vulkan. We just translate here literally and
+    // let SPIRV-Tools opt to do the legalization work.
+    //
+    // Note: legalization specific code
+    theBuilder.createStore(lhsPtr, rhsVal);
+    needsLegalization = true;
+  } else if (TypeTranslator::isAKindOfStructuredOrByteBuffer(lhsValType)) {
+    // The rhs should be a pointer and the lhs should be a pointer-to-pointer.
+    // Directly store the pointer here and let SPIRV-Tools opt to do the clean
+    // up.
+    //
+    // Note: legalization specific code
+    theBuilder.createStore(lhsPtr, rhsVal);
+    needsLegalization = true;
+
+    // For ConstantBuffers/TextureBuffers, we decompose and assign each field
+    // recursively like normal structs using the following logic.
     //
-    // TODO: this is very hacky, especially using isConstantTextureBufferLoad()
-    // to check whether we are loading from a ConstantBuffer or TextureBuffer.
     // The frontend forbids declaring ConstantBuffer<T> or TextureBuffer<T>
     // variables as function parameters/returns/variables, but happily accepts
     // assignments/returns from ConstantBuffer<T>/TextureBuffer<T> to function
     // parameters/returns/variables of type T. And ConstantBuffer<T> is not
-    // represented differently as struct T. Those together cause a great amount
-    // of nastiness.
-    theBuilder.createStore(lhsPtr, rhsVal);
-    needsLegalization = true;
+    // represented differently as struct T.
   } else if (const auto *recordType = lhsValType->getAs<RecordType>()) {
     uint32_t index = 0;
     for (const auto *decl : recordType->getDecl()->decls()) {

tools/clang/lib/SPIRV/SPIRVEmitter.h

@@ -66,7 +66,7 @@ public:
   /// not be wrapped in ImplicitCastExpr (LValueToRValue) when appearing in
   /// HLSLVectorElementExpr since the generated HLSLVectorElementExpr itself can
   /// be lvalue or rvalue.
-  SpirvEvalInfo loadIfGLValue(const Expr *expr);
+  inline SpirvEvalInfo loadIfGLValue(const Expr *expr);
 
   /// Casts the given value from fromType to toType. fromType and toType should
   /// both be scalar or vector types of the same size.
@@ -106,6 +106,11 @@ private:
   SpirvEvalInfo doMemberExpr(const MemberExpr *expr);
   SpirvEvalInfo doUnaryOperator(const UnaryOperator *expr);
 
+  /// Overload with pre computed SpirvEvalInfo.
+  ///
+  /// The given expr will not be evaluated again.
+  SpirvEvalInfo loadIfGLValue(const Expr *expr, SpirvEvalInfo info);
+
   /// Loads the pointer of the aliased-to-variable if the given expression is a
   /// DeclRefExpr referencing an alias variable. See DeclResultIdMapper for
   /// more explanation regarding this.
@@ -127,14 +132,13 @@ private:
   /// lhs again.
   SpirvEvalInfo processAssignment(const Expr *lhs, const SpirvEvalInfo &rhs,
                                   bool isCompoundAssignment,
-                                  SpirvEvalInfo lhsPtr = 0,
-                                  const Expr *rhsExpr = nullptr);
+                                  SpirvEvalInfo lhsPtr = 0);
 
   /// Generates SPIR-V instructions to store rhsVal into lhsPtr. This will be
   /// recursive if lhsValType is a composite type. rhsExpr will be used as a
   /// reference to adjust the CodeGen if not nullptr.
   void storeValue(const SpirvEvalInfo &lhsPtr, const SpirvEvalInfo &rhsVal,
-                  QualType lhsValType, const Expr *rhsExpr = nullptr);
+                  QualType lhsValType);
 
   /// Generates the necessary instructions for conducting the given binary
   /// operation on lhs and rhs. If lhsResultId is not nullptr, the evaluated
@@ -854,6 +858,10 @@ void SPIRVEmitter::doDeclStmt(const DeclStmt *declStmt) {
     doDecl(decl);
 }
 
+SpirvEvalInfo SPIRVEmitter::loadIfGLValue(const Expr *expr) {
+  return loadIfGLValue(expr, doExpr(expr));
+}
+
 } // end namespace spirv
 } // end namespace clang
 

tools/clang/test/CodeGenSPIRV/spirv.legal.cbuffer.hlsl

@@ -1,9 +1,5 @@
 // Run: %dxc -T ps_6_0 -E main
 
-// Note: The following is invalid SPIR-V code.
-//
-// * The assignment ignores storage class (and thus layout) difference.
-
 // %type_ConstantBuffer_S is the type for myCBuffer. With layout decoration.
 // %S is the type for myASBuffer elements. With layout decoration.
 // %S_0 is the type for function local variables. Without layout decoration.
@@ -31,15 +27,22 @@ float4 doStuff(S buffer) {
 
 float4 main(in float4 pos : SV_Position) : SV_Target
 {
+// Initializing a T with a ConstantBuffer<T> is a copy
 // CHECK:      [[val:%\d+]] = OpLoad %type_ConstantBuffer_S %myCBuffer
-// CHECK-NEXT:                OpStore %buffer1 [[val]]
+// CHECK-NEXT: [[vec:%\d+]] = OpCompositeExtract %v4float [[val]] 0
+// CHECK-NEXT: [[ptr:%\d+]] = OpAccessChain %_ptr_Function_v4float %buffer1 %uint_0
+// CHECK-NEXT:                OpStore [[ptr]] [[vec]]
     S buffer1 = myCBuffer;
 
+// Assigning a ConstantBuffer<T> to a T is a copy
 // CHECK:      [[val:%\d+]] = OpLoad %type_ConstantBuffer_S %myCBuffer
-// CHECK-NEXT:                OpStore %buffer2 [[val]]
+// CHECK-NEXT: [[vec:%\d+]] = OpCompositeExtract %v4float [[val]] 0
+// CHECK-NEXT: [[ptr:%\d+]] = OpAccessChain %_ptr_Function_v4float %buffer2 %uint_0
+// CHECK-NEXT:                OpStore [[ptr]] [[vec]]
     S buffer2;
     buffer2 = myCBuffer;
 
+// We have the same struct type here
 // CHECK:      [[val:%\d+]] = OpFunctionCall %S_0 %retStuff
 // CHECK-NEXT:                OpStore %buffer3 [[val]]
     S buffer3;
@@ -53,14 +56,20 @@ float4 main(in float4 pos : SV_Position) : SV_Target
 // CHECK-NEXT:                OpStore [[adr]] [[vec]]
     myASBuffer.Append(myCBuffer);
 
+// Passing a ConstantBuffer<T> to a T parameter is a copy
 // CHECK:      [[val:%\d+]] = OpLoad %type_ConstantBuffer_S %myCBuffer
-// CHECK-NEXT:                OpStore %param_var_buffer [[val]]
+// CHECK-NEXT: [[vec:%\d+]] = OpCompositeExtract %v4float [[val]] 0
+// CHECK-NEXT: [[ptr:%\d+]] = OpAccessChain %_ptr_Function_v4float %param_var_buffer %uint_0
+// CHECK-NEXT:                OpStore [[ptr]] [[vec]]
     return doStuff(myCBuffer);
 }
 
 S retStuff() {
+// Returning a ConstantBuffer<T> as a T is a copy
 // CHECK:      [[val:%\d+]] = OpLoad %type_ConstantBuffer_S %myCBuffer
-// CHECK-NEXT:                OpStore %temp_var_ret [[val]]
+// CHECK-NEXT: [[vec:%\d+]] = OpCompositeExtract %v4float [[val]] 0
+// CHECK-NEXT: [[ptr:%\d+]] = OpAccessChain %_ptr_Function_v4float %temp_var_ret %uint_0
+// CHECK-NEXT:                OpStore [[ptr]] [[vec]]
 // CHECK-NEXT: [[ret:%\d+]] = OpLoad %S_0 %temp_var_ret
 // CHECK-NEXT:                OpReturnValue [[ret]]
     return myCBuffer;

tools/clang/test/CodeGenSPIRV/spirv.legal.tbuffer.hlsl

@@ -4,7 +4,7 @@
 //
 // * The assignment ignores storage class (and thus layout) difference.
 
-// %type_TextureBuffer_S is the type for myCBuffer. With layout decoration.
+// %type_TextureBuffer_S is the type for myTBuffer. With layout decoration.
 // %S is the type for myASBuffer elements. With layout decoration.
 // %S_0 is the type for function local variables. Without layout decoration.
 
@@ -19,8 +19,8 @@ struct S {
     float4 f;
 };
 
-// CHECK: %myCBuffer = OpVariable %_ptr_Uniform_type_TextureBuffer_S Uniform
-TextureBuffer<S>         myCBuffer;
+// CHECK: %myTBuffer = OpVariable %_ptr_Uniform_type_TextureBuffer_S Uniform
+TextureBuffer<S>          myTBuffer;
 AppendStructuredBuffer<S> myASBuffer;
 
 S retStuff();
@@ -31,15 +31,22 @@ float4 doStuff(S buffer) {
 
 float4 main(in float4 pos : SV_Position) : SV_Target
 {
-// CHECK:      [[val:%\d+]] = OpLoad %type_TextureBuffer_S %myCBuffer
-// CHECK-NEXT:                OpStore %buffer1 [[val]]
-    S buffer1 = myCBuffer;
+// Initializing a T with a TextureBuffer<T> is a copy
+// CHECK:      [[val:%\d+]] = OpLoad %type_TextureBuffer_S %myTBuffer
+// CHECK-NEXT: [[vec:%\d+]] = OpCompositeExtract %v4float [[val]] 0
+// CHECK-NEXT: [[ptr:%\d+]] = OpAccessChain %_ptr_Function_v4float %buffer1 %uint_0
+// CHECK-NEXT:                OpStore [[ptr]] [[vec]]
+    S buffer1 = myTBuffer;
 
-// CHECK:      [[val:%\d+]] = OpLoad %type_TextureBuffer_S %myCBuffer
-// CHECK-NEXT:                OpStore %buffer2 [[val]]
+// Assigning a TextureBuffer<T> to a T is a copy
+// CHECK:      [[val:%\d+]] = OpLoad %type_TextureBuffer_S %myTBuffer
+// CHECK-NEXT: [[vec:%\d+]] = OpCompositeExtract %v4float [[val]] 0
+// CHECK-NEXT: [[ptr:%\d+]] = OpAccessChain %_ptr_Function_v4float %buffer2 %uint_0
+// CHECK-NEXT:                OpStore [[ptr]] [[vec]]
     S buffer2;
-    buffer2 = myCBuffer;
+    buffer2 = myTBuffer;
 
+// We have the same struct type here
 // CHECK:      [[val:%\d+]] = OpFunctionCall %S_0 %retStuff
 // CHECK-NEXT:                OpStore %buffer3 [[val]]
     S buffer3;
@@ -47,19 +54,25 @@ float4 main(in float4 pos : SV_Position) : SV_Target
 
 // The underlying struct type has the same layout. Can write out as a whole.
 // CHECK:      [[ptr:%\d+]] = OpAccessChain %_ptr_Uniform_S %myASBuffer %uint_0 {{%\d+}}
-// CHECK-NEXT: [[val:%\d+]] = OpLoad %type_TextureBuffer_S %myCBuffer
+// CHECK-NEXT: [[val:%\d+]] = OpLoad %type_TextureBuffer_S %myTBuffer
 // CHECK-NEXT:                OpStore [[ptr]] [[val]]
-    myASBuffer.Append(myCBuffer);
+    myASBuffer.Append(myTBuffer);
 
-// CHECK:      [[val:%\d+]] = OpLoad %type_TextureBuffer_S %myCBuffer
-// CHECK-NEXT:                OpStore %param_var_buffer [[val]]
-    return doStuff(myCBuffer);
+// Passing a TextureBuffer<T> to a T parameter is a copy
+// CHECK:      [[val:%\d+]] = OpLoad %type_TextureBuffer_S %myTBuffer
+// CHECK-NEXT: [[vec:%\d+]] = OpCompositeExtract %v4float [[val]] 0
+// CHECK-NEXT: [[ptr:%\d+]] = OpAccessChain %_ptr_Function_v4float %param_var_buffer %uint_0
+// CHECK-NEXT:                OpStore [[ptr]] [[vec]]
+    return doStuff(myTBuffer);
 }
 
 S retStuff() {
-// CHECK:      [[val:%\d+]] = OpLoad %type_TextureBuffer_S %myCBuffer
-// CHECK-NEXT:                OpStore %temp_var_ret [[val]]
+// Returning a TextureBuffer<T> as a T is a copy
+// CHECK:      [[val:%\d+]] = OpLoad %type_TextureBuffer_S %myTBuffer
+// CHECK-NEXT: [[vec:%\d+]] = OpCompositeExtract %v4float [[val]] 0
+// CHECK-NEXT: [[ptr:%\d+]] = OpAccessChain %_ptr_Function_v4float %temp_var_ret %uint_0
+// CHECK-NEXT:                OpStore [[ptr]] [[vec]]
 // CHECK-NEXT: [[ret:%\d+]] = OpLoad %S_0 %temp_var_ret
 // CHECK-NEXT:                OpReturnValue [[ret]]
-    return myCBuffer;
+    return myTBuffer;
 }

tools/clang/unittests/SPIRV/CodeGenSPIRVTest.cpp

@@ -984,9 +984,7 @@ TEST_F(FileTest, SpirvLegalizationStructuredBufferCounter) {
               /*runValidation=*/false);
 }
 TEST_F(FileTest, SpirvLegalizationConstantBuffer) {
-  runFileTest("spirv.legal.cbuffer.hlsl", Expect::Success,
-              // The generated SPIR-V needs legalization.
-              /*runValidation=*/false);
+  runFileTest("spirv.legal.cbuffer.hlsl");
 }
 TEST_F(FileTest, SpirvLegalizationTextureBuffer) {
   runFileTest("spirv.legal.tbuffer.hlsl", Expect::Success,