1 files changed, 8 insertions, 3 deletions
diff --git a/llvm/lib/MC/DXContainerRootSignature.cpp b/llvm/lib/MC/DXContainerRootSignature.cpp
index b9ebb7a9e789..2338370d8438 100644
--- a/llvm/lib/MC/DXContainerRootSignature.cpp
+++ b/llvm/lib/MC/DXContainerRootSignature.cpp
@@ -32,10 +32,12 @@ static uint32_t rewriteOffsetToCurrentByte(raw_svector_ostream &Stream,
 
 size_t RootSignatureDesc::getSize() const {
   uint32_t StaticSamplersOffset = computeStaticSamplersOffset();
-  size_t StaticSamplersSize =
-      StaticSamplers.size() * sizeof(dxbc::RTS0::v1::StaticSampler);
+  size_t StaticSamplersSize = sizeof(dxbc::RTS0::v1::StaticSampler);
+  if (Version > 2)
+    StaticSamplersSize = sizeof(dxbc::RTS0::v3::StaticSampler);
 
-  return size_t(StaticSamplersOffset) + StaticSamplersSize;
+  return size_t(StaticSamplersOffset) +
+         (StaticSamplersSize * StaticSamplers.size());
 }
 
 uint32_t RootSignatureDesc::computeRootParametersOffset() const {
@@ -171,6 +173,9 @@ void RootSignatureDesc::write(raw_ostream &OS) const {
     support::endian::write(BOS, S.ShaderRegister, llvm::endianness::little);
     support::endian::write(BOS, S.RegisterSpace, llvm::endianness::little);
     support::endian::write(BOS, S.ShaderVisibility, llvm::endianness::little);
+
+    if (Version > 2)
+      support::endian::write(BOS, S.Flags, llvm::endianness::little);
   }
   assert(Storage.size() == getSize());
   OS.write(Storage.data(), Storage.size());