[mlir] use storage uniquer for side effects

[ftynse]

Summary

The Effect class hierarchy in the side effect infrastructure currently only
supports non-parametric effects uniquely identified by their type. Only one
instance of each effect is stored in method and returned by-pointer upon
request. Hook the Effect class hierarchy to StorageUniquer instead. This makes
Effect and derived classes value-types, similarly to types and attributes, and
prepares the ground for parametric effets.

[mehdi_amini]

mlir/lib/IR/MLIRContext.cpp

27

I don't really understand the layering here: I'm not convinced that the Context should be made aware of this?

[ftynse]

mlir/lib/IR/MLIRContext.cpp

27

I am not happy with this either, this wouldn't scale to user-defined side effects. The other things that have uniquers are organized as follows:

• types and attributes are registered by the dialect;
• affine constructs (arguably they are part of the IR, not dialect) are registered here.

Any suggestions? @rriddle ?

[ftynse]

Replaced with another approach - https://reviews.llvm.org/D91493.

mlir/include/mlir/IR/MLIRContext.h

public:
   /// Returns the storage uniquer used for constructing type storage instances.
   /// This should not be used directly.
   StorageUniquer &getTypeUniquer();
 
   /// Returns the storage uniquer used for constructing attribute storage
   /// instances. This should not be used directly.
   StorageUniquer &getAttributeUniquer();
 
