[mlir][spirv] Introduce OwningSPIRVModuleRef for ownership

[antiagainst]

Summary

Similar to OwningModuleRef, OwningSPIRVModuleRef signals ownership
transfer clearly. This is useful for APIs like spirv::deserialize,
where a spirv::ModuleOp is returned by deserializing SPIR-V binary
module.

This addresses the ASAN error as reported in
https://bugs.llvm.org/show_bug.cgi?id=46272

[antiagainst]

I also tried to turn the existing OwningModuleRef into a template class OwningModuleRefBase and using OwningModuleRef = OwningModuleRefBase<mlir::ModuleOp>. But it means we need now to include the Module.h header where we forward declare class OwningModuleRef. That seems more intrusive than the current approach, which means a bit code duplication. But I'm happy to go back to the template approach if you think it's better. Just let me know. :)

[rriddle]

I was somewhat strongly against this in the past because it sets the wrong usage pattern expectations. Though the example you have here makes more sense to me than say FuncOp. That being said, I would prefer we don't duplicate this code and instead use a templated version.

[rriddle]

In D81652#2087519, @antiagainst wrote:

\But it means we need now to include the Module.h header where we forward declare class OwningModuleRef.

That's only an issue if you have a 'using' statement. I would expect that you could use inheritance and achieve the same result. That would only break down in a situation where you wanted to templatize based on the BaseRef class, but I don't foresee such a usecase right now.

[mehdi_amini]

In D81652#2087976, @rriddle wrote:

I was somewhat strongly against this in the past because it sets the wrong usage pattern expectations.

Can you clarify what you mean with "pattern expectations" here?

[rriddle]

In D81652#2088372, @mehdi_amini wrote:

In D81652#2087976, @rriddle wrote:

I was somewhat strongly against this in the past because it sets the wrong usage pattern expectations.

Can you clarify what you mean with "pattern expectations" here?

The "usage pattern expectations" I mean are that we shouldn't be using this type of API within passes, patterns, or a majority of the infra that people usually work in. IMO this is more for the top-level layer, i.e., things like parsing an IR blob. In most normal use cases, passing a builder should be preferred over manually constructing and managing the lifetime of the IR.

[mehdi_amini]

In D81652#2088817, @rriddle wrote:

In D81652#2088372, @mehdi_amini wrote:

In D81652#2087976, @rriddle wrote:

I was somewhat strongly against this in the past because it sets the wrong usage pattern expectations.

Can you clarify what you mean with "pattern expectations" here?

The "usage pattern expectations" I mean are that we shouldn't be using this type of API within passes, patterns, or a majority of the infra that people usually work in. IMO this is more for the top-level layer, i.e., things like parsing an IR blob. In most normal use cases, passing a builder should be preferred over manually constructing and managing the lifetime of the IR.

Oh right: builders are much better. We still have a bunch of "unsafe" APIs though, and even though aren't the most common to use I am not sure why they couldn't get more type safety?

[antiagainst]

In D81652#2087976, @rriddle wrote:

I was somewhat strongly against this in the past because it sets the wrong usage pattern expectations. Though the example you have here makes more sense to me than say FuncOp. That being said, I would prefer we don't duplicate this code and instead use a templated version.

SG. I'll change later then.

[antiagainst]

Address comments

[antiagainst]

This is now changed to use inheritance as suggested by River. Let me know if this is okay now.

[rriddle]

mlir/include/mlir/IR/ModuleRefBase.h

21 ↗

