[Mlir-commits] [mlir] [mlir][SPIR-V] Add support for SPV_INTEL_long_composites extension (PR #195685)

[Jakub Kuderski]

================
@@ -229,3 +230,223 @@ TEST_F(SerializationTest, DoesNotContainSymbolName) {
   };
   EXPECT_FALSE(scanInstruction(hasVarName));
 }
+
+//===----------------------------------------------------------------------===//
+// SPV_INTEL_long_composites: composites whose binary form would exceed the
+// SPIR-V 16-bit word-count limit are split into a parent + *ContinuedINTEL ops
+// on serialization, and merged back on deserialization. These tests build the
+// large composites programmatically so that the IR doesn't have to expand
+// thousands of operands literally.
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+// Picked to comfortably exceed kMaxWordCount = 65535 for any of the splittable
+// composite/struct opcodes -- each one packs at most kMaxWordCount - {1,2,3}
+// operands into the parent word, so 65540 always triggers a split.
+constexpr unsigned kLongCompositeSize = 65540;
+
+bool hasOpcode(SmallVectorImpl<uint32_t> &binary, spirv::Opcode target) {
+  size_t offset = spirv::kHeaderWordCount;
+  while (offset < binary.size()) {
+    uint32_t wordCount = binary[offset] >> 16;
+    if (!wordCount || offset + wordCount > binary.size())
+      return false;
+    auto op = static_cast<spirv::Opcode>(binary[offset] & 0xffff);
+    if (op == target)
+      return true;
+    offset += wordCount;
+  }
+  return false;
+}
+
+bool hasLongCompositesCapabilityAndExtension(SmallVectorImpl<uint32_t> &b) {
+  bool foundCap = false;
+  bool foundExt = false;
+  size_t offset = spirv::kHeaderWordCount;
+  while (offset < b.size()) {
+    uint32_t wordCount = b[offset] >> 16;
+    if (!wordCount || offset + wordCount > b.size())
+      break;
+    auto op = static_cast<spirv::Opcode>(b[offset] & 0xffff);
+    ArrayRef<uint32_t> operands(b.data() + offset + 1, wordCount - 1);
+    if (op == spirv::Opcode::OpCapability && !operands.empty() &&
+        operands[0] ==
+            static_cast<uint32_t>(spirv::Capability::LongCompositesINTEL))
+      foundCap = true;
+    if (op == spirv::Opcode::OpExtension) {
+      unsigned idx = 0;
+      if (spirv::decodeStringLiteral(operands, idx) ==
+          spirv::stringifyExtension(
+              spirv::Extension::SPV_INTEL_long_composites))
+        foundExt = true;
+    }
+    offset += wordCount;
+  }
+  return foundCap && foundExt;
+}
+
+// Verifies that no instruction in the binary has a word count exceeding the
+// SPIR-V 16-bit limit (which would mean the splitting logic failed).
+bool allInstructionsWithinWordLimit(SmallVectorImpl<uint32_t> &b) {
+  size_t offset = spirv::kHeaderWordCount;
+  while (offset < b.size()) {
+    uint32_t wordCount = b[offset] >> 16;
+    if (!wordCount || wordCount > spirv::kMaxWordCount)
+      return false;
+    offset += wordCount;
+  }
+  return true;
+}
+
+} // namespace
+
+TEST_F(SerializationTest, LongTypeStructIsSplit) {
+  // Build a struct with kLongCompositeSize i8 members and reference it via a
+  // global variable so the type gets serialized.
+  OpBuilder builder(module->getRegion());
+  Type i8Type = builder.getIntegerType(8);
----------------
kuhar wrote:

These tests build long composites out of `i8`, but the fixture creates the module with an empty VCE triple. The serializer then emits `OpTypeInt 8` without the `Int8` capability, so the tests can pass while exercising invalid SPIR-V for reasons unrelated to long-composites.

Please either use a type that is valid with the baseline capabilities, such as `i32`, or initialize these test modules with the required `Int8` capability.

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