[spirv] Use owning module ref to avoid leaks and fix ASAN tests

mlir/lib/Dialect/SPIRV/Serialization/Deserializer.cpp

public:
   /// Creates a deserializer for the given SPIR-V `binary` module.
   /// The SPIR-V ModuleOp will be created into `context.
   explicit Deserializer(ArrayRef<uint32_t> binary, MLIRContext *context);
 
   /// Deserializes the remembered SPIR-V binary module.
   LogicalResult deserialize();
 
   /// Collects the final SPIR-V ModuleOp.
-  Optional<spirv::ModuleOp> collect();
+  spirv::OwningSPIRVModuleRef collect();
 
 private:
   //===--------------------------------------------------------------------===//
   // Module structure
   //===--------------------------------------------------------------------===//
 
   /// Initializes the `module` ModuleOp in this deserializer instance.
-  spirv::ModuleOp createModuleOp();
+  spirv::OwningSPIRVModuleRef createModuleOp();
 
   /// Processes SPIR-V module header in `binary`.
   LogicalResult processHeader();
 
   /// Processes the SPIR-V OpCapability with `operands` and updates bookkeeping
   /// in the deserializer.
   LogicalResult processCapability(ArrayRef<uint32_t> operands);
 
private:
 
   /// MLIRContext to create SPIR-V ModuleOp into.
   MLIRContext *context;
 
   // TODO: create Location subclass for binary blob
   Location unknownLoc;
 
   /// The SPIR-V ModuleOp.
-  Optional<spirv::ModuleOp> module;
+  spirv::OwningSPIRVModuleRef module;
 
   /// The current function under construction.
   Optional<spirv::FuncOp> curFunction;
 
   /// The current block under construction.
   Block *curBlock = nullptr;
 
   OpBuilder opBuilder;
LogicalResult Deserializer::deserialize() {
   }
 
   attachVCETriple();
 
   LLVM_DEBUG(llvm::dbgs() << "+++ completed deserialization +++\n");
   return success();
 }
 
-Optional<spirv::ModuleOp> Deserializer::collect() { return module; }
+spirv::OwningSPIRVModuleRef Deserializer::collect() {
+  return std::move(module);
+}
 
 //===----------------------------------------------------------------------===//
 // Module structure
 //===----------------------------------------------------------------------===//
 
-spirv::ModuleOp Deserializer::createModuleOp() {
+spirv::OwningSPIRVModuleRef Deserializer::createModuleOp() {
   OpBuilder builder(context);
   OperationState state(unknownLoc, spirv::ModuleOp::getOperationName());
   spirv::ModuleOp::build(builder, state);
   return cast<spirv::ModuleOp>(Operation::create(state));
 }
 
 LogicalResult Deserializer::processHeader() {
   if (binary.size() < spirv::kHeaderWordCount)
for (auto *block : constructBlocks) {
     }
     for (auto &op : *block)
       newBlock->push_back(op.clone(mapper));
   }
 
   // Go through all ops and remap the operands.
   auto remapOperands = [&](Operation *op) {
     for (auto &operand : op->getOpOperands())
-      if (auto mappedOp = mapper.lookupOrNull(operand.get()))
+      if (Value mappedOp = mapper.lookupOrNull(operand.get()))
         operand.set(mappedOp);
     for (auto &succOp : op->getBlockOperands())
-      if (auto mappedOp = mapper.lookupOrNull(succOp.get()))
+      if (Block *mappedOp = mapper.lookupOrNull(succOp.get()))
         succOp.set(mappedOp);
   };
   for (auto &block : body) {
     block.walk(remapOperands);
   }
 
   // We have created the SelectionOp/LoopOp and "moved" all blocks belonging to
   // the selection/loop construct into its region. Next we need to fix the

 template <>
 LogicalResult
 Deserializer::processOp<spirv::EntryPointOp>(ArrayRef<uint32_t> words) {
   unsigned wordIndex = 0;
   if (wordIndex >= words.size()) {
     return emitError(unknownLoc,
                      "missing Execution Model specification in OpEntryPoint");
   }
-  auto exec_model = opBuilder.getI32IntegerAttr(words[wordIndex++]);
+  auto execModel = opBuilder.getI32IntegerAttr(words[wordIndex++]);
   if (wordIndex >= words.size()) {
     return emitError(unknownLoc, "missing <id> in OpEntryPoint");
   }
   // Get the function <id>
   auto fnID = words[wordIndex++];
   // Get the function name
   auto fnName = decodeStringLiteral(words, wordIndex);
   // Verify that the function <id> matches the fnName
while (wordIndex < words.size()) {
     auto arg = getGlobalVariable(words[wordIndex]);
     if (!arg) {
       return emitError(unknownLoc, "undefined result <id> ")
              << words[wordIndex] << " while decoding OpEntryPoint";
     }
     interface.push_back(opBuilder.getSymbolRefAttr(arg.getOperation()));
     wordIndex++;
   }
-  opBuilder.create<spirv::EntryPointOp>(unknownLoc, exec_model,
+  opBuilder.create<spirv::EntryPointOp>(unknownLoc, execModel,
                                         opBuilder.getSymbolRefAttr(fnName),
                                         opBuilder.getArrayAttr(interface));
   return success();
 }
 
 template <>
 LogicalResult
 Deserializer::processOp<spirv::ExecutionModeOp>(ArrayRef<uint32_t> words) {

 
 spirv::OwningSPIRVModuleRef spirv::deserialize(ArrayRef<uint32_t> binary,
                                                MLIRContext *context) {
   Deserializer deserializer(binary, context);
 
   if (failed(deserializer.deserialize()))
     return nullptr;
 
-  return deserializer.collect().getValueOr(nullptr);
+  return deserializer.collect();
 }

mlir/unittests/Dialect/SPIRV/SerializationTest.cpp

 // covered by normal serialization and deserialization roundtripping.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/Serialization.h"
 #include "mlir/Dialect/SPIRV/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/SPIRVBinaryUtils.h"
 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/SPIRVModule.h"
 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/SPIRVTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/MLIRContext.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Sequence.h"
state.addAttribute("vce_triple",
                        spirv::VerCapExtAttr::get(
                            spirv::Version::V_1_0, ArrayRef<spirv::Capability>(),
                            ArrayRef<spirv::Extension>(), &context));
     spirv::ModuleOp::build(builder, state);
     module = cast<spirv::ModuleOp>(Operation::create(state));
   }
 
   Type getFloatStructType() {
-    OpBuilder opBuilder(module.body());
+    OpBuilder opBuilder(module->body());
     llvm::SmallVector<Type, 1> elementTypes{opBuilder.getF32Type()};
     llvm::SmallVector<spirv::StructType::OffsetInfo, 1> offsetInfo{0};
     auto structType = spirv::StructType::get(elementTypes, offsetInfo);
     return structType;
   }
 
   void addGlobalVar(Type type, llvm::StringRef name) {
-    OpBuilder opBuilder(module.body());
+    OpBuilder opBuilder(module->body());
     auto ptrType = spirv::PointerType::get(type, spirv::StorageClass::Uniform);
     opBuilder.create<spirv::GlobalVariableOp>(
         UnknownLoc::get(&context), TypeAttr::get(ptrType),
         opBuilder.getStringAttr(name), nullptr);
   }
 
   bool findInstruction(llvm::function_ref<bool(spirv::Opcode opcode,
                                                ArrayRef<uint32_t> operands)>
while (currOffset < binarySize) {
       }
       currOffset += wordCount;
     }
     return false;
   }
 
 protected:
   MLIRContext context;
-  spirv::ModuleOp module;
+  spirv::OwningSPIRVModuleRef module;
   SmallVector<uint32_t, 0> binary;
 };
 
 //===----------------------------------------------------------------------===//
 // Block decoration
 //===----------------------------------------------------------------------===//
 
 TEST_F(SerializationTest, BlockDecorationTest) {
   auto structType = getFloatStructType();
   addGlobalVar(structType, "var0");
-  ASSERT_TRUE(succeeded(spirv::serialize(module, binary)));
+  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);
   };
   EXPECT_TRUE(findInstruction(hasBlockDecoration));
 }