[mlir][spirv] Avoid duplicated Block decoration during serialization

[antiagainst]

Summary

It's legal per the Vulkan / SPIR-V spec; still it's better to avoid
such duplication to have cleaner blob and reduce the binary size.

Depends On D115531

mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp

 Serializer::processGlobalVariableOp(spirv::GlobalVariableOp varOp) {
   // Get TypeID.
   uint32_t resultTypeID = 0;
   SmallVector<StringRef, 4> elidedAttrs;
   if (failed(processType(varOp.getLoc(), varOp.type(), resultTypeID))) {
     return failure();
   }
 
-  if (isInterfaceStructPtrType(varOp.type())) {
-    auto structType = varOp.type()
-                          .cast<spirv::PointerType>()
-                          .getPointeeType()
-                          .cast<spirv::StructType>();
-    if (failed(
-            emitDecoration(getTypeID(structType), spirv::Decoration::Block))) {
-      return varOp.emitError("cannot decorate ")
-             << structType << " with Block decoration";
-    }
-  }
-
   elidedAttrs.push_back("type");
   SmallVector<uint32_t, 4> operands;
   operands.push_back(resultTypeID);
   auto resultID = getNextID();
 
   // Encode the name.
   auto varName = varOp.sym_name();
   elidedAttrs.push_back(SymbolTable::getSymbolAttrName());

mlir/lib/Target/SPIRV/Serialization/Serializer.cpp

LogicalResult Serializer::processType(Location loc, Type type,
   SetVector<StringRef> serializationCtx;
   return processTypeImpl(loc, type, typeID, serializationCtx);
 }
 
 LogicalResult
 Serializer::processTypeImpl(Location loc, Type type, uint32_t &typeID,
                             SetVector<StringRef> &serializationCtx) {
   typeID = getTypeID(type);
-  if (typeID) {
+  if (typeID)
     return success();
-  }
+
   typeID = getNextID();
   SmallVector<uint32_t, 4> operands;
 
   operands.push_back(typeID);
   auto typeEnum = spirv::Opcode::OpTypeVoid;
   bool deferSerialization = false;
 
   if ((type.isa<FunctionType>() &&
if (pointeeStruct && pointeeStruct.isIdentified() &&
       if (failed(processTypeImpl(loc, ptrType.getPointeeType(), pointeeTypeID,
                                  serializationCtx)))
         return failure();
     }
 
     typeEnum = spirv::Opcode::OpTypePointer;
     operands.push_back(static_cast<uint32_t>(ptrType.getStorageClass()));
     operands.push_back(pointeeTypeID);
+
+    if (isInterfaceStructPtrType(ptrType)) {
+      if (failed(emitDecoration(getTypeID(pointeeStruct),
+                                spirv::Decoration::Block)))
+        return emitError(loc, "cannot decorate ")
+               << pointeeStruct << " with Block decoration";
+    }
+
     return success();
   }
 
   if (auto runtimeArrayType = type.dyn_cast<spirv::RuntimeArrayType>()) {
     uint32_t elementTypeID = 0;
     if (failed(processTypeImpl(loc, runtimeArrayType.getElementType(),
                                elementTypeID, serializationCtx))) {
       return failure();

mlir/unittests/Dialect/SPIRV/SerializationTest.cpp

protected:
   spirv::GlobalVariableOp addGlobalVar(Type type, llvm::StringRef name) {
     OpBuilder builder(module->getRegion());
     auto ptrType = spirv::PointerType::get(type, spirv::StorageClass::Uniform);
     return builder.create<spirv::GlobalVariableOp>(
         UnknownLoc::get(&context), TypeAttr::get(ptrType),
         builder.getStringAttr(name), nullptr);
   }
 
+  /// Handles a SPIR-V instruction with the given `opcode` and `operand`.
+  /// Returns true to interrupt.
+  using HandleFn = llvm::function_ref<bool(spirv::Opcode opcode,
+                                           ArrayRef<uint32_t> operands)>;
+
   /// Returns true if we can find a matching instruction in the SPIR-V blob.
-  bool findInstruction(llvm::function_ref<bool(spirv::Opcode opcode,
-                                               ArrayRef<uint32_t> operands)>
-                           matchFn) {
+  bool scanInstruction(HandleFn handleFn) {
     auto binarySize = binary.size();
     auto *begin = binary.begin();
     auto currOffset = spirv::kHeaderWordCount;
 
     while (currOffset < binarySize) {
       auto wordCount = binary[currOffset] >> 16;
-      if (!wordCount || (currOffset + wordCount > binarySize)) {
+      if (!wordCount || (currOffset + wordCount > binarySize))
         return false;
-      }
+
       spirv::Opcode opcode =
           static_cast<spirv::Opcode>(binary[currOffset] & 0xffff);
-
-      if (matchFn(opcode,
-                  llvm::ArrayRef<uint32_t>(begin + currOffset + 1,
-                                           begin + currOffset + wordCount))) {
+      llvm::ArrayRef<uint32_t> operands(begin + currOffset + 1,
+                                        begin + currOffset + wordCount);
+      if (handleFn(opcode, operands))
         return true;
-      }
+
       currOffset += wordCount;
     }
     return false;
   }
 
 protected:
   MLIRContext context;
   OwningOpRef<spirv::ModuleOp> module;
   SmallVector<uint32_t, 0> binary;
 };
 
 //===----------------------------------------------------------------------===//
 // Block decoration
 //===----------------------------------------------------------------------===//
 
 TEST_F(SerializationTest, ContainsBlockDecoration) {
   auto structType = getFloatStructType();
   addGlobalVar(structType, "var0");
 
   ASSERT_TRUE(succeeded(spirv::serialize(module.get(), binary)));
 
   auto hasBlockDecoration = [](spirv::Opcode opcode,
-                               ArrayRef<uint32_t> operands) -> bool {
-    if (opcode != spirv::Opcode::OpDecorate || operands.size() != 2)
-      return false;
-    return operands[1] == static_cast<uint32_t>(spirv::Decoration::Block);
+                               ArrayRef<uint32_t> operands) {
+    return opcode == spirv::Opcode::OpDecorate && operands.size() == 2 &&
+           operands[1] == static_cast<uint32_t>(spirv::Decoration::Block);
+  };
+  EXPECT_TRUE(scanInstruction(hasBlockDecoration));
+}
+
+TEST_F(SerializationTest, ContainsNoDuplicatedBlockDecoration) {
+  auto structType = getFloatStructType();
+  // Two global variables using the same type should not decorate the type with
+  // duplicated `Block` decorations.
+  addGlobalVar(structType, "var0");
+  addGlobalVar(structType, "var1");
+
+  ASSERT_TRUE(succeeded(spirv::serialize(module.get(), binary)));
+
+  unsigned count = 0;
+  auto countBlockDecoration = [&count](spirv::Opcode opcode,
+                                       ArrayRef<uint32_t> operands) {
+    if (opcode == spirv::Opcode::OpDecorate && operands.size() == 2 &&
+        operands[1] == static_cast<uint32_t>(spirv::Decoration::Block))
+      ++count;
   };
-  EXPECT_TRUE(findInstruction(hasBlockDecoration));
+  ASSERT_FALSE(scanInstruction(countBlockDecoration));
+  EXPECT_EQ(count, 1u);
 }
 
 TEST_F(SerializationTest, ContainsSymbolName) {
   auto structType = getFloatStructType();
   addGlobalVar(structType, "var0");
 
   spirv::SerializationOptions options;
   options.emitSymbolName = true;
   ASSERT_TRUE(succeeded(spirv::serialize(module.get(), binary, options)));
 
   auto hasVarName = [](spirv::Opcode opcode, ArrayRef<uint32_t> operands) {
     unsigned index = 1; // Skip the result <id>
     return opcode == spirv::Opcode::OpName &&
            spirv::decodeStringLiteral(operands, index) == "var0";
   };
-  EXPECT_TRUE(findInstruction(hasVarName));
+  EXPECT_TRUE(scanInstruction(hasVarName));
 }
 
 TEST_F(SerializationTest, DoesNotContainSymbolName) {
   auto structType = getFloatStructType();
   addGlobalVar(structType, "var0");
 
   spirv::SerializationOptions options;
   options.emitSymbolName = false;
   ASSERT_TRUE(succeeded(spirv::serialize(module.get(), binary, options)));
 
   auto hasVarName = [](spirv::Opcode opcode, ArrayRef<uint32_t> operands) {
     unsigned index = 1; // Skip the result <id>
     return opcode == spirv::Opcode::OpName &&
            spirv::decodeStringLiteral(operands, index) == "var0";
   };
-  EXPECT_FALSE(findInstruction(hasVarName));
+  EXPECT_FALSE(scanInstruction(hasVarName));
 }