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

[Davide Grohmann]

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

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

>From 54172aac4f2be65ec10fd484b47d9600c208d9e3 Mon Sep 17 00:00:00 2001
From: Davide Grohmann <davide.grohmann at arm.com>
Date: Wed, 29 Jan 2025 12:56:47 +0100
Subject: [PATCH] [mlir][spirv] Truncate Literal String size at max number
 words

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

Change-Id: I54c4e54d6ad081861b524d4ae236a1e5080b88c4
Signed-off-by: Davide Grohmann <davide.grohmann at arm.com>
---
 .../include/mlir/Target/SPIRV/SPIRVBinaryUtils.h |  6 ++++++
 mlir/lib/Target/SPIRV/SPIRVBinaryUtils.cpp       | 16 +++++++++++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

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);
 }