[mlir][spirv] Truncate Literal String size at max number words

[davidegrohmann]

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

[github-actions]

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-spirv

Author: Davide Grohmann (davidegrohmann)

Changes

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

[IgWod-IMG]

Thank you for your contribution!

If possible, could you please try to add a test to exercise this change? I believe a target test should do the trick (see mlir/test/Target/SPIRV). Those tests serialize given MLIR into SPIR-V and then deserialize the SPIR-V back to MLIR. So, I would hope you could see a truncation after the roundtrip.

[davidegrohmann]

Thank you for your contribution!

If possible, could you please try to add a test to exercise this change? I believe a target test should do the trick (see mlir/test/Target/SPIRV). Those tests serialize given MLIR into SPIR-V and then deserialize the SPIR-V back to MLIR. So, I would hope you could see a truncation after the roundtrip.

It has been tricky to add a good test for this. Something needs to trigger the encoding of a string literal with more than 262140 chars.
Any suggestion on which spirv operation to use in such a test?

[IgWod-IMG]

Any suggestion on which spirv operation to use in such a test?

I thought that maybe spirv.EntryPoint would work but no success. I also tried spirv.ExecutionMode. None of my attempts reach your code. So, I’m afraid that I cannot help.

[davidegrohmann]

Any suggestion on which spirv operation to use in such a test?

I thought that maybe spirv.EntryPoint would work but no success. I also tried spirv.ExecutionMode. None of my attempts reach your code. So, I’m afraid that I cannot help.

I have tried the same plus playing around with location information without success.

[kuhar]

LGTM % nits

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[github-actions]

@davidegrohmann Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!