[Mlir-commits] [mlir] [mlir][spirv] Truncate Literal String size at max number words (PR #142916)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-spirv

Author: Davide Grohmann (davidegrohmann)

<details>
<summary>Changes</summary>

If not truncated the SPIRV serialization would not fail but instead produce an invalid SPIR-V module.

---
Full diff: https://github.com/llvm/llvm-project/pull/142916.diff


2 Files Affected:

- (modified) mlir/include/mlir/Target/SPIRV/SPIRVBinaryUtils.h (+6) 
- (modified) mlir/lib/Target/SPIRV/SPIRVBinaryUtils.cpp (+15-1) 


``````````diff
diff --git a/mlir/include/mlir/Target/SPIRV/SPIRVBinaryUtils.h b/mlir/include/mlir/Target/SPIRV/SPIRVBinaryUtils.h
index e46a576f1d48e..d3847ae3d3bb2 100644
--- a/mlir/include/mlir/Target/SPIRV/SPIRVBinaryUtils.h
+++ b/mlir/include/mlir/Target/SPIRV/SPIRVBinaryUtils.h
@@ -30,6 +30,12 @@ constexpr uint32_t kMagicNumber = 0x07230203;
 /// The serializer tool ID registered to the Khronos Group
 constexpr uint32_t kGeneratorNumber = 22;
 
+// Max number of words
+constexpr uint32_t kMaxWordCount = 65535;
+
+// Max number of words for literal
+constexpr uint32_t kMaxLiteralWordCount = kMaxWordCount - 3;
+
 /// Appends a SPRI-V module header to `header` with the given `version` and
 /// `idBound`.
 void appendModuleHeader(SmallVectorImpl<uint32_t> &header,
diff --git a/mlir/lib/Target/SPIRV/SPIRVBinaryUtils.cpp b/mlir/lib/Target/SPIRV/SPIRVBinaryUtils.cpp
index 31205d8f408f1..4d4d67a012ae1 100644
--- a/mlir/lib/Target/SPIRV/SPIRVBinaryUtils.cpp
+++ b/mlir/lib/Target/SPIRV/SPIRVBinaryUtils.cpp
@@ -13,6 +13,9 @@
 #include "mlir/Target/SPIRV/SPIRVBinaryUtils.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "llvm/Config/llvm-config.h" // for LLVM_VERSION_MAJOR
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "spirv-binary-utils"
 
 using namespace mlir;
 
@@ -68,7 +71,18 @@ void spirv::encodeStringLiteralInto(SmallVectorImpl<uint32_t> &binary,
                                     StringRef literal) {
   // We need to encode the literal and the null termination.
   auto encodingSize = literal.size() / 4 + 1;
+  auto sizeOfDataToCopy = literal.size();
+  if (encodingSize >= kMaxLiteralWordCount) {
+    // reserve one word for the null termination
+    encodingSize = kMaxLiteralWordCount - 1;
+    // do not override the last word (null termination) when copying
+    sizeOfDataToCopy = (encodingSize - 1) * 4;
+    LLVM_DEBUG(llvm::dbgs() << "Truncating string literal to max size ("
+                            << std::to_string(kMaxLiteralWordCount - 1)
+                            << "): " << literal << "\n");
+  }
   auto bufferStartSize = binary.size();
   binary.resize(bufferStartSize + encodingSize, 0);
-  std::memcpy(binary.data() + bufferStartSize, literal.data(), literal.size());
+  std::memcpy(binary.data() + bufferStartSize, literal.data(),
+              sizeOfDataToCopy);
 }

``````````

</details>


https://github.com/llvm/llvm-project/pull/142916