[spirv] Fix type mismatch when loading Constant/Texture Buffer (#874)

ConstantBuffer/TextureBuffer are rerepresented the same as normal
structs upon which they are parameterized. So, using the normal
translation path via TypeTranslator, we have no way to tell whether
a struct is indeed just a normal struct or a Constant/Texture
buffer. Thus, we cannot apply the Block/BufferBlock decoration
appropriately. This causes type mismatch for OpLoad instructions.

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -370,6 +370,9 @@ uint32_t DeclResultIdMapper::createVarOfExplicitLayoutStruct(
   const uint32_t structType =
       theBuilder.getStructType(fieldTypes, typeName, fieldNames, decorations);
 
+  // Register the <type-id> for this decl
+  ctBufferPCTypeIds[decl] = structType;
+
   const auto sc = usageKind == ContextUsageKind::PushConstant
                       ? spv::StorageClass::PushConstant
                       : spv::StorageClass::Uniform;
@@ -480,6 +483,13 @@ uint32_t DeclResultIdMapper::createCounterVar(const ValueDecl *decl) {
   return counterVars[decl] = counterId;
 }
 
+uint32_t
+DeclResultIdMapper::getCTBufferPushConstantTypeId(const DeclContext *decl) {
+  const auto found = ctBufferPCTypeIds.find(decl);
+  assert(found != ctBufferPCTypeIds.end());
+  return found->second;
+}
+
 std::vector<uint32_t> DeclResultIdMapper::collectStageVars() const {
   std::vector<uint32_t> vars;
 

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -266,6 +266,19 @@ public:
   /// {RW|Append|Consume}StructuredBuffer variable.
   uint32_t getOrCreateCounterId(const ValueDecl *decl);
 
+  /// \brief Returns the <type-id> for the given cbuffer, tbuffer,
+  /// ConstantBuffer, TextureBuffer, or push constant block.
+  ///
+  /// Note: we need this method because constant/texture buffers and push
+  /// constant blocks are all represented as normal struct types upon which
+  /// they are parameterized. That is different from structured buffers,
+  /// for which we can tell they are not normal structs by investigating
+  /// the name. But for constant/texture buffers and push constant blocks,
+  /// we need to have the additional Block/BufferBlock decoration to keep
+  /// type consistent. Normal translation path for structs via TypeTranslator
+  /// won't attach Block/BufferBlock decoration.
+  uint32_t getCTBufferPushConstantTypeId(const DeclContext *decl);
+
   /// \brief Returns all defined stage (builtin/input/ouput) variables in this
   /// mapper.
   std::vector<uint32_t> collectStageVars() const;
@@ -468,6 +481,10 @@ private:
   /// counter variables
   llvm::DenseMap<const NamedDecl *, uint32_t> counterVars;
 
+  /// Mapping from cbuffer/tbuffer/ConstantBuffer/TextureBufer/push-constant
+  /// to the <type-id>
+  llvm::DenseMap<const DeclContext *, uint32_t> ctBufferPCTypeIds;
+
 public:
   /// The gl_PerVertex structs for both input and output
   GlPerVertex glPerVertex;

tools/clang/lib/SPIRV/SPIRVEmitter.cpp

@@ -168,6 +168,37 @@ const Expr *isStructuredBufferLoad(const Expr *expr, const Expr **index) {
   return nullptr;
 }
 
+/// Returns the referenced variable's DeclContext if the given expr is
+/// a DeclRefExpr referencing a ConstantBuffer/TextureBuffer. Otherwise,
+/// returns nullptr.
+const DeclContext *isConstantTextureBufferDeclRef(const Expr *expr) {
+  if (const auto *declRefExpr = dyn_cast<DeclRefExpr>(expr->IgnoreParenCasts()))
+    if (const auto *varDecl = dyn_cast<VarDecl>(declRefExpr->getFoundDecl()))
+      if (const auto *bufferDecl =
+              dyn_cast<HLSLBufferDecl>(varDecl->getDeclContext()))
+        // Make sure we are not returning true for VarDecls inside
+        // cbuffer/tbuffer.
+        if (bufferDecl->isConstantBufferView())
+          return varDecl->getType()->getAs<RecordType>()->getDecl();
+
+  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;
+}
+
 bool spirvToolsLegalize(std::vector<uint32_t> *module, std::string *messages) {
   spvtools::Optimizer optimizer(SPV_ENV_VULKAN_1_0);
 
@@ -562,8 +593,16 @@ SpirvEvalInfo SPIRVEmitter::loadIfGLValue(const Expr *expr) {
   auto info = doExpr(expr);
 
   if (!info.isRValue()) {
-    const uint32_t valType =
-        typeTranslator.translateType(expr->getType(), info.getLayoutRule());
+    uint32_t valType = 0;
+    // TODO: Ouch. Very hacky. We need special path to get the value type if
+    // we are loading a whole ConstantBuffer/TextureBuffer since the normal
+    // type translation path won't work.
+    if (const auto *declContext = isConstantTextureBufferDeclRef(expr)) {
+      valType = declIdMapper.getCTBufferPushConstantTypeId(declContext);
+    } else {
+      valType =
+          typeTranslator.translateType(expr->getType(), info.getLayoutRule());
+    }
     info.setResultId(theBuilder.createLoad(valType, info)).setRValue();
   }
 
@@ -840,11 +879,11 @@ void SPIRVEmitter::doVarDecl(const VarDecl *decl) {
 
     }
     // Function local variables. Just emit OpStore at the current insert point.
-    else if (decl->hasInit()) {
-      if (const auto constId = tryToEvaluateAsConst(decl->getInit()))
+    else if (const Expr *init = decl->getInit()) {
+      if (const auto constId = tryToEvaluateAsConst(init))
         theBuilder.createStore(varId, constId);
       else
-        storeValue(varId, loadIfGLValue(decl->getInit()), decl->getType());
+        storeValue(varId, loadIfGLValue(init), decl->getType(), init);
     }
   } else {
     varId = declIdMapper.createExternVar(decl);
@@ -1298,7 +1337,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);
+      storeValue(tempVar, retInfo, retType, retVal);
 
       theBuilder.createReturnValue(theBuilder.createLoad(valType, tempVar));
     } else {
@@ -1415,7 +1454,8 @@ SpirvEvalInfo SPIRVEmitter::doBinaryOperator(const BinaryOperator *expr) {
   // For other binary operations, we need to evaluate lhs before rhs.
   if (opcode == BO_Assign) {
     return processAssignment(expr->getLHS(), loadIfGLValue(expr->getRHS()),
-                             false);
+                             /*isCompoundAssignment=*/false, /*lhsPtr=*/0,
+                             expr->getRHS());
   }
 
   // Try to optimize floatMxN * float and floatN * float case
@@ -2581,6 +2621,13 @@ 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 {
@@ -3390,7 +3437,8 @@ case BO_##kind : {                                                             \
 SpirvEvalInfo SPIRVEmitter::processAssignment(const Expr *lhs,
                                               const SpirvEvalInfo &rhs,
                                               const bool isCompoundAssignment,
-                                              SpirvEvalInfo lhsPtr) {
+                                              SpirvEvalInfo lhsPtr,
+                                              const Expr *rhsExpr) {
   // Assigning to vector swizzling should be handled differently.
   if (SpirvEvalInfo result = tryToAssignToVectorElements(lhs, rhs))
     return result;
@@ -3408,7 +3456,7 @@ SpirvEvalInfo SPIRVEmitter::processAssignment(const Expr *lhs,
   if (!lhsPtr)
     lhsPtr = doExpr(lhs);
 
-  storeValue(lhsPtr, rhs, lhs->getType());
+  storeValue(lhsPtr, rhs, lhs->getType(), rhsExpr);
 
   // Plain assignment returns a rvalue, while compound assignment returns
   // lvalue.
@@ -3417,17 +3465,18 @@ SpirvEvalInfo SPIRVEmitter::processAssignment(const Expr *lhs,
 
 void SPIRVEmitter::storeValue(const SpirvEvalInfo &lhsPtr,
                               const SpirvEvalInfo &rhsVal,
-                              const QualType valType) {
-  if (typeTranslator.isScalarType(valType) ||
-      typeTranslator.isVectorType(valType) ||
-      typeTranslator.isMxNMatrix(valType)) {
+                              const QualType lhsValType, const Expr *rhsExpr) {
+  if (typeTranslator.isScalarType(lhsValType) ||
+      typeTranslator.isVectorType(lhsValType) ||
+      typeTranslator.isMxNMatrix(lhsValType)) {
     theBuilder.createStore(lhsPtr, rhsVal);
   } else if (lhsPtr.getLayoutRule() == rhsVal.getLayoutRule()) {
     // If lhs and rhs has the same memory layout, we should be safe to load
     // 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(valType)) {
+  } else if (hlsl::IsHLSLResourceType(lhsValType) ||
+             isConstantTextureBufferLoad(rhsExpr)) {
     // Resource types are represented using RecordType in the AST.
     // Handle them before the general RecordType.
     //
@@ -3435,9 +3484,18 @@ void SPIRVEmitter::storeValue(const SpirvEvalInfo &lhsPtr,
     // 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.
+    //
+    // 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;
-  } else if (const auto *recordType = valType->getAs<RecordType>()) {
+  } else if (const auto *recordType = lhsValType->getAs<RecordType>()) {
     uint32_t index = 0;
     for (const auto *decl : recordType->getDecl()->decls()) {
       // Ignore implicit generated struct declarations/constructors/destructors.
@@ -3462,7 +3520,7 @@ void SPIRVEmitter::storeValue(const SpirvEvalInfo &lhsPtr,
       ++index;
     }
   } else if (const auto *arrayType =
-                 astContext.getAsConstantArrayType(valType)) {
+                 astContext.getAsConstantArrayType(lhsValType)) {
     const auto elemType = arrayType->getElementType();
     // TODO: handle extra large array size?
     const auto size =
@@ -3483,7 +3541,7 @@ void SPIRVEmitter::storeValue(const SpirvEvalInfo &lhsPtr,
                  rhsVal.substResultId(subRhsVal), elemType);
     }
   } else {
-    emitError("storing value of type %0 unimplemented", {}) << valType;
+    emitError("storing value of type %0 unimplemented", {}) << lhsValType;
   }
 }
 

tools/clang/lib/SPIRV/SPIRVEmitter.h

@@ -120,12 +120,14 @@ private:
   /// lhs again.
   SpirvEvalInfo processAssignment(const Expr *lhs, const SpirvEvalInfo &rhs,
                                   bool isCompoundAssignment,
-                                  SpirvEvalInfo lhsPtr = 0);
+                                  SpirvEvalInfo lhsPtr = 0,
+                                  const Expr *rhsExpr = nullptr);
 
   /// Generates SPIR-V instructions to store rhsVal into lhsPtr. This will be
-  /// recursive if valType is a composite type.
+  /// 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 valType);
+                  QualType lhsValType, const Expr *rhsExpr = nullptr);
 
   /// Generates the necessary instructions for conducting the given binary
   /// operation on lhs and rhs. If lhsResultId is not nullptr, the evaluated

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

@@ -0,0 +1,67 @@
+// 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.
+
+// CHECK:     OpMemberDecorate %type_ConstantBuffer_S 0 Offset 0
+// CHECK:     OpMemberDecorate %S 0 Offset 0
+// CHECK-NOT: OpMemberDecorate %S_0 0 Offset 0
+
+// CHECK:     %type_ConstantBuffer_S = OpTypeStruct %v4float
+// CHECK:     %S = OpTypeStruct %v4float
+// CHECK:     %S_0 = OpTypeStruct %v4float
+struct S {
+    float4 f;
+};
+
+// CHECK: %myCBuffer = OpVariable %_ptr_Uniform_type_ConstantBuffer_S Uniform
+ConstantBuffer<S>         myCBuffer;
+AppendStructuredBuffer<S> myASBuffer;
+
+S retStuff();
+
+float4 doStuff(S buffer) {
+    return buffer.f;
+}
+
+float4 main(in float4 pos : SV_Position) : SV_Target
+{
+// CHECK:      [[val:%\d+]] = OpLoad %type_ConstantBuffer_S %myCBuffer
+// CHECK-NEXT:                OpStore %buffer1 [[val]]
+    S buffer1 = myCBuffer;
+
+// CHECK:      [[val:%\d+]] = OpLoad %type_ConstantBuffer_S %myCBuffer
+// CHECK-NEXT:                OpStore %buffer2 [[val]]
+    S buffer2;
+    buffer2 = myCBuffer;
+
+// CHECK:      [[val:%\d+]] = OpFunctionCall %S_0 %retStuff
+// CHECK-NEXT:                OpStore %buffer3 [[val]]
+    S buffer3;
+    buffer3 = retStuff();
+
+// Write out each component recursively
+// CHECK:      [[ptr:%\d+]] = OpAccessChain %_ptr_Uniform_S %myASBuffer %uint_0 {{%\d+}}
+// CHECK-NEXT: [[val:%\d+]] = OpLoad %type_ConstantBuffer_S %myCBuffer
+// CHECK-NEXT: [[vec:%\d+]] = OpCompositeExtract %v4float [[val]] 0
+// CHECK-NEXT: [[adr:%\d+]] = OpAccessChain %_ptr_Uniform_v4float [[ptr]] %uint_0
+// CHECK-NEXT:                OpStore [[adr]] [[vec]]
+    myASBuffer.Append(myCBuffer);
+
+// CHECK:      [[val:%\d+]] = OpLoad %type_ConstantBuffer_S %myCBuffer
+// CHECK-NEXT:                OpStore %param_var_buffer [[val]]
+    return doStuff(myCBuffer);
+}
+
+S retStuff() {
+// CHECK:      [[val:%\d+]] = OpLoad %type_ConstantBuffer_S %myCBuffer
+// CHECK-NEXT:                OpStore %temp_var_ret [[val]]
+// CHECK-NEXT: [[ret:%\d+]] = OpLoad %S_0 %temp_var_ret
+// CHECK-NEXT:                OpReturnValue [[ret]]
+    return myCBuffer;
+}

tools/clang/test/CodeGenSPIRV/binary-op.assign.resource.hlsl → tools/clang/test/CodeGenSPIRV/spirv.legal.sbuffer.hlsl

@@ -18,9 +18,6 @@ struct ResourceBundle {
     RWBuffer<float3>           rwbuf;
     Texture2D<float>           tex2d;
     RWTexture2D<float4>        rwtex2d;
-    // The frontend does not allow ConstantBuffers and TextureBuffers to be in a struct.
-    //ConstantBuffer<T>          cbuf;
-    //TextureBuffer<T>           tbuf;
     StructuredBuffer<T>        sbuf;
     RWStructuredBuffer<T>      rwsbuf;
     AppendStructuredBuffer<T>  asbuf;
@@ -34,8 +31,6 @@ Buffer<float3>             g_buf;
 RWBuffer<float3>           g_rwbuf;
 Texture2D<float>           g_tex2d;
 RWTexture2D<float4>        g_rwtex2d;
-//ConstantBuffer<T>          g_cbuf;
-//TextureBuffer<T>           g_tbuf;
 StructuredBuffer<T>        g_sbuf;
 RWStructuredBuffer<T>      g_rwsbuf;
 AppendStructuredBuffer<T>  g_asbuf;
@@ -71,9 +66,6 @@ void main() {
 // CHECK-NEXT:                OpStore [[ptr]] [[val]]
     b.rwtex2d = g_rwtex2d;
 
-    //b.cbuf    = g_cbuf;
-    //b.tbuf    = g_tbuf;
-
 // CHECK-NEXT: [[val:%\d+]] = OpLoad %type_StructuredBuffer_T %g_sbuf
 // CHECK-NEXT: [[ptr:%\d+]] = OpAccessChain %_ptr_Function_type_StructuredBuffer_T_0 %b %int_5
 // CHECK-NEXT:                OpStore [[ptr]] [[val]]

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

@@ -0,0 +1,65 @@
+// 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_TextureBuffer_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.
+
+// CHECK:     OpMemberDecorate %type_TextureBuffer_S 0 Offset 0
+// CHECK:     OpMemberDecorate %S 0 Offset 0
+// CHECK-NOT: OpMemberDecorate %S_0 0 Offset 0
+
+// CHECK:     %type_TextureBuffer_S = OpTypeStruct %v4float
+// CHECK:     %S = OpTypeStruct %v4float
+// CHECK:     %S_0 = OpTypeStruct %v4float
+struct S {
+    float4 f;
+};
+
+// CHECK: %myCBuffer = OpVariable %_ptr_Uniform_type_TextureBuffer_S Uniform
+TextureBuffer<S>         myCBuffer;
+AppendStructuredBuffer<S> myASBuffer;
+
+S retStuff();
+
+float4 doStuff(S buffer) {
+    return buffer.f;
+}
+
+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;
+
+// CHECK:      [[val:%\d+]] = OpLoad %type_TextureBuffer_S %myCBuffer
+// CHECK-NEXT:                OpStore %buffer2 [[val]]
+    S buffer2;
+    buffer2 = myCBuffer;
+
+// CHECK:      [[val:%\d+]] = OpFunctionCall %S_0 %retStuff
+// CHECK-NEXT:                OpStore %buffer3 [[val]]
+    S buffer3;
+    buffer3 = retStuff();
+
+// 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:                OpStore [[ptr]] [[val]]
+    myASBuffer.Append(myCBuffer);
+
+// CHECK:      [[val:%\d+]] = OpLoad %type_TextureBuffer_S %myCBuffer
+// CHECK-NEXT:                OpStore %param_var_buffer [[val]]
+    return doStuff(myCBuffer);
+}
+
+S retStuff() {
+// CHECK:      [[val:%\d+]] = OpLoad %type_TextureBuffer_S %myCBuffer
+// CHECK-NEXT:                OpStore %temp_var_ret [[val]]
+// CHECK-NEXT: [[ret:%\d+]] = OpLoad %S_0 %temp_var_ret
+// CHECK-NEXT:                OpReturnValue [[ret]]
+    return myCBuffer;
+}

tools/clang/unittests/SPIRV/CodeGenSPIRVTest.cpp

@@ -139,11 +139,6 @@ TEST_F(FileTest, BinaryOpAssign) { runFileTest("binary-op.assign.hlsl"); }
 TEST_F(FileTest, BinaryOpAssignComposite) {
   runFileTest("binary-op.assign.composite.hlsl");
 }
-TEST_F(FileTest, BinaryOpAssignResource) {
-  runFileTest("binary-op.assign.resource.hlsl", Expect::Success,
-              // It generates invalid SPIR-V code right now.
-              /*runValidation=*/false);
-}
 
 // For comma binary operator
 TEST_F(FileTest, BinaryOpComma) { runFileTest("binary-op.comma.hlsl"); }
@@ -959,6 +954,22 @@ TEST_F(FileTest, SpirvInterpolationError) {
   runFileTest("spirv.interpolation.error.hlsl", FileTest::Expect::Failure);
 }
 
+TEST_F(FileTest, SpirvLegalizationStructuredBuffer) {
+  runFileTest("spirv.legal.sbuffer.hlsl", Expect::Success,
+              // The generated SPIR-V needs legalization.
+              /*runValidation=*/false);
+}
+TEST_F(FileTest, SpirvLegalizationConstantBuffer) {
+  runFileTest("spirv.legal.cbuffer.hlsl", Expect::Success,
+              // The generated SPIR-V needs legalization.
+              /*runValidation=*/false);
+}
+TEST_F(FileTest, SpirvLegalizationTextureBuffer) {
+  runFileTest("spirv.legal.tbuffer.hlsl", Expect::Success,
+              // The generated SPIR-V needs legalization.
+              /*runValidation=*/false);
+}
+
 TEST_F(FileTest, VulkanAttributeErrors) {
   runFileTest("vk.attribute.error.hlsl", FileTest::Expect::Failure);
 }