+  /// Returns the storage uniquer used for constructing uniqued side effects.
+  /// This should not be used directly.
+  StorageUniquer &getEffectUniquer();
+
   /// These APIs are tracking whether the context will be used in a
   /// multithreading environment: this has no effect other than enabling
   /// assertions on misuses of some APIs.
   void enterMultiThreadedExecution();
   void exitMultiThreadedExecution();
 
 private:
   const std::unique_ptr<MLIRContextImpl> impl;

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

 // querying the side effects of an operation.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INTERFACES_SIDEEFFECTS_H
 #define MLIR_INTERFACES_SIDEEFFECTS_H
 
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/StorageUniquerSupport.h"
 
 namespace mlir {
 namespace SideEffects {
+
+namespace detail {
+class EffectUniquer;
+} // namespace detail
+
+/// The default storage class for Effects. Only contains the type ID of the
+/// effect.
+class EffectStorage : public StorageUniquer::BaseStorage {
+  friend StorageUniquer;
+  friend MLIRContext;
+
+public:
+  /// Constructs a storage with the given type ID.
+  explicit EffectStorage(const TypeID &id) : typeID(id) {}
+
+  /// Returns type ID of the storage.
+  TypeID getTypeID() const { return typeID; }
+
+private:
+  /// Friends can construct a storage and initialize it later.
+  EffectStorage() {}
+  void initialize(const TypeID &id) { typeID = id; }
+
+  /// Type ID of the effect.
+  TypeID typeID;
+};
+
+namespace detail {
+/// A utility class to get or create unique instances of side effects within an
+/// MLIRContext.
+class EffectUniquer {
+public:
+  /// Get a uniqued non-parametric side effect.
+  // TODO: support parametric side effects
+  template <typename T>
+  static T get(MLIRContext *ctx) {
+    return ctx->getEffectUniquer().get<typename T::ImplType>(T::getTypeID());
+  }
+};
+} // namespace detail
+
 //===----------------------------------------------------------------------===//
 // Effects
 //===----------------------------------------------------------------------===//
 
 /// This class represents a base class for a specific effect type.
 class Effect {
 public:
-  /// This base class is used for derived effects that are non-parametric.
-  template <typename DerivedEffect, typename BaseEffect = Effect>
-  class Base : public BaseEffect {
-  public:
-    using BaseT = Base<DerivedEffect>;
+  template <typename DerivedEffect, typename BaseEffect = Effect,
+            typename StorageType = EffectStorage>
+  using Base =
+      mlir::detail::StorageUserBase<DerivedEffect, BaseEffect, StorageType,
+                                    detail::EffectUniquer>;
+  using ImplType = EffectStorage;
 
-    /// Return the unique identifier for the base effects class.
-    static TypeID getEffectID() { return TypeID::get<DerivedEffect>(); }
-
-    /// 'classof' used to support llvm style cast functionality.
-    static bool classof(const ::mlir::SideEffects::Effect *effect) {
-      return effect->getEffectID() == BaseT::getEffectID();
-    }
+  TypeID getTypeID() const { return impl->getTypeID(); }
 
-    /// Returns a unique instance for the derived effect class.
-    static DerivedEffect *get() {
-      return BaseEffect::template get<DerivedEffect>();
-    }
-    using BaseEffect::get;
+  /// Support for LLVM-style casting.
+  template <typename U>
+  bool isa() const;
+  template <typename First, typename Second, typename... Rest>
+  bool isa() const;
+  template <typename U>
+  U dyn_cast() const;
+  template <typename U>
+  U dyn_cast_or_null() const;
+  template <typename U>
+  U cast() const;
+
+  /// Support casting to itself.
+  static bool classof(Effect) { return true; }
+
+  Effect() : impl(nullptr) {}
+  /*implicit*/ Effect(ImplType *impl) : impl(impl) {}
+  Effect(const Effect &) = default;
 
-  protected:
-    Base() : BaseEffect(BaseT::getEffectID()) {}
-  };
+  ImplType *impl;
 
-  /// Return the unique identifier for the base effects class.
-  TypeID getEffectID() const { return id; }
-
-  /// Returns a unique instance for the given effect class.
-  template <typename DerivedEffect> static DerivedEffect *get() {
-    static_assert(std::is_base_of<Effect, DerivedEffect>::value,
-                  "expected DerivedEffect to inherit from Effect");
-
-    static DerivedEffect instance;
-    return &instance;
-  }
+template <typename U>
+bool Effect::isa() const {
+  assert(impl && "isa<> used on a null attribute.");
+  return U::classof(*this);
 
-protected:
-  Effect(TypeID id) : id(id) {}
+template <typename First, typename Second, typename... Rest>
+bool Effect::isa() const {
+  return isa<First>() || isa<Second, Rest...>();
+}
 
-private:
-  /// The id of the derived effect class.
-  TypeID id;
-};
+template <typename U>
+U Effect::dyn_cast() const {
+  return isa<U>() ? U(impl) : U(nullptr);
+}
+template <typename U>
+U Effect::dyn_cast_or_null() const {
+  return (impl && isa<U>()) ? U(impl) : U(nullptr);
+}
+template <typename U>
+U Effect::cast() const {
+  assert(isa<U>());
+  return U(impl);
+}
 
 //===----------------------------------------------------------------------===//
 // Resources
 //===----------------------------------------------------------------------===//
 
 /// This class represents a specific resource that an effect applies to. This
 /// class represents an abstract interface for a given resource.
 class Resource {

 };
 
 /// This class represents a specific instance of an effect. It contains the
 /// effect being applied, a resource that corresponds to where the effect is
 /// applied, and an optional value(either operand, result, or region entry
 /// argument) that the effect is applied to.
 template <typename EffectT> class EffectInstance {
 public:
-  EffectInstance(EffectT *effect, Resource *resource = DefaultResource::get())
+  EffectInstance(EffectT effect, Resource *resource = DefaultResource::get())
       : effect(effect), resource(resource) {}
-  EffectInstance(EffectT *effect, Value value,
+  EffectInstance(EffectT effect, Value value,
                  Resource *resource = DefaultResource::get())
       : effect(effect), resource(resource), value(value) {}
 
   /// Return the effect being applied.
-  EffectT *getEffect() const { return effect; }
+  EffectT getEffect() const { return effect; }
 
   /// Return the value the effect is applied on, or nullptr if there isn't a
   /// known value being affected.
   Value getValue() const { return value; }
 
   /// Return the resource that the effect applies to.
   Resource *getResource() const { return resource; }
 
 private:
   /// The specific effect being applied.
-  EffectT *effect;
+  EffectT effect;
 
   /// The resource that the given value resides in.
   Resource *resource;
 
   /// The value that the effect applies to. This is optionally null.
   Value value;
 };
 } // namespace SideEffects
struct Effect : public SideEffects::Effect {
 
   static bool classof(const SideEffects::Effect *effect);
 };
 using EffectInstance = SideEffects::EffectInstance<Effect>;
 
 /// The following effect indicates that the operation allocates from some
 /// resource. An 'allocate' effect implies only allocation of the resource, and
 /// not any visible mutation or dereference.
-struct Allocate : public Effect::Base<Allocate> {};
+struct Allocate : public Effect::Base<Allocate> {
+  using Base::Base;
+};
 
 /// The following effect indicates that the operation frees some resource that
 /// has been allocated. An 'allocate' effect implies only de-allocation of the
 /// resource, and not any visible allocation, mutation or dereference.
-struct Free : public Effect::Base<Free> {};
+struct Free : public Effect::Base<Free> {
+  using Base::Base;
+};
 
 /// The following effect indicates that the operation reads from some resource.
 /// A 'read' effect implies only dereferencing of the resource, and not any
 /// visible mutation.
-struct Read : public Effect::Base<Read> {};
+struct Read : public Effect::Base<Read> {
+  using Base::Base;
+};
 
 /// The following effect indicates that the operation writes to some resource. A
 /// 'write' effect implies only mutating a resource, and not any visible
 /// dereference or read.
-struct Write : public Effect::Base<Write> {};
+struct Write : public Effect::Base<Write> {
+  using Base::Base;
+};
 } // namespace MemoryEffects
 
 //===----------------------------------------------------------------------===//
 // SideEffect Utilities
 //===----------------------------------------------------------------------===//
 
 /// Return true if the given operation is unused, and has no side effects on
 /// memory that prevent erasing.

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

class EffectOpInterfaceBase<string name, string baseEffect>
   let extraClassDeclaration = [{
     /// Collect all of the effect instances that correspond to the given
     /// `Effect` and place them in 'effects'.
     template <typename Effect> void getEffects(
       SmallVectorImpl<SideEffects::EffectInstance<
                                               }] # baseEffect # [{>> &effects) {
       getEffects(effects);
       llvm::erase_if(effects, [&](auto &it) {
-        return !llvm::isa<Effect>(it.getEffect());
+        return !it.getEffect().template isa<Effect>();
       });
     }
 
     /// Returns true if this operation exhibits the given effect.
     template <typename Effect> bool hasEffect() {
       SmallVector<SideEffects::EffectInstance<}] # baseEffect # [{>, 4> effects;
       getEffects(effects);
       return llvm::any_of(effects, [](const auto &it) {
-        return llvm::isa<Effect>(it.getEffect());
+        return it.getEffect().template isa<Effect>();
       });
     }
 
     /// Returns true if this operation only has the given effect.
     template <typename Effect> bool onlyHasEffect() {
       SmallVector<SideEffects::EffectInstance<}] # baseEffect # [{>, 4> effects;
       getEffects(effects);
       return !effects.empty() && llvm::all_of(effects, [](const auto &it) {
-        return isa<Effect>(it.getEffect());
+        return it.getEffect().template isa<Effect>();
       });
     }
 
     /// Returns true if this operation has no effects.
     bool hasNoEffect() {
       SmallVector<SideEffects::EffectInstance<}] # baseEffect # [{>, 4> effects;
       getEffects(effects);
       return effects.empty();

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

if (auto memEffect = dyn_cast<MemoryEffectOpInterface>(defOp)) {
       // Collect all memory effects on `v`.
       SmallVector<MemoryEffects::EffectInstance, 1> effects;
       memEffect.getEffectsOnValue(v, effects);
 
       // If we have the 'Allocate' memory effect on `v`, then `v` should be the
       // original buffer.
       if (llvm::any_of(
               effects, [](const MemoryEffects::EffectInstance &instance) {
-                return isa<MemoryEffects::Allocate>(instance.getEffect());
+                return instance.getEffect().isa<MemoryEffects::Allocate>();
               }))
         return v;
     }
 
     if (auto viewLikeOp = dyn_cast<ViewLikeOpInterface>(defOp)) {
       auto it =
           aliases.insert(std::make_pair(v, find(viewLikeOp.getViewSource())));
       return it.first->second;

mlir/lib/IR/MLIRContext.cpp

 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Identifier.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/Types.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
#endif
   StorageUniquer attributeUniquer;
 
   /// Cached Attribute Instances.
   BoolAttr falseAttr, trueAttr;
   UnitAttr unitAttr;
   UnknownLoc unknownLocAttr;
   DictionaryAttr emptyDictionaryAttr;
 
+  //===--------------------------------------------------------------------===//
+  // Side effect uniquing
+  //===--------------------------------------------------------------------===//
+
+  /// Effect uniquing.
+  StorageUniquer effectUniquer;
+
 public:
   MLIRContextImpl() : identifiers(identifierAllocator) {}
   ~MLIRContextImpl() {
     for (auto typeMapping : registeredTypes)
       typeMapping.second->~AbstractType();
     for (auto attrMapping : registeredAttributes)
       attrMapping.second->~AbstractAttribute();
   }
impl->emptyDictionaryAttr =
       AttributeUniquer::get<DictionaryAttr>(this, ArrayRef<NamedAttribute>());
 
   // Register the affine storage objects with the uniquer.
   impl->affineUniquer
       .registerParametricStorageType<AffineBinaryOpExprStorage>();
   impl->affineUniquer
       .registerParametricStorageType<AffineConstantExprStorage>();
   impl->affineUniquer.registerParametricStorageType<AffineDimExprStorage>();
+
+  // Register the memory effects with the uniquer.
+  auto registerMemoryEffect = [this](TypeID typeID) {
+    impl->effectUniquer
+        .registerSingletonStorageType<SideEffects::EffectStorage>(
+            typeID, [typeID](SideEffects::EffectStorage *storage) {
+              storage->initialize(typeID);
+            });
+  };
+  registerMemoryEffect(MemoryEffects::Allocate::getTypeID());
+  registerMemoryEffect(MemoryEffects::Free::getTypeID());
+  registerMemoryEffect(MemoryEffects::Read::getTypeID());
+  registerMemoryEffect(MemoryEffects::Write::getTypeID());
 }
 
 MLIRContext::~MLIRContext() {}
 
 /// Copy the specified array of elements into memory managed by the provided
 /// bump pointer allocator.  This assumes the elements are all PODs.
 template <typename T>
 static ArrayRef<T> copyArrayRefInto(llvm::BumpPtrAllocator &allocator,
Location UnknownLoc::get(MLIRContext *context) {
   return context->getImpl().unknownLocAttr;
 }
 
 /// Return empty dictionary.
 DictionaryAttr DictionaryAttr::getEmpty(MLIRContext *context) {
   return context->getImpl().emptyDictionaryAttr;
 }
 
+StorageUniquer &MLIRContext::getEffectUniquer() {
+  return getImpl().effectUniquer;
+}
+
 //===----------------------------------------------------------------------===//
 // AffineMap uniquing
 //===----------------------------------------------------------------------===//
 
 StorageUniquer &MLIRContext::getAffineUniquer() {
   return getImpl().affineUniquer;
 }
 

mlir/lib/Interfaces/SideEffectInterfaces.cpp

 /// Include the definitions of the side effect interfaces.
 #include "mlir/Interfaces/SideEffectInterfaces.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // MemoryEffects
 //===----------------------------------------------------------------------===//
 
 bool MemoryEffects::Effect::classof(const SideEffects::Effect *effect) {
-  return isa<Allocate, Free, Read, Write>(effect);
+  return effect->isa<Allocate, Free, Read, Write>();
 }
 
 //===----------------------------------------------------------------------===//
 // SideEffect Utilities
 //===----------------------------------------------------------------------===//
 
 bool mlir::isOpTriviallyDead(Operation *op) {
   return op->use_empty() && wouldOpBeTriviallyDead(op);
while (!effectingOps.empty()) {
     if (auto effectInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
       // Check to see if this op either has no effects, or only allocates/reads
       // memory.
       SmallVector<MemoryEffects::EffectInstance, 1> effects;
       effectInterface.getEffects(effects);
       if (!llvm::all_of(effects, [op](const MemoryEffects::EffectInstance &it) {
             // We can drop allocations if the value is a result of the
             // operation.
-            if (isa<MemoryEffects::Allocate>(it.getEffect()))
+            if (it.getEffect().isa<MemoryEffects::Allocate>())
               return it.getValue() && it.getValue().getDefiningOp() == op;
             // Otherwise, the effect must be a read.
-            return isa<MemoryEffects::Read>(it.getEffect());
+            return it.getEffect().isa<MemoryEffects::Read>();
           })) {
         return false;
       }
       continue;
 
       // Otherwise, if the op has recursive side effects we can treat the
       // operation itself as having no effects.
     } else if (hasRecursiveEffects) {

mlir/lib/Transforms/BufferPlacement.cpp

operation->walk([&](MemoryEffectOpInterface opInterface) {
       SmallVector<MemoryEffects::EffectInstance, 2> effects;
       opInterface.getEffects(effects);
 
       SmallVector<MemoryEffects::EffectInstance, 2> allocateResultEffects;
       llvm::copy_if(
           effects, std::back_inserter(allocateResultEffects),
           [=](MemoryEffects::EffectInstance &it) {
             Value value = it.getValue();
-            return isa<MemoryEffects::Allocate>(it.getEffect()) && value &&
+            return it.getEffect().isa<MemoryEffects::Allocate>() && value &&
                    value.isa<OpResult>() &&
                    it.getResource() !=
                        SideEffects::AutomaticAllocationScopeResource::get();
           });
       // If there is one result only, we will be able to move the allocation and
       // (possibly existing) deallocation ops.
       if (allocateResultEffects.size() != 1)
         return;
for (auto &entry : allocs) {
             if (!effectInterface)
               return false;
             // Try to find a free effect that is applied to one of our values
             // that will be automatically freed by our pass.
             SmallVector<MemoryEffects::EffectInstance, 2> effects;
             effectInterface.getEffectsOnValue(entry.allocValue, effects);
             return llvm::any_of(
                 effects, [&](MemoryEffects::EffectInstance &it) {
-                  return isa<MemoryEffects::Free>(it.getEffect());
+                  return it.getEffect().isa<MemoryEffects::Free>();
                 });
           });
       // Assign the associated dealloc operation (if any).
       if (userIt != entry.allocValue.user_end())
         entry.deallocOperation = *userIt;
     }
   }
 

mlir/test/lib/Dialect/Test/TestDialect.cpp

if (!effectsAttr)
     return;
 
   // If there is one, it is an array of dictionary attributes that hold
   // information on the effects of this operation.
   for (Attribute element : effectsAttr) {
     DictionaryAttr effectElement = element.cast<DictionaryAttr>();
 
     // Get the specific memory effect.
-    MemoryEffects::Effect *effect =
-        llvm::StringSwitch<MemoryEffects::Effect *>(
+    MemoryEffects::Effect effect =
+        llvm::StringSwitch<MemoryEffects::Effect>(
             effectElement.get("effect").cast<StringAttr>().getValue())
-            .Case("allocate", MemoryEffects::Allocate::get())
-            .Case("free", MemoryEffects::Free::get())
-            .Case("read", MemoryEffects::Read::get())
-            .Case("write", MemoryEffects::Write::get());
+            .Case("allocate", MemoryEffects::Allocate::get(getContext()))
+            .Case("free", MemoryEffects::Free::get(getContext()))
+            .Case("read", MemoryEffects::Read::get(getContext()))
+            .Case("write", MemoryEffects::Write::get(getContext()));
 
     // Check for a result to affect.
     Value value;
     if (effectElement.get("on_result"))
       value = getResult();
 
     // Check for a non-default resource to use.
     SideEffects::Resource *resource = SideEffects::DefaultResource::get();

mlir/test/lib/IR/TestSideEffects.cpp

module.walk([&](MemoryEffectOpInterface op) {
       if (effects.empty()) {
         op.emitRemark() << "operation has no memory effects";
         return;
       }
 
       for (MemoryEffects::EffectInstance instance : effects) {
         auto diag = op.emitRemark() << "found an instance of ";
 
-        if (isa<MemoryEffects::Allocate>(instance.getEffect()))
+        if (instance.getEffect().isa<MemoryEffects::Allocate>())
           diag << "'allocate'";
-        else if (isa<MemoryEffects::Free>(instance.getEffect()))
+        else if (instance.getEffect().isa<MemoryEffects::Free>())
           diag << "'free'";
-        else if (isa<MemoryEffects::Read>(instance.getEffect()))
+        else if (instance.getEffect().isa<MemoryEffects::Read>())
           diag << "'read'";
-        else if (isa<MemoryEffects::Write>(instance.getEffect()))
+        else if (instance.getEffect().isa<MemoryEffects::Write>())
           diag << "'write'";
 
         if (instance.getValue())
           diag << " on a value,";
 
         diag << " on resource '" << instance.getResource()->getName() << "'";
       }
     });

mlir/test/mlir-tblgen/op-side-effects.td

def SideEffectOpA : TEST_Op<"side_effect_op_a"> {
   let results = (outs Res<AnyMemRef, "", [MemAlloc<CustomResource>]>);
 }
 
 def SideEffectOpB : TEST_Op<"side_effect_op_b",
     [MemoryEffects<[MemWrite<CustomResource>]>]>;
 
 // CHECK: void SideEffectOpA::getEffects
 // CHECK:   for (::mlir::Value value : getODSOperands(0))
-// CHECK:     effects.emplace_back(MemoryEffects::Read::get(), value, ::mlir::SideEffects::DefaultResource::get());
+// CHECK:     effects.emplace_back(MemoryEffects::Read::get(getContext()), value, ::mlir::SideEffects::DefaultResource::get());
 // CHECK:   for (::mlir::Value value : getODSResults(0))
-// CHECK:     effects.emplace_back(MemoryEffects::Allocate::get(), value, CustomResource::get());
+// CHECK:     effects.emplace_back(MemoryEffects::Allocate::get(getContext()), value, CustomResource::get());
 
 // CHECK: void SideEffectOpB::getEffects
-// CHECK:   effects.emplace_back(MemoryEffects::Write::get(), CustomResource::get());
+// CHECK:   effects.emplace_back(MemoryEffects::Write::get(getContext()), CustomResource::get());

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

for (auto &it : interfaceEffects) {
     for (auto &location : it.second) {
       if (location.kind != EffectKind::Static) {
         body << "  for (::mlir::Value value : getODS"
              << (location.kind == EffectKind::Operand ? "Operands" : "Results")
              << "(" << location.index << "))\n  ";
       }
 
       body << "  effects.emplace_back(" << location.effect.getName()
-           << "::get()";
+           << "::get(getContext())";
 
       // If the effect isn't static, it has a specific value attached to it.
       if (location.kind != EffectKind::Static)
         body << ", value";
       body << ", " << location.effect.getResource() << "::get());\n";
     }
   }
 }