[mlir] Introduce OwningFuncRef following the same pattern as OwningModuleRef

[tberghammer]

Summary

Between the creating of a function op and adding it to a module it has no
explicit owner what can lead to memory leaks. The new class is intended to be
used as a convenient RAII owner for that time.

This change generalizes the code for OwningModuleRef to support arbitrary Op
types and makes OwningModuleRef and OwningFuncRef a simple shim around the
generic class to still enable forward declaring these types.

[tberghammer]

Fix code formatting

[tberghammer]

Add missing include

[rriddle]

Hi Tamas, thanks for the patch. ModuleOp is a bit special in that it is the top-level IR entity. What situation are you in where you can't insert the FuncOp directly on creation?

[tberghammer]

Hi River, I have two (somewhat related) usecases for having a FuncOp without adding it to a module:

1. I would like to write an mlir::OwingFuncRef BuildFunc(...) function what takes some internal/domain specific input (e.g. IR from the front end) and returns the newly built function. I could just pass in the ModuleOp as an argument and add the FuncOp to it but returning an OwningFuncRef feels like a nicer API
2. My code what builds a FuncOp might fail (e.g. due to invalid input) and in that case I want to report an error without modifying the existing ModuleOp as adding a broken function to it would make the module op broken as well. My approach for it is to create an OwningFuncRef, use it to own the FuncOp during the building process (so in case of a failure it will correctly destroy it) and then in case of success I would move the FuncOp from the OwningFuncRef to the ModuleOp.

[rriddle]

In D75649#1914044, @tberghammer wrote:

Hi River, I have two (somewhat related) usecases for having a FuncOp without adding it to a module:

1. I would like to write an mlir::OwingFuncRef BuildFunc(...) function what takes some internal/domain specific input (e.g. IR from the front end) and returns the newly built function. I could just pass in the ModuleOp as an argument and add the FuncOp to it but returning an OwningFuncRef feels like a nicer API
2. My code what builds a FuncOp might fail (e.g. due to invalid input) and in that case I want to report an error without modifying the existing ModuleOp as adding a broken function to it would make the module op broken as well. My approach for it is to create an OwningFuncRef, use it to own the FuncOp during the building process (so in case of a failure it will correctly destroy it) and then in case of success I would move the FuncOp from the OwningFuncRef to the ModuleOp.

I would rather not encourage this type of usage. FuncOp::create is something that I've been meaning to remove for a long time, but I keep getting distracted. It is a relic of the time when FuncOp was Function, and not represented as an Operation. You also don't necessarily need to pass in a ModuleOp, using OpBuilder::create is much more idiomatic with the rest of the infra.

[mehdi_amini]

using OpBuilder::create is much more idiomatic with the rest of the infra.

The builder APIs don't express clear ownership though, having a RAII type would be nice!
Could the OpBuilder::create always return an OwningOpRef<OpTy> ?

[rriddle]

In D75649#1916222, @mehdi_amini wrote:

using OpBuilder::create is much more idiomatic with the rest of the infra.

The builder APIs don't express clear ownership though, having a RAII type would be nice!
Could the OpBuilder::create always return an OwningOpRef<OpTy> ?

That doesn't really make sense to me, OpBuilder::create generally inserts the operation into a block. I'm not sure what ownership semantics you are assigning to it here. Am I missing something? OpBuilder::create isn't a static method.

[tberghammer]

In D75649#1916265, @rriddle wrote:

In D75649#1916222, @mehdi_amini wrote:

using OpBuilder::create is much more idiomatic with the rest of the infra.

The builder APIs don't express clear ownership though, having a RAII type would be nice!
Could the OpBuilder::create always return an OwningOpRef<OpTy> ?

That doesn't really make sense to me, OpBuilder::create generally inserts the operation into a block. I'm not sure what ownership semantics you are assigning to it here. Am I missing something? OpBuilder::create isn't a static method.

I think there is still value in being able to create a FuncOp and any other "block type" operation without adding it to the parent block for error reporting and error recovery. Also I can think of cases where you don't know the parent block at the time when we started to create a block (e.g. we add the FuncOp to a different module based on some property we derive later).

A valid point from Mehdi is that we should change FuncOp::create to return OwningFuncRef so it makes the semantics clear but I would prefer to do that separately as it touches quite a few files.

P.S.: Looking at mlir/examples/toy/Ch3/mlir/MLIRGen.cpp it uses the exact pattern I want to avoid where in case of a failure we have to manually erase the FuncOp to avoid a memory leak.

[mehdi_amini]

In D75649#1916265, @rriddle wrote:

That doesn't really make sense to me, OpBuilder::create generally inserts the operation into a block. I'm not sure what ownership semantics you are assigning to it here. Am I missing something? OpBuilder::create isn't a static method.

Ah of course you're right! I was thinking about Operation::create actually :)
It is less often used directly though because of OpBuilder...

mlir/include/mlir/IR/Function.h

 
 #ifndef MLIR_IR_FUNCTION_H
 #define MLIR_IR_FUNCTION_H
 
 #include "mlir/Analysis/CallInterfaces.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OwningOpRef.h"
 #include "mlir/IR/SymbolTable.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
 //===--------------------------------------------------------------------===//
 // Function Operation.
 //===--------------------------------------------------------------------===//
 
private:
   LogicalResult verifyType() {
     auto type = getTypeAttr().getValue();
     if (!type.isa<FunctionType>())
       return emitOpError("requires '" + getTypeAttrName() +
                          "' attribute of function type");
     return success();
   }
 };
+
+/// This class acts as an owning reference to a function, and will automatically
+/// destroy the held function if valid.
+class OwningFuncRef : public OwningOpRef<FuncOp> {
+ public:
+  using OwningOpRef::OwningOpRef;
+};
+
 } // end namespace mlir
 
 namespace llvm {
 
 // Functions hash just like pointers.
 template <> struct DenseMapInfo<mlir::FuncOp> {
   static mlir::FuncOp getEmptyKey() {
     auto pointer = llvm::DenseMapInfo<void *>::getEmptyKey();

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/SymbolTable.h"
+#include "mlir/IR/OwningOpRef.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
 class ModuleTerminatorOp;
 
 //===----------------------------------------------------------------------===//
 // Module Operation.
 //===----------------------------------------------------------------------===//

 public:
   using Op::Op;
   static StringRef getOperationName() { return "module_terminator"; }
   static void build(Builder *, OperationState &) {}
 };
 
 /// This class acts as an owning reference to a module, and will automatically
 /// destroy the held module if valid.
-class OwningModuleRef {
+class OwningModuleRef : public OwningOpRef<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 OwningOpRef::OwningOpRef;
 };
 
 } // end namespace mlir
 
 namespace llvm {
 
 /// Allow stealing the low bits of ModuleOp.
 template <> struct PointerLikeTypeTraits<mlir::ModuleOp> {

mlir/include/mlir/IR/OwningOpRef.h

+//===- OwningOpRef.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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_OWNING_OP_REF_H
+#define MLIR_IR_OWNING_OP_REF_H
+
+#include <cstddef>
+
+namespace mlir {
+
+/// This class acts as an owning reference to an op, and will automatically
+/// destroy the held op if valid.
+template <typename T>
+class OwningOpRef {
+public:
+  OwningOpRef(std::nullptr_t = nullptr) {}
+  OwningOpRef(T op) : op(op) {}
+  OwningOpRef(OwningOpRef &&other) : op(other.release()) {}
+  ~OwningOpRef() {
+    if (op)
+      op.erase();
+  }
+
+  // Assign from another op reference.
+  OwningOpRef &operator=(OwningOpRef &&other) {
+    if (op)
+      op.erase();
+    op = other.release();
+    return *this;
+  }
+
+  /// Allow accessing the internal op.
+  T get() const { return op; }
+  T operator*() const { return op; }
+  T *operator->() { return &op; }
+  explicit operator bool() const { return op; }
+
+  /// Release the referenced op.
+  T release() {
+    T released;
+    std::swap(released, op);
+    return released;
+  }
+
+private:
+  T op;
+};
+
+} // end namespace mlir
+
+#endif // MLIR_IR_OWNING_OP_REF_H