[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

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

nit: I would likely just rename to OwningOpRefBase instead.

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

2527

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/ModuleRefBase.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 OwningModuleRefBase<spirv::ModuleOp> {
+public:
+  using OwningModuleRefBase<spirv::ModuleOp>::OwningModuleRefBase;
+};
+
+} // 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/ModuleRefBase.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 OwningModuleRefBase<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 OwningModuleRefBase<ModuleOp>::OwningModuleRefBase;
 };
 
 } // end namespace mlir
 
 namespace llvm {
 
 /// Allow stealing the low bits of ModuleOp.
 template <> struct PointerLikeTypeTraits<mlir::ModuleOp> {

mlir/include/mlir/IR/ModuleRefBase.h

+//===- ModuleRefBase.h - MLIR OwningModuleRefBase ---------------*- 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 a module-like op, and will
+/// automatically destroy the held module on destruction if the held module is
+/// valid.
+template <typename ModuleTy>
+class OwningModuleRefBase {
+public:
+  OwningModuleRefBase(std::nullptr_t = nullptr) {}
+  OwningModuleRefBase(ModuleTy module) : module(module) {}
+  OwningModuleRefBase(OwningModuleRefBase &&other) : module(other.release()) {}
+  ~OwningModuleRefBase() {
+    if (module)
+      module.erase();
+  }
+
+  // Assign from another module reference.
+  OwningModuleRefBase &operator=(OwningModuleRefBase &&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"

 }
 
 // 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,
+spirv::OwningSPIRVModuleRef spirv::deserialize(ArrayRef<uint32_t> binary,
                                                MLIRContext *context) {
   Deserializer deserializer(binary, context);
 
   if (failed(deserializer.deserialize()))
-    return llvm::None;
+    return nullptr;
+
+  Optional<spirv::ModuleOp> module = deserializer.collect();
+  if (!module)
+    return nullptr;
 
-  return deserializer.collect();
+  return *module;
 }

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