(On Diff #275884)

Can you add that OpBuilder and related functionality should be highly preferred instead, and this should only be used in situations where existing solutions are not viable.

24 ↗

(On Diff #275884)

nit: I would likely just rename to OwningOpRefBase instead.

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

1750

nit: return deserializer.collect.getValueOr(nullptr);

[antiagainst]

Address comments: renamed class and file and add comments

[antiagainst]

Fix accidentially removed code

[antiagainst]

Rename variables too

mlir/include/mlir/Dialect/SPIRV/SPIRVModule.h

+//===- SPIRVModule.h - SPIR-V Module Utilities ------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPIRV_SPIRVMODULE_H
+#define MLIR_DIALECT_SPIRV_SPIRVMODULE_H
+
+#include "mlir/Dialect/SPIRV/SPIRVOps.h"
+#include "mlir/IR/OwningOpRefBase.h"
+
+namespace mlir {
+namespace spirv {
+
+/// This class acts as an owning reference to a SPIR-V module, and will
+/// automatically destroy the held module on destruction if the held module
+/// is valid.
+class OwningSPIRVModuleRef : public OwningOpRefBase<spirv::ModuleOp> {
+public:
+  using OwningOpRefBase<spirv::ModuleOp>::OwningOpRefBase;
+};
+
+} // end namespace spirv
+} // end namespace mlir
+
+#endif // MLIR_DIALECT_SPIRV_SPIRVMODULE_H

mlir/include/mlir/Dialect/SPIRV/Serialization.h

 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
 struct LogicalResult;
 class MLIRContext;
 
 namespace spirv {
 class ModuleOp;
+class OwningSPIRVModuleRef;
 
 /// Serializes the given SPIR-V `module` and writes to `binary`. On failure,
 /// reports errors to the error handler registered with the MLIR context for
 /// `module`.
 LogicalResult serialize(ModuleOp module, SmallVectorImpl<uint32_t> &binary,
                         bool emitDebugInfo = false);
 
 /// Deserializes the given SPIR-V `binary` module and creates a MLIR ModuleOp
 /// in the given `context`. Returns the ModuleOp on success; otherwise, reports
-/// errors to the error handler registered with `context` and returns
-/// llvm::None.
-Optional<ModuleOp> deserialize(ArrayRef<uint32_t> binary, MLIRContext *context);
+/// errors to the error handler registered with `context` and returns a null
+/// module.
+OwningSPIRVModuleRef deserialize(ArrayRef<uint32_t> binary,
+                                 MLIRContext *context);
 
 } // end namespace spirv
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_SPIRV_SERIALIZATION_H_

mlir/include/mlir/IR/Module.h

 //===- Module.h - MLIR Module Class -----------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Module is the top-level container for code in an MLIR program.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_MODULE_H
 #define MLIR_IR_MODULE_H
 
+#include "mlir/IR/OwningOpRefBase.h"
 #include "mlir/IR/SymbolTable.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
 class ModuleTerminatorOp;
 
 //===----------------------------------------------------------------------===//
 // Module Operation.
: public Op<ModuleTerminatorOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
                 OpTrait::HasParent<ModuleOp>::Impl, OpTrait::IsTerminator> {
 public:
   using Op::Op;
   static StringRef getOperationName() { return "module_terminator"; }
   static void build(OpBuilder &, OperationState &) {}
 };
 
 /// This class acts as an owning reference to a module, and will automatically
-/// destroy the held module if valid.
-class OwningModuleRef {
+/// destroy the held module on destruction if the held module is valid.
+class OwningModuleRef : public OwningOpRefBase<ModuleOp> {
 public:
-  OwningModuleRef(std::nullptr_t = nullptr) {}
-  OwningModuleRef(ModuleOp module) : module(module) {}
-  OwningModuleRef(OwningModuleRef &&other) : module(other.release()) {}
-  ~OwningModuleRef() {
-    if (module)
-      module.erase();
-  }
-
-  // Assign from another module reference.
-  OwningModuleRef &operator=(OwningModuleRef &&other) {
-    if (module)
-      module.erase();
-    module = other.release();
-    return *this;
-  }
-
-  /// Allow accessing the internal module.
-  ModuleOp get() const { return module; }
-  ModuleOp operator*() const { return module; }
-  ModuleOp *operator->() { return &module; }
-  explicit operator bool() const { return module; }
-
-  /// Release the referenced module.
-  ModuleOp release() {
-    ModuleOp released;
-    std::swap(released, module);
-    return released;
-  }
-
-private:
-  ModuleOp module;
+  using OwningOpRefBase<ModuleOp>::OwningOpRefBase;
 };
 
 } // end namespace mlir
 
 namespace llvm {
 
 /// Allow stealing the low bits of ModuleOp.
 template <> struct PointerLikeTypeTraits<mlir::ModuleOp> {

mlir/include/mlir/IR/OwningOpRefBase.h

+//===- OwningOpRefBase.h - MLIR OwningOpRefBase -----------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file provides a base class for owning module refs.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_MODULEREFBASE_H
+#define MLIR_IR_MODULEREFBASE_H
+
+#include <utility>
+
+namespace mlir {
+
+/// This class acts as an owning reference to an op, and will automatically
+/// destroy the held module on destruction if the held module is valid.
+///
+/// Note that OpBuilder and related functionality should be highly preferred
+/// instead, and this should only be used in situations where existing solutions
+/// are not viable.
+template <typename ModuleTy>
+class OwningOpRefBase {
+public:
+  OwningOpRefBase(std::nullptr_t = nullptr) {}
+  OwningOpRefBase(ModuleTy module) : module(module) {}
+  OwningOpRefBase(OwningOpRefBase &&other) : module(other.release()) {}
+  ~OwningOpRefBase() {
+    if (module)
+      module.erase();
+  }
+
+  // Assign from another module reference.
+  OwningOpRefBase &operator=(OwningOpRefBase &&other) {
+    if (module)
+      module.erase();
+    module = other.release();
+    return *this;
+  }
+
+  /// Allow accessing the internal module.
+  ModuleTy get() const { return module; }
+  ModuleTy operator*() const { return module; }
+  ModuleTy *operator->() { return &module; }
+  explicit operator bool() const { return module; }
+
+  /// Release the referenced module.
+  ModuleTy release() {
+    ModuleTy released;
+    std::swap(released, module);
+    return released;
+  }
+
+private:
+  ModuleTy module;
+};
+
+} // end namespace mlir
+
+#endif // MLIR_IR_MODULEREFBASE_H

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

 //===- Deserializer.cpp - MLIR SPIR-V Deserialization ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the SPIR-V binary to MLIR SPIR-V module deserialization.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/Serialization.h"
 
 #include "mlir/Dialect/SPIRV/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/SPIRVBinaryUtils.h"
+#include "mlir/Dialect/SPIRV/SPIRVModule.h"
 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/SPIRVTypes.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Location.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Sequence.h"
LogicalResult Deserializer::processPhi(ArrayRef<uint32_t> operands) {
   BlockArgument blockArg = curBlock->addArgument(blockArgType);
   valueMap[operands[1]] = blockArg;
   LLVM_DEBUG(llvm::dbgs() << "[phi] created block argument " << blockArg
                           << " id = " << operands[1] << " of type "
                           << blockArgType << '\n');
 
   // For each (value, predecessor) pair, insert the value to the predecessor's
   // blockPhiInfo entry so later we can fix the block argument there.
-  for (unsigned i = 2, e = operands.size(); i < e; i += 2) {
-    uint32_t value = operands[i];
-    Block *predecessor = getOrCreateBlock(operands[i + 1]);
-    blockPhiInfo[predecessor].push_back(value);
-    LLVM_DEBUG(llvm::dbgs() << "[phi] predecessor @ " << predecessor
-                            << " with arg id = " << value << '\n');
-  }
-
-  return success();
-}
-
-namespace {
-/// A class for putting all blocks in a structured selection/loop in a
-/// spv.selection/spv.loop op.
-class ControlFlowStructurizer {
-public:
-  /// Structurizes the loop at the given `headerBlock`.
-  ///
-  /// This method will create an spv.loop op in the `mergeBlock` and move all
-  /// blocks in the structured loop into the spv.loop's region. All branches to
-  /// the `headerBlock` will be redirected to the `mergeBlock`.
-  /// This method will also update `mergeInfo` by remapping all blocks inside to
-  /// the newly cloned ones inside structured control flow op's regions.
-  static LogicalResult structurize(Location loc, BlockMergeInfoMap &mergeInfo,
-                                   Block *headerBlock, Block *mergeBlock,
-                                   Block *continueBlock) {
-    return ControlFlowStructurizer(loc, mergeInfo, headerBlock, mergeBlock,
-                                   continueBlock)
-        .structurizeImpl();
-  }
-
-private:
-  ControlFlowStructurizer(Location loc, BlockMergeInfoMap &mergeInfo,
-                          Block *header, Block *merge, Block *cont)
-      : location(loc), blockMergeInfo(mergeInfo), headerBlock(header),
-        mergeBlock(merge), continueBlock(cont) {}
-
-  /// Creates a new spv.selection op at the beginning of the `mergeBlock`.
-  spirv::SelectionOp createSelectionOp();
-
-  /// Creates a new spv.loop op at the beginning of the `mergeBlock`.
-  spirv::LoopOp createLoopOp();
-
-  /// Collects all blocks reachable from `headerBlock` except `mergeBlock`.
-  void collectBlocksInConstruct();
-
-  LogicalResult structurizeImpl();
-
-  Location location;
-
-  BlockMergeInfoMap &blockMergeInfo;
-
-  Block *headerBlock;
-  Block *mergeBlock;
-  Block *continueBlock; // nullptr for spv.selection
-
-  llvm::SetVector<Block *> constructBlocks;
-};
-} // namespace
-
-spirv::SelectionOp ControlFlowStructurizer::createSelectionOp() {
-  // Create a builder and set the insertion point to the beginning of the
-  // merge block so that the newly created SelectionOp will be inserted there.
-  OpBuilder builder(&mergeBlock->front());
-
-  auto control = builder.getI32IntegerAttr(
-      static_cast<uint32_t>(spirv::SelectionControl::None));
-  auto selectionOp = builder.create<spirv::SelectionOp>(location, control);
-  selectionOp.addMergeBlock();
-
-  return selectionOp;
-}
-
-spirv::LoopOp ControlFlowStructurizer::createLoopOp() {
-  // Create a builder and set the insertion point to the beginning of the
-  // merge block so that the newly created LoopOp will be inserted there.
-  OpBuilder builder(&mergeBlock->front());
-
-  // TODO(antiagainst): handle loop control properly
-  auto loopOp = builder.create<spirv::LoopOp>(location);
-  loopOp.addEntryAndMergeBlock();
-
-  return loopOp;
-}
-
-void ControlFlowStructurizer::collectBlocksInConstruct() {
-  assert(constructBlocks.empty() && "expected empty constructBlocks");
-
-  // Put the header block in the work list first.
-  constructBlocks.insert(headerBlock);
-
-  // For each item in the work list, add its successors excluding the merge
-  // block.
-  for (unsigned i = 0; i < constructBlocks.size(); ++i) {
-    for (auto *successor : constructBlocks[i]->getSuccessors())
-      if (successor != mergeBlock)
-        constructBlocks.insert(successor);
-  }
-}
-
-LogicalResult ControlFlowStructurizer::structurizeImpl() {
-  Operation *op = nullptr;
-  bool isLoop = continueBlock != nullptr;
-  if (isLoop) {
-    if (auto loopOp = createLoopOp())
-      op = loopOp.getOperation();
-  } else {
-    if (auto selectionOp = createSelectionOp())
-      op = selectionOp.getOperation();
-  }
-  if (!op)
-    return failure();
-  Region &body = op->getRegion(0);
-
-  BlockAndValueMapping mapper;
-  // All references to the old merge block should be directed to the
-  // selection/loop merge block in the SelectionOp/LoopOp's region.
-  mapper.map(mergeBlock, &body.back());
-
-  collectBlocksInConstruct();
-
-  // We've identified all blocks belonging to the selection/loop's region. Now
-  // need to "move" them into the selection/loop. Instead of really moving the
-  // blocks, in the following we copy them and remap all values and branches.
-  // This is because:
-  // * Inserting a block into a region requires the block not in any region
-  //   before. But selections/loops can nest so we can create selection/loop ops
-  //   in a nested manner, which means some blocks may already be in a
-  //   selection/loop region when to be moved again.
-  // * It's much trickier to fix up the branches into and out of the loop's
-  //   region: we need to treat not-moved blocks and moved blocks differently:
-  //   Not-moved blocks jumping to the loop header block need to jump to the
-  //   merge point containing the new loop op but not the loop continue block's
-  //   back edge. Moved blocks jumping out of the loop need to jump to the
-  //   merge block inside the loop region but not other not-moved blocks.
-  //   We cannot use replaceAllUsesWith clearly and it's harder to follow the
-  //   logic.
-
-  // Create a corresponding block in the SelectionOp/LoopOp's region for each
-  // block in this loop construct.
-  OpBuilder builder(body);
-  for (auto *block : constructBlocks) {
-    // Create a block and insert it before the selection/loop merge block in the
-    // SelectionOp/LoopOp's region.
-    auto *newBlock = builder.createBlock(&body.back());
-    mapper.map(block, newBlock);
-    LLVM_DEBUG(llvm::dbgs() << "[cf] cloned block " << newBlock
-                            << " from block " << block << "\n");
-    if (!isFnEntryBlock(block)) {
-      for (BlockArgument blockArg : block->getArguments()) {
-        auto newArg = newBlock->addArgument(blockArg.getType());
-        mapper.map(blockArg, newArg);
-        LLVM_DEBUG(llvm::dbgs() << "[cf] remapped block argument " << blockArg
-                                << " to " << newArg << '\n');
-      }
-    } else {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "[cf] block " << block << " is a function entry block\n");
-    }
-    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()))
-        operand.set(mappedOp);
-    for (auto &succOp : op->getBlockOperands())
-      if (auto 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
-  // connections between this new SelectionOp/LoopOp with existing blocks.
-
-  // All existing incoming branches should go to the merge block, where the
-  // SelectionOp/LoopOp resides right now.
-  headerBlock->replaceAllUsesWith(mergeBlock);
-
-  if (isLoop) {
-    // The loop selection/loop header block may have block arguments. Since now
-    // we place the selection/loop op inside the old merge block, we need to
-    // make sure the old merge block has the same block argument list.
-    assert(mergeBlock->args_empty() && "OpPhi in loop merge block unsupported");
-    for (BlockArgument blockArg : headerBlock->getArguments()) {
-      mergeBlock->addArgument(blockArg.getType());
-    }
-
-    // If the loop header block has block arguments, make sure the spv.branch op
-    // matches.
-    SmallVector<Value, 4> blockArgs;
-    if (!headerBlock->args_empty())
-      blockArgs = {mergeBlock->args_begin(), mergeBlock->args_end()};
-
-    // The loop entry block should have a unconditional branch jumping to the
-    // loop header block.
-    builder.setInsertionPointToEnd(&body.front());
-    builder.create<spirv::BranchOp>(location, mapper.lookupOrNull(headerBlock),
-                                    ArrayRef<Value>(blockArgs));
-  }
-
-  // All the blocks cloned into the SelectionOp/LoopOp's region can now be
-  // cleaned up.
-  LLVM_DEBUG(llvm::dbgs() << "[cf] cleaning up blocks after clone\n");
-  // First we need to drop all operands' references inside all blocks. This is
-  // needed because we can have blocks referencing SSA values from one another.
-  for (auto *block : constructBlocks)
-    block->dropAllReferences();
-
-  // Then erase all old blocks.
-  for (auto *block : constructBlocks) {
-    // We've cloned all blocks belonging to this construct into the structured
-    // control flow op's region. Among these blocks, some may compose another
-    // selection/loop. If so, they will be recorded within blockMergeInfo.
-    // We need to update the pointers there to the newly remapped ones so we can
-    // continue structurizing them later.
-    // TODO(antiagainst): The asserts in the following assumes input SPIR-V blob
-    // forms correctly nested selection/loop constructs. We should relax this
-    // and support error cases better.
-    auto it = blockMergeInfo.find(block);
-    if (it != blockMergeInfo.end()) {
-      Block *newHeader = mapper.lookupOrNull(block);
-      assert(newHeader && "nested loop header block should be remapped!");
-
-      Block *newContinue = it->second.continueBlock;
-      if (newContinue) {
-        newContinue = mapper.lookupOrNull(newContinue);
-        assert(newContinue && "nested loop continue block should be remapped!");
-      }
-
-      Block *newMerge = it->second.mergeBlock;
-      if (Block *mappedTo = mapper.lookupOrNull(newMerge))
-        newMerge = mappedTo;
-
-      // Keep original location for nested selection/loop ops.
-      Location loc = it->second.loc;
-      // The iterator should be erased before adding a new entry into
-      // blockMergeInfo to avoid iterator invalidation.
-      blockMergeInfo.erase(it);
-      blockMergeInfo.try_emplace(newHeader, loc, newMerge, newContinue);
-    }
-
-    // The structured selection/loop's entry block does not have arguments.
-    // If the function's header block is also part of the structured control
-    // flow, we cannot just simply erase it because it may contain arguments
-    // matching the function signature and used by the cloned blocks.
-    if (isFnEntryBlock(block)) {
-      LLVM_DEBUG(llvm::dbgs() << "[cf] changing entry block " << block
-                              << " to only contain a spv.Branch op\n");
-      // Still keep the function entry block for the potential block arguments,
-      // but replace all ops inside with a branch to the merge block.
-      block->clear();
-      builder.setInsertionPointToEnd(block);
-      builder.create<spirv::BranchOp>(location, mergeBlock);
-    } else {
-      LLVM_DEBUG(llvm::dbgs() << "[cf] erasing block " << block << "\n");
-      block->erase();
-    }
-  }
-
-  LLVM_DEBUG(
-      llvm::dbgs() << "[cf] after structurizing construct with header block "
-                   << headerBlock << ":\n"
-                   << *op << '\n');
-
-  return success();
-}
-
-LogicalResult Deserializer::wireUpBlockArgument() {
-  LLVM_DEBUG(llvm::dbgs() << "[phi] start wiring up block arguments\n");
-
-  OpBuilder::InsertionGuard guard(opBuilder);
-
-  for (const auto &info : blockPhiInfo) {
-    Block *block = info.first;
-    const BlockPhiInfo &phiInfo = info.second;
-    LLVM_DEBUG(llvm::dbgs() << "[phi] block " << block << "\n");
-    LLVM_DEBUG(llvm::dbgs() << "[phi] before creating block argument:\n");
-    LLVM_DEBUG(block->getParentOp()->print(llvm::dbgs()));
-    LLVM_DEBUG(llvm::dbgs() << '\n');
-
-    // Set insertion point to before this block's terminator early because we
-    // may materialize ops via getValue() call.
-    auto *op = block->getTerminator();
-    opBuilder.setInsertionPoint(op);
-
-    SmallVector<Value, 4> blockArgs;
-    blockArgs.reserve(phiInfo.size());
-    for (uint32_t valueId : phiInfo) {
-      if (Value value = getValue(valueId)) {
-        blockArgs.push_back(value);
-        LLVM_DEBUG(llvm::dbgs() << "[phi] block argument " << value
-                                << " id = " << valueId << '\n');
-      } else {
-        return emitError(unknownLoc, "OpPhi references undefined value!");
-      }
-    }
-
-    if (auto branchOp = dyn_cast<spirv::BranchOp>(op)) {
-      // Replace the previous branch op with a new one with block arguments.
-      opBuilder.create<spirv::BranchOp>(branchOp.getLoc(), branchOp.getTarget(),
-                                        blockArgs);
-      branchOp.erase();
-    } else {
-      return emitError(unknownLoc, "unimplemented terminator for Phi creation");
-    }
-
-    LLVM_DEBUG(llvm::dbgs() << "[phi] after creating block argument:\n");
-    LLVM_DEBUG(block->getParentOp()->print(llvm::dbgs()));
-    LLVM_DEBUG(llvm::dbgs() << '\n');
-  }
-  blockPhiInfo.clear();
-
-  LLVM_DEBUG(llvm::dbgs() << "[phi] completed wiring up block arguments\n");
-  return success();
-}
-
-LogicalResult Deserializer::structurizeControlFlow() {
-  LLVM_DEBUG(llvm::dbgs() << "[cf] start structurizing control flow\n");
-
-  while (!blockMergeInfo.empty()) {
-    Block *headerBlock = blockMergeInfo.begin()->first;
-    BlockMergeInfo mergeInfo = blockMergeInfo.begin()->second;
-
-    LLVM_DEBUG(llvm::dbgs() << "[cf] header block " << headerBlock << ":\n");
-    LLVM_DEBUG(headerBlock->print(llvm::dbgs()));
-
-    auto *mergeBlock = mergeInfo.mergeBlock;
-    assert(mergeBlock && "merge block cannot be nullptr");
-    if (!mergeBlock->args_empty())
-      return emitError(unknownLoc, "OpPhi in loop merge block unimplemented");
-    LLVM_DEBUG(llvm::dbgs() << "[cf] merge block " << mergeBlock << ":\n");
-    LLVM_DEBUG(mergeBlock->print(llvm::dbgs()));
-
-    auto *continueBlock = mergeInfo.continueBlock;
-    if (continueBlock) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "[cf] continue block " << continueBlock << ":\n");
-      LLVM_DEBUG(continueBlock->print(llvm::dbgs()));
-    }
-    // Erase this case before calling into structurizer, who will update
-    // blockMergeInfo.
-    blockMergeInfo.erase(blockMergeInfo.begin());
-    if (failed(ControlFlowStructurizer::structurize(mergeInfo.loc,
-                                                    blockMergeInfo, headerBlock,
-                                                    mergeBlock, continueBlock)))
-      return failure();
-  }
-
-  LLVM_DEBUG(llvm::dbgs() << "[cf] completed structurizing control flow\n");
-  return success();
-}
-
-//===----------------------------------------------------------------------===//
-// Debug
-//===----------------------------------------------------------------------===//
-
-Location Deserializer::createFileLineColLoc(OpBuilder opBuilder) {
-  if (!debugLine)
-    return unknownLoc;
-
-  auto fileName = debugInfoMap.lookup(debugLine->fileID).str();
-  if (fileName.empty())
-    fileName = "<unknown>";
-  return opBuilder.getFileLineColLoc(opBuilder.getIdentifier(fileName),
-                                     debugLine->line, debugLine->col);
-}
-
-LogicalResult Deserializer::processDebugLine(ArrayRef<uint32_t> operands) {
-  // According to SPIR-V spec:
-  // "This location information applies to the instructions physically
-  // following this instruction, up to the first occurrence of any of the
-  // following: the next end of block, the next OpLine instruction, or the next
-  // OpNoLine instruction."
-  if (operands.size() != 3)
-    return emitError(unknownLoc, "OpLine must have 3 operands");
-  debugLine = DebugLine(operands[0], operands[1], operands[2]);
-  return success();
-}
-
-LogicalResult Deserializer::clearDebugLine() {
-  debugLine = llvm::None;
-  return success();
-}
-
-LogicalResult Deserializer::processDebugString(ArrayRef<uint32_t> operands) {
-  if (operands.size() < 2)
-    return emitError(unknownLoc, "OpString needs at least 2 operands");
-
-  if (!debugInfoMap.lookup(operands[0]).empty())
-    return emitError(unknownLoc,
-                     "duplicate debug string found for result <id> ")
-           << operands[0];
-
-  unsigned wordIndex = 1;
-  StringRef debugString = decodeStringLiteral(operands, wordIndex);
-  if (wordIndex != operands.size())
-    return emitError(unknownLoc,
-                     "unexpected trailing words in OpString instruction");
-
-  debugInfoMap[operands[0]] = debugString;
-  return success();
-}
-
-//===----------------------------------------------------------------------===//
-// Instruction
-//===----------------------------------------------------------------------===//
-
-Value Deserializer::getValue(uint32_t id) {
-  if (auto constInfo = getConstant(id)) {
-    // Materialize a `spv.constant` op at every use site.
-    return opBuilder.create<spirv::ConstantOp>(unknownLoc, constInfo->second,
-                                               constInfo->first);
-  }
-  if (auto varOp = getGlobalVariable(id)) {
-    auto addressOfOp = opBuilder.create<spirv::AddressOfOp>(
-        unknownLoc, varOp.type(),
-        opBuilder.getSymbolRefAttr(varOp.getOperation()));
-    return addressOfOp.pointer();
-  }
-  if (auto constOp = getSpecConstant(id)) {
-    auto referenceOfOp = opBuilder.create<spirv::ReferenceOfOp>(
-        unknownLoc, constOp.default_value().getType(),
-        opBuilder.getSymbolRefAttr(constOp.getOperation()));
-    return referenceOfOp.reference();
-  }
-  if (auto undef = getUndefType(id)) {
-    return opBuilder.create<spirv::UndefOp>(unknownLoc, undef);
-  }
-  return valueMap.lookup(id);
-}
-
-LogicalResult
-Deserializer::sliceInstruction(spirv::Opcode &opcode,
-                               ArrayRef<uint32_t> &operands,
-                               Optional<spirv::Opcode> expectedOpcode) {
-  auto binarySize = binary.size();
-  if (curOffset >= binarySize) {
-    return emitError(unknownLoc, "expected ")
-           << (expectedOpcode ? spirv::stringifyOpcode(*expectedOpcode)
-                              : "more")
-           << " instruction";
-  }
-
-  // For each instruction, get its word count from the first word to slice it
-  // from the stream properly, and then dispatch to the instruction handler.
-
-  uint32_t wordCount = binary[curOffset] >> 16;
-
-  if (wordCount == 0)
-    return emitError(unknownLoc, "word count cannot be zero");
-
-  uint32_t nextOffset = curOffset + wordCount;
-  if (nextOffset > binarySize)
-    return emitError(unknownLoc, "insufficient words for the last instruction");
-
-  opcode = extractOpcode(binary[curOffset]);
-  operands = binary.slice(curOffset + 1, wordCount - 1);
-  curOffset = nextOffset;
-  return success();
-}
-
-LogicalResult Deserializer::processInstruction(spirv::Opcode opcode,
-                                               ArrayRef<uint32_t> operands,
-                                               bool deferInstructions) {
-  LLVM_DEBUG(llvm::dbgs() << "[inst] processing instruction "
-                          << spirv::stringifyOpcode(opcode) << "\n");
-
-  // First dispatch all the instructions whose opcode does not correspond to
-  // those that have a direct mirror in the SPIR-V dialect
-  switch (opcode) {
-  case spirv::Opcode::OpCapability:
-    return processCapability(operands);
-  case spirv::Opcode::OpExtension:
-    return processExtension(operands);
-  case spirv::Opcode::OpExtInst:
-    return processExtInst(operands);
-  case spirv::Opcode::OpExtInstImport:
-    return processExtInstImport(operands);
-  case spirv::Opcode::OpMemberName:
-    return processMemberName(operands);
-  case spirv::Opcode::OpMemoryModel:
-    return processMemoryModel(operands);
-  case spirv::Opcode::OpEntryPoint:
-  case spirv::Opcode::OpExecutionMode:
-    if (deferInstructions) {
-      deferredInstructions.emplace_back(opcode, operands);
-      return success();
-    }
-    break;
-  case spirv::Opcode::OpVariable:
-    if (isa<spirv::ModuleOp>(opBuilder.getBlock()->getParentOp())) {
-      return processGlobalVariable(operands);
-    }
-    break;
-  case spirv::Opcode::OpLine:
-    return processDebugLine(operands);
-  case spirv::Opcode::OpNoLine:
-    return clearDebugLine();
-  case spirv::Opcode::OpName:
-    return processName(operands);
-  case spirv::Opcode::OpString:
-    return processDebugString(operands);
-  case spirv::Opcode::OpModuleProcessed:
-  case spirv::Opcode::OpSource:
-  case spirv::Opcode::OpSourceContinued:
-  case spirv::Opcode::OpSourceExtension:
-    // TODO: This is debug information embedded in the binary which should be
-    // translated into the spv.module.
-    return success();
-  case spirv::Opcode::OpTypeVoid:
-  case spirv::Opcode::OpTypeBool:
-  case spirv::Opcode::OpTypeInt:
-  case spirv::Opcode::OpTypeFloat:
-  case spirv::Opcode::OpTypeVector:
-  case spirv::Opcode::OpTypeMatrix:
-  case spirv::Opcode::OpTypeArray:
-  case spirv::Opcode::OpTypeFunction:
-  case spirv::Opcode::OpTypeRuntimeArray:
-  case spirv::Opcode::OpTypeStruct:
-  case spirv::Opcode::OpTypePointer:
-  case spirv::Opcode::OpTypeCooperativeMatrixNV:
-    return processType(opcode, operands);
-  case spirv::Opcode::OpConstant:
-    return processConstant(operands, /*isSpec=*/false);
-  case spirv::Opcode::OpSpecConstant:
-    return processConstant(operands, /*isSpec=*/true);
-  case spirv::Opcode::OpConstantComposite:
-    return processConstantComposite(operands);
-  case spirv::Opcode::OpConstantTrue:
-    return processConstantBool(/*isTrue=*/true, operands, /*isSpec=*/false);
-  case spirv::Opcode::OpSpecConstantTrue:
-    return processConstantBool(/*isTrue=*/true, operands, /*isSpec=*/true);
-  case spirv::Opcode::OpConstantFalse:
-    return processConstantBool(/*isTrue=*/false, operands, /*isSpec=*/false);
-  case spirv::Opcode::OpSpecConstantFalse:
-    return processConstantBool(/*isTrue=*/false, operands, /*isSpec=*/true);
-  case spirv::Opcode::OpConstantNull:
-    return processConstantNull(operands);
-  case spirv::Opcode::OpDecorate:
-    return processDecoration(operands);
-  case spirv::Opcode::OpMemberDecorate:
-    return processMemberDecoration(operands);
-  case spirv::Opcode::OpFunction:
-    return processFunction(operands);
-  case spirv::Opcode::OpLabel:
-    return processLabel(operands);
-  case spirv::Opcode::OpBranch:
-    return processBranch(operands);
-  case spirv::Opcode::OpBranchConditional:
-    return processBranchConditional(operands);
-  case spirv::Opcode::OpSelectionMerge:
-    return processSelectionMerge(operands);
-  case spirv::Opcode::OpLoopMerge:
-    return processLoopMerge(operands);
-  case spirv::Opcode::OpPhi:
-    return processPhi(operands);
-  case spirv::Opcode::OpUndef:
-    return processUndef(operands);
-  default:
-    break;
-  }
-  return dispatchToAutogenDeserialization(opcode, operands);
-}
-
-LogicalResult Deserializer::processUndef(ArrayRef<uint32_t> operands) {
-  if (operands.size() != 2) {
-    return emitError(unknownLoc, "OpUndef instruction must have two operands");
-  }
-  auto type = getType(operands[0]);
-  if (!type) {
-    return emitError(unknownLoc, "unknown type <id> with OpUndef instruction");
-  }
-  undefMap[operands[1]] = type;
-  return success();
-}
-
-LogicalResult Deserializer::processExtInst(ArrayRef<uint32_t> operands) {
-  if (operands.size() < 4) {
-    return emitError(unknownLoc,
-                     "OpExtInst must have at least 4 operands, result type "
-                     "<id>, result <id>, set <id> and instruction opcode");
-  }
-  if (!extendedInstSets.count(operands[2])) {
-    return emitError(unknownLoc, "undefined set <id> in OpExtInst");
-  }
-  SmallVector<uint32_t, 4> slicedOperands;
-  slicedOperands.append(operands.begin(), std::next(operands.begin(), 2));
-  slicedOperands.append(std::next(operands.begin(), 4), operands.end());
-  return dispatchToExtensionSetAutogenDeserialization(
-      extendedInstSets[operands[2]], operands[3], slicedOperands);
-}
-
-namespace {
-
-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++]);
-  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
-  auto parsedFunc = getFunction(fnID);
-  if (!parsedFunc) {
-    return emitError(unknownLoc, "no function matching <id> ") << fnID;
-  }
-  if (parsedFunc.getName() != fnName) {
-    return emitError(unknownLoc, "function name mismatch between OpEntryPoint "
-                                 "and OpFunction with <id> ")
-           << fnID << ": " << fnName << " vs. " << parsedFunc.getName();
-  }
-  SmallVector<Attribute, 4> interface;
-  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.getSymbolRefAttr(fnName),
-                                        opBuilder.getArrayAttr(interface));
-  return success();
-}
-
-template <>
-LogicalResult
-Deserializer::processOp<spirv::ExecutionModeOp>(ArrayRef<uint32_t> words) {
-  unsigned wordIndex = 0;
-  if (wordIndex >= words.size()) {
-    return emitError(unknownLoc,
-                     "missing function result <id> in OpExecutionMode");
-  }
-  // Get the function <id> to get the name of the function
-  auto fnID = words[wordIndex++];
-  auto fn = getFunction(fnID);
-  if (!fn) {
-    return emitError(unknownLoc, "no function matching <id> ") << fnID;
-  }
-  // Get the Execution mode
-  if (wordIndex >= words.size()) {
-    return emitError(unknownLoc, "missing Execution Mode in OpExecutionMode");
-  }
-  auto execMode = opBuilder.getI32IntegerAttr(words[wordIndex++]);
-
-  // Get the values
-  SmallVector<Attribute, 4> attrListElems;
-  while (wordIndex < words.size()) {
-    attrListElems.push_back(opBuilder.getI32IntegerAttr(words[wordIndex++]));
-  }
-  auto values = opBuilder.getArrayAttr(attrListElems);
-  opBuilder.create<spirv::ExecutionModeOp>(
-      unknownLoc, opBuilder.getSymbolRefAttr(fn.getName()), execMode, values);
-  return success();
-}
-
-template <>
-LogicalResult
-Deserializer::processOp<spirv::ControlBarrierOp>(ArrayRef<uint32_t> operands) {
-  if (operands.size() != 3) {
-    return emitError(
-        unknownLoc,
-        "OpControlBarrier must have execution scope <id>, memory scope <id> "
-        "and memory semantics <id>");
-  }
-
-  SmallVector<IntegerAttr, 3> argAttrs;
-  for (auto operand : operands) {
-    auto argAttr = getConstantInt(operand);
-    if (!argAttr) {
-      return emitError(unknownLoc,
-                       "expected 32-bit integer constant from <id> ")
-             << operand << " for OpControlBarrier";
-    }
-    argAttrs.push_back(argAttr);
-  }
-
-  opBuilder.create<spirv::ControlBarrierOp>(unknownLoc, argAttrs[0],
-                                            argAttrs[1], argAttrs[2]);
-  return success();
-}
-
-template <>
-LogicalResult
-Deserializer::processOp<spirv::FunctionCallOp>(ArrayRef<uint32_t> operands) {
-  if (operands.size() < 3) {
-    return emitError(unknownLoc,
-                     "OpFunctionCall must have at least 3 operands");
-  }
-
-  Type resultType = getType(operands[0]);
-  if (!resultType) {
-    return emitError(unknownLoc, "undefined result type from <id> ")
-           << operands[0];
-  }
-
-  // Use null type to mean no result type.
-  if (isVoidType(resultType))
-    resultType = nullptr;
-
-  auto resultID = operands[1];
-  auto functionID = operands[2];
-
-  auto functionName = getFunctionSymbol(functionID);
-
-  SmallVector<Value, 4> arguments;
-  for (auto operand : llvm::drop_begin(operands, 3)) {
-    auto value = getValue(operand);
-    if (!value) {
-      return emitError(unknownLoc, "unknown <id> ")
-             << operand << " used by OpFunctionCall";
-    }
-    arguments.push_back(value);
-  }
-
-  auto opFunctionCall = opBuilder.create<spirv::FunctionCallOp>(
-      unknownLoc, resultType, opBuilder.getSymbolRefAttr(functionName),
-      arguments);
-
-  if (resultType)
-    valueMap[resultID] = opFunctionCall.getResult(0);
-  return success();
-}
-
-template <>
-LogicalResult
-Deserializer::processOp<spirv::MemoryBarrierOp>(ArrayRef<uint32_t> operands) {
-  if (operands.size() != 2) {
-    return emitError(unknownLoc, "OpMemoryBarrier must have memory scope <id> "
-                                 "and memory semantics <id>");
-  }
-
-  SmallVector<IntegerAttr, 2> argAttrs;
-  for (auto operand : operands) {
-    auto argAttr = getConstantInt(operand);
-    if (!argAttr) {
-      return emitError(unknownLoc,
-                       "expected 32-bit integer constant from <id> ")
-             << operand << " for OpMemoryBarrier";
-    }
-    argAttrs.push_back(argAttr);
-  }
-
-  opBuilder.create<spirv::MemoryBarrierOp>(unknownLoc, argAttrs[0],
-                                           argAttrs[1]);
-  return success();
-}
-
-// Pull in auto-generated Deserializer::dispatchToAutogenDeserialization() and
-// various Deserializer::processOp<...>() specializations.
-#define GET_DESERIALIZATION_FNS
-#include "mlir/Dialect/SPIRV/SPIRVSerialization.inc"
-} // namespace
-
-Optional<spirv::ModuleOp> spirv::deserialize(ArrayRef<uint32_t> binary,
-                                             MLIRContext *context) {
-  Deserializer deserializer(binary, context);
-
-  if (failed(deserializer.deserialize()))
-    return llvm::None;
-
-  return deserializer.collect();
-}
+  for (unsigned i

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

 //===- TranslateRegistration.cpp - hooks to mlir-translate ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a translation from SPIR-V binary module to MLIR SPIR-V
 // ModuleOp.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/SPIRV/SPIRVModule.h"
 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/Serialization.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/Parser.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Translation.h"
if (size % sizeof(uint32_t) != 0) {
     emitError(UnknownLoc::get(context))
         << "SPIR-V binary module must contain integral number of 32-bit words";
     return {};
   }
 
   auto binary = llvm::makeArrayRef(reinterpret_cast<const uint32_t *>(start),
                                    size / sizeof(uint32_t));
 
-  auto spirvModule = spirv::deserialize(binary, context);
+  spirv::OwningSPIRVModuleRef spirvModule = spirv::deserialize(binary, context);
   if (!spirvModule)
     return {};
 
   OwningModuleRef module(ModuleOp::create(FileLineColLoc::get(
       input->getBufferIdentifier(), /*line=*/0, /*column=*/0, context)));
-  module->getBody()->push_front(spirvModule->getOperation());
+  module->getBody()->push_front(spirvModule.release());
 
   return module;
 }
 
 namespace mlir {
 void registerFromSPIRVTranslation() {
   TranslateToMLIRRegistration fromBinary(
       "deserialize-spirv",
static LogicalResult roundTripModule(llvm::SourceMgr &sourceMgr,
 
   if (std::next(spirvModules.begin()) != spirvModules.end())
     return srcModule->emitError("found more than one 'spv.module' op");
 
   if (failed(spirv::serialize(*spirvModules.begin(), binary, emitDebugInfo)))
     return failure();
 
   // Then deserialize to get back a SPIR-V module.
-  auto spirvModule = spirv::deserialize(binary, context);
+  spirv::OwningSPIRVModuleRef spirvModule = spirv::deserialize(binary, context);
   if (!spirvModule)
     return failure();
 
   // Wrap around in a new MLIR module.
   OwningModuleRef dstModule(ModuleOp::create(FileLineColLoc::get(
       /*filename=*/"", /*line=*/0, /*column=*/0, context)));
-  dstModule->getBody()->push_front(spirvModule->getOperation());
+  dstModule->getBody()->push_front(spirvModule.release());
   dstModule->print(output);
 
   return mlir::success();
 }
 
 namespace mlir {
 void registerTestRoundtripSPIRV() {
   TranslateRegistration roundtrip(

mlir/unittests/Dialect/SPIRV/DeserializationTest.cpp

 //===- DeserializationTest.cpp - SPIR-V Deserialization Tests -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // The purpose of this file is to provide negative deserialization tests.
 // For positive deserialization tests, please use serialization and
 // deserialization for roundtripping.
 //
 //===----------------------------------------------------------------------===//
 
 #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/Serialization.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/MLIRContext.h"
 #include "gmock/gmock.h"
 
 #include <memory>
 
DeserializationTest() {
     // Register a diagnostic handler to capture the diagnostic so that we can
     // check it later.
     context.getDiagEngine().registerHandler([&](Diagnostic &diag) {
       diagnostic.reset(new Diagnostic(std::move(diag)));
     });
   }
 
   /// Performs deserialization and returns the constructed spv.module op.
-  Optional<spirv::ModuleOp> deserialize() {
+  spirv::OwningSPIRVModuleRef deserialize() {
     return spirv::deserialize(binary, &context);
   }
 
   /// Checks there is a diagnostic generated with the given `errorMessage`.
   void expectDiagnostic(StringRef errorMessage) {
     ASSERT_NE(nullptr, diagnostic.get());
 
     // TODO(antiagainst): check error location too.
protected:
   std::unique_ptr<Diagnostic> diagnostic;
 };
 
 //===----------------------------------------------------------------------===//
 // Basics
 //===----------------------------------------------------------------------===//
 
 TEST_F(DeserializationTest, EmptyModuleFailure) {
-  ASSERT_EQ(llvm::None, deserialize());
+  ASSERT_FALSE(deserialize());
   expectDiagnostic("SPIR-V binary module must have a 5-word header");
 }
 
 TEST_F(DeserializationTest, WrongMagicNumberFailure) {
   addHeader();
   binary.front() = 0xdeadbeef; // Change to a wrong magic number
-  ASSERT_EQ(llvm::None, deserialize());
+  ASSERT_FALSE(deserialize());
   expectDiagnostic("incorrect magic number");
 }
 
 TEST_F(DeserializationTest, OnlyHeaderSuccess) {
   addHeader();
-  EXPECT_NE(llvm::None, deserialize());
+  EXPECT_TRUE(deserialize());
 }
 
 TEST_F(DeserializationTest, ZeroWordCountFailure) {
   addHeader();
   binary.push_back(0); // OpNop with zero word count
 
-  ASSERT_EQ(llvm::None, deserialize());
+  ASSERT_FALSE(deserialize());
   expectDiagnostic("word count cannot be zero");
 }
 
 TEST_F(DeserializationTest, InsufficientWordFailure) {
   addHeader();
   binary.push_back((2u << 16) |
                    static_cast<uint32_t>(spirv::Opcode::OpTypeVoid));
   // Missing word for type <id>
 
-  ASSERT_EQ(llvm::None, deserialize());
+  ASSERT_FALSE(deserialize());
   expectDiagnostic("insufficient words for the last instruction");
 }
 
 //===----------------------------------------------------------------------===//
 // Types
 //===----------------------------------------------------------------------===//
 
 TEST_F(DeserializationTest, IntTypeMissingSignednessFailure) {
   addHeader();
   addInstruction(spirv::Opcode::OpTypeInt, {nextID++, 32});
 
-  ASSERT_EQ(llvm::None, deserialize());
+  ASSERT_FALSE(deserialize());
   expectDiagnostic("OpTypeInt must have bitwidth and signedness parameters");
 }
 
 //===----------------------------------------------------------------------===//
 // StructType
 //===----------------------------------------------------------------------===//
 
 TEST_F(DeserializationTest, OpMemberNameSuccess) {
TEST_F(DeserializationTest, OpMemberNameSuccess) {
   spirv::encodeStringLiteralInto(operands1, "i1");
   addInstruction(spirv::Opcode::OpMemberName, operands1);
 
   SmallVector<uint32_t, 5> operands2 = {structType, 1};
   spirv::encodeStringLiteralInto(operands2, "i2");
   addInstruction(spirv::Opcode::OpMemberName, operands2);
 
   binary.append(typeDecl.begin(), typeDecl.end());
-  EXPECT_NE(llvm::None, deserialize());
+  EXPECT_TRUE(deserialize());
 }
 
 TEST_F(DeserializationTest, OpMemberNameMissingOperands) {
   addHeader();
   SmallVector<uint32_t, 5> typeDecl;
   std::swap(typeDecl, binary);
 
   auto int32Type = addIntType(32);
   auto int64Type = addIntType(64);
   auto structType = addStructType({int32Type, int64Type});
   std::swap(typeDecl, binary);
 
   SmallVector<uint32_t, 5> operands1 = {structType};
   addInstruction(spirv::Opcode::OpMemberName, operands1);
 
   binary.append(typeDecl.begin(), typeDecl.end());
-  ASSERT_EQ(llvm::None, deserialize());
+  ASSERT_FALSE(deserialize());
   expectDiagnostic("OpMemberName must have at least 3 operands");
 }
 
 TEST_F(DeserializationTest, OpMemberNameExcessOperands) {
   addHeader();
   SmallVector<uint32_t, 5> typeDecl;
   std::swap(typeDecl, binary);
 
   auto int32Type = addIntType(32);
   auto structType = addStructType({int32Type});
   std::swap(typeDecl, binary);
 
   SmallVector<uint32_t, 5> operands = {structType, 0};
   spirv::encodeStringLiteralInto(operands, "int32");
   operands.push_back(42);
   addInstruction(spirv::Opcode::OpMemberName, operands);
 
   binary.append(typeDecl.begin(), typeDecl.end());
-  ASSERT_EQ(llvm::None, deserialize());
+  ASSERT_FALSE(deserialize());
   expectDiagnostic("unexpected trailing words in OpMemberName instruction");
 }
 
 //===----------------------------------------------------------------------===//
 // Functions
 //===----------------------------------------------------------------------===//
 
 TEST_F(DeserializationTest, FunctionMissingEndFailure) {
   addHeader();
   auto voidType = addVoidType();
   auto fnType = addFunctionType(voidType, {});
   addFunction(voidType, fnType);
   // Missing OpFunctionEnd
 
-  ASSERT_EQ(llvm::None, deserialize());
+  ASSERT_FALSE(deserialize());
   expectDiagnostic("expected OpFunctionEnd instruction");
 }
 
 TEST_F(DeserializationTest, FunctionMissingParameterFailure) {
   addHeader();
   auto voidType = addVoidType();
   auto i32Type = addIntType(32);
   auto fnType = addFunctionType(voidType, {i32Type});
   addFunction(voidType, fnType);
   // Missing OpFunctionParameter
 
-  ASSERT_EQ(llvm::None, deserialize());
+  ASSERT_FALSE(deserialize());
   expectDiagnostic("expected OpFunctionParameter instruction");
 }
 
 TEST_F(DeserializationTest, FunctionMissingLabelForFirstBlockFailure) {
   addHeader();
   auto voidType = addVoidType();
   auto fnType = addFunctionType(voidType, {});
   addFunction(voidType, fnType);
   // Missing OpLabel
   addReturn();
   addFunctionEnd();
 
-  ASSERT_EQ(llvm::None, deserialize());
+  ASSERT_FALSE(deserialize());
   expectDiagnostic("a basic block must start with OpLabel");
 }
 
 TEST_F(DeserializationTest, FunctionMalformedLabelFailure) {
   addHeader();
   auto voidType = addVoidType();
   auto fnType = addFunctionType(voidType, {});
   addFunction(voidType, fnType);
   addInstruction(spirv::Opcode::OpLabel, {}); // Malformed OpLabel
   addReturn();
   addFunctionEnd();
 
-  ASSERT_EQ(llvm::None, deserialize());
+  ASSERT_FALSE(deserialize());
   expectDiagnostic("OpLabel should only have result <id>");
 }