[mlir] Support for mutable types

[ftynse]

Summary

Introduce support for mutable storage in the StorageUniquer infrastructure.
This makes MLIR have key-value storage instead of just uniqued key storage. A
storage instance now contains a unique immutable key and a mutable value, both
stored in the arena allocator that belongs to the context. This is a
preconditio for supporting recursive types that require delayed initialization,
in particular LLVM structure types. The functionality is exercised in the test
pass with trivial self-recursive type. So far, recursive types can only be
printed in parsed in a closed type system. Removing this restriction is left
for future work.

[rriddle]

Thanks Alex, I think this is an elegant solution.

Could you update https://mlir.llvm.org/docs/Tutorials/DefiningAttributesAndTypes/? It's gotten quite crusty and needs a good cleanup, so I can add docs when I clean it up if you don't want to do it now. This also might be something that we could include in the Toy dialect.

Just one conceptual question, what are the intended use cases for a failed mutation? In the current use cases we have, the type would realistically assert that the body was either not set or identical to the newly provided body.
For the cases where you envision this being used, what would a type do with the result of mutate?

Adding approval here, because the result of the above discussion doesn't change much of anything user facing.

mlir/include/mlir/IR/TypeSupport.h

136

Can you add the same to the AttributeUniquer for uniformity?

mlir/test/lib/IR/TestTypes.cpp

27

nit: Can you move the body of this function out to a new function?

[ftynse]

Address review

[ftynse]

In D84171#2165365, @rriddle wrote:

Thanks Alex, I think this is an elegant solution.

Could you update https://mlir.llvm.org/docs/Tutorials/DefiningAttributesAndTypes/? It's gotten quite crusty and needs a good cleanup, so I can add docs when I clean it up if you don't want to do it now. This also might be something that we could include in the Toy dialect.

Sure, I was waiting to make sure we agree on the approach before updating that. I added a description of types with mutable components, but haven't touched the rest. I want to make sure LLVM dialect type rework is on tracks first, and then can look into updating the rest if you don't beat me to it.

Just one conceptual question, what are the intended use cases for a failed mutation? In the current use cases we have, the type would realistically assert that the body was either not set or identical to the newly provided body.
For the cases where you envision this being used, what would a type do with the result of mutate?

One thing is to have a graceful failure mechanism similar to getChecked. A bit stretched example: imagine creating two types with the same name by parsing them from strings (we intend to expose this through C API), it's better to be able to return null to the caller instead of asserting inside the library. One could argue that it's a wrong use of the library, but it can be triggered by user input so having an option not to assert sounds better.

Another thing is partly provisioning for the type auto-renaming scheme along the lines of

Type getTypeWithRenaming(StringRef name, ...) {
  Type t;
  do {
    t = RecursiveType::get(name);
    if (succeeded(t.setBody(...)) break;
    name = nextPossibleName(name);
  } while (true);
  return t ;
}

[ftynse]

upd

[ftynse]

rebase

mlir/include/mlir/IR/StorageUniquerSupport.h

protected:
   template <typename LocationT, typename... Args>
   static ConcreteT getChecked(LocationT loc, unsigned kind, Args... args) {
     // If the construction invariants fail then we return a null attribute.
     if (failed(ConcreteT::verifyConstructionInvariants(loc, args...)))
       return ConcreteT();
     return UniquerT::template get<ConcreteT>(loc.getContext(), kind, args...);
   }
 
+  /// Mutate the current storage instance. This will not change the unique key.
+  /// The arguments are forwarded to 'ConcreteT::mutate'.
+  template <typename... Args>
+  LogicalResult mutate(Args &&... args) {
+    return UniquerT::mutate(this->getContext(), getImpl(),
+                            std::forward<Args>(args)...);
+  }
+
   /// Default implementation that just returns success.
   template <typename... Args>
   static LogicalResult verifyConstructionInvariants(Args... args) {
     return success();
   }
 
   /// Utility for easy access to the storage instance.
   ImplType *getImpl() const { return static_cast<ImplType *>(this->impl); }
 };
 } // namespace detail
 } // namespace mlir
 
 #endif

mlir/include/mlir/IR/TypeSupport.h

struct TypeUniquer {
   template <typename T, typename... Args>
   static T get(MLIRContext *ctx, unsigned kind, Args &&... args) {
     return ctx->getTypeUniquer().get<typename T::ImplType>(
         [&](TypeStorage *storage) {
           storage->initialize(AbstractType::lookup(T::getTypeID(), ctx));
         },
         kind, std::forward<Args>(args)...);
   }
+
+  /// Change the mutable component of the given type instance in the provided
+  /// context.
+  template <typename ImplType, typename... Args>
+  static LogicalResult mutate(MLIRContext *ctx, ImplType *impl,
+                              Args &&... args) {
+    assert(impl && "cannot mutate null type");
+    return ctx->getTypeUniquer().mutate(impl, std::forward<Args>(args)...);
+  }
 };
 } // namespace detail
 
 } // end namespace mlir
 
 #endif

mlir/include/mlir/IR/Types.h

 ///      * 'args' must correspond with the arguments passed into the
 ///        'TypeBase::get' call after the type kind.
 ///
 ///
 /// Type storage objects inherit from TypeStorage and contain the following:
 ///    - The type kind (for LLVM-style RTTI).
 ///    - The dialect that defined the type.
 ///    - Any parameters of the type.
+///    - An optional mutable component.
 /// For non-parametric types, a convenience DefaultTypeStorage is provided.
 /// Parametric storage types must derive TypeStorage and respect the following:
 ///    - Define a type alias, KeyTy, to a type that uniquely identifies the
 ///      instance of the type within its kind.
 ///      * The key type must be constructible from the values passed into the
 ///        detail::TypeUniquer::get call after the type kind.
 ///      * If the KeyTy does not have an llvm::DenseMapInfo specialization, the
 ///        storage class must define a hashing method:
 ///         'static unsigned hashKey(const KeyTy &)'
 ///
 ///    - Provide a method, 'bool operator==(const KeyTy &) const', to
 ///      compare the storage instance against an instance of the key type.
 ///
-///    - Provide a construction method:
+///    - Provide a static construction method:
 ///        'DerivedStorage *construct(TypeStorageAllocator &, const KeyTy &key)'
 ///      that builds a unique instance of the derived storage. The arguments to
 ///      this function are an allocator to store any uniqued data within the
 ///      context and the key type for this storage.
+///
+///    - If they have a mutable component, this component must not be a part of
+//       the key.
 class Type {
 public:
   /// Integer identifier for all the concrete type kinds.
   /// Note: This is not an enum class as each dialect will likely define a
   /// separate enumeration for the specific types that they define. Not being an
   /// enum class also simplifies the handling of type kinds by not requiring
   /// casts for each use.
   enum Kind {

mlir/include/mlir/Support/StorageUniquer.h

 //===- StorageUniquer.h - Common Storage Class Uniquer ----------*- 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_SUPPORT_STORAGEUNIQUER_H
 #define MLIR_SUPPORT_STORAGEUNIQUER_H
 
 #include "mlir/Support/LLVM.h"
+#include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/Allocator.h"
 
 namespace mlir {
 namespace detail {
 struct StorageUniquerImpl;
 
 /// Trait to check if ImplTy provides a 'getKey' method with types 'Args'.

 ///      this function are an allocator to store any uniqued data and the key
 ///      type for this storage.
 ///
 ///    - Provide a cleanup method:
 ///        'void cleanup()'
 ///      that is called when erasing a storage instance. This should cleanup any
 ///      fields of the storage as necessary and not attempt to free the memory
 ///      of the storage itself.
+///
+/// Storage classes may have an optional mutable component, which must not take
+/// part in the unique immutable key. In this case, storage classes may be
+/// mutated with `mutate` and must additionally respect the following:
+///    - Provide a mutation method:
+///        'LogicalResult mutate(StorageAllocator &, <...>)'
+///      that is called when mutating a storage instance. The first argument is
+///      an allocator to store any mutable data, and the remaining arguments are
+///      forwarded from the call site. The storage can be mutated at any time
+///      after creation. Care must be taken to avoid excessive mutation since
+///      the allocated storage can keep containing previous states. The return
+///      value of the function is used to indicate whether the mutation was
+///      successful, e.g., to limit the number of mutations or enable deferred
+///      one-time assignment of the mutable component.
 class StorageUniquer {
 public:
   StorageUniquer();
   ~StorageUniquer();
 
   /// Set the flag specifying if multi-threading is disabled within the uniquer.
   void disableMultithreading(bool disable = true);
 
auto ctorFn = [&](StorageAllocator &allocator) {
       auto *storage = new (allocator.allocate<Storage>()) Storage();
       if (initFn)
         initFn(storage);
       return storage;
     };
     return static_cast<Storage *>(getImpl(kind, ctorFn));
   }
 
+  /// Changes the mutable component of 'storage' by forwarding the trailing
+  /// arguments to the 'mutate' function of the derived class.
+  template <typename Storage, typename... Args>
+  LogicalResult mutate(Storage *storage, Args &&... args) {
+    auto mutationFn = [&](StorageAllocator &allocator) -> LogicalResult {
+      return static_cast<Storage &>(*storage).mutate(
+          allocator, std::forward<Args>(args)...);
+    };
+    return mutateImpl(mutationFn);
+  }
+
   /// Erases a uniqued instance of 'Storage'. This function is used for derived
   /// types that have complex storage or uniquing constraints.
   template <typename Storage, typename Arg, typename... Args>
   void erase(unsigned kind, Arg &&arg, Args &&... args) {
     // Construct a value of the derived key type.
     auto derivedKey =
         getKey<Storage>(std::forward<Arg>(arg), std::forward<Args>(args)...);
 
BaseStorage *getImpl(unsigned kind,
                        function_ref<BaseStorage *(StorageAllocator &)> ctorFn);
 
   /// Implementation for erasing an instance of a derived type with complex
   /// storage.
   void eraseImpl(unsigned kind, unsigned hashValue,
                  function_ref<bool(const BaseStorage *)> isEqual,
                  function_ref<void(BaseStorage *)> cleanupFn);
 
+  /// Implementation for mutating an instance of a derived storage.
+  LogicalResult
+  mutateImpl(function_ref<LogicalResult(StorageAllocator &)> mutationFn);
+
   /// The internal implementation class.
   std::unique_ptr<detail::StorageUniquerImpl> impl;
 
   //===--------------------------------------------------------------------===//
   // Key Construction
   //===--------------------------------------------------------------------===//
 
   /// Used to construct an instance of 'ImplTy::KeyTy' if there is an

mlir/lib/Support/StorageUniquer.cpp

void erase(unsigned kind, unsigned hashValue,
     if (existing == storageTypes.end())
       return;
 
     // Cleanup the storage and remove it from the map.
     cleanupFn(existing->storage);
     storageTypes.erase(existing);
   }
 
+  /// Mutates an instance of a derived storage in a thread-safe way.
+  LogicalResult
+  mutate(function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
+    if (!threadingIsEnabled)
+      return mutationFn(allocator);
+
+    llvm::sys::SmartScopedWriter<true> lock(mutex);
+    return mutationFn(allocator);
+  }
+
   //===--------------------------------------------------------------------===//
   // Instance Storage
   //===--------------------------------------------------------------------===//
 
   /// Utility to create and initialize a storage instance.
   BaseStorage *
   initializeStorage(unsigned kind,
                     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {

 
 /// Implementation for erasing an instance of a derived type with complex
 /// storage.
 void StorageUniquer::eraseImpl(unsigned kind, unsigned hashValue,
                                function_ref<bool(const BaseStorage *)> isEqual,
                                function_ref<void(BaseStorage *)> cleanupFn) {
   impl->erase(kind, hashValue, isEqual, cleanupFn);
 }
+
+/// Implementation for mutating an instance of a derived storage.
+LogicalResult StorageUniquer::mutateImpl(
+    function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
+  return impl->mutate(mutationFn);
+}

mlir/test/IR/recursive-type.mlir

+// RUN: mlir-opt %s -test-recursive-types | FileCheck %s
+
+// CHECK-LABEL: @roundtrip
+func @roundtrip() {
+  // CHECK: !test.test_rec<a, test_rec<b, test_type>>
+  "test.dummy_op_for_roundtrip"() : () -> !test.test_rec<a, test_rec<b, test_type>>
+  // CHECK: !test.test_rec<c, test_rec<c>>
+  "test.dummy_op_for_roundtrip"() : () -> !test.test_rec<c, test_rec<c>>
+  return
+}
+
+// CHECK-LABEL: @create
+func @create() {
+  // CHECK: !test.test_rec<some_long_and_unique_name, test_rec<some_long_and_unique_name>>
+  return
+}

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

 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringSwitch.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // TestDialect Interfaces
 //===----------------------------------------------------------------------===//
 

 TestDialect::TestDialect(MLIRContext *context)
     : Dialect(getDialectNamespace(), context) {
   addOperations<
 #define GET_OP_LIST
 #include "TestOps.cpp.inc"
       >();
   addInterfaces<TestOpAsmInterface, TestOpFolderDialectInterface,
                 TestInlinerInterface>();
-  addTypes<TestType>();
+  addTypes<TestType, TestRecursiveType>();
   allowUnknownOperations();
 }
 
-Type TestDialect::parseType(DialectAsmParser &parser) const {
-  if (failed(parser.parseKeyword("test_type")))
+static Type parseTestType(DialectAsmParser &parser,
+                          llvm::SetVector<Type> &stack) {
+  StringRef typeTag;
+  if (failed(parser.parseKeyword(&typeTag)))
+    return Type();
+
+  if (typeTag == "test_type")
+    return TestType::get(parser.getBuilder().getContext());
+
+  if (typeTag != "test_rec")
+    return Type();
+
+  StringRef name;
+  if (parser.parseLess() || parser.parseKeyword(&name))
     return Type();
-  return TestType::get(getContext());
+  auto rec = TestRecursiveType::create(parser.getBuilder().getContext(), name);
+
+  // If this type already has been parsed above in the stack, expect just the
+  // name.
+  if (stack.contains(rec)) {
+    if (failed(parser.parseGreater()))
+      return Type();
+    return rec;
   }
 
-void TestDialect::printType(Type type, DialectAsmPrinter &printer) const {
-  assert(type.isa<TestType>() && "unexpected type");
+  // Otherwise, parse the body and update the type.
+  if (failed(parser.parseComma()))
+    return Type();
+  stack.insert(rec);
+  Type subtype = parseTestType(parser, stack);
+  stack.pop_back();
+  if (!subtype || failed(parser.parseGreater()) || failed(rec.setBody(subtype)))
+    return Type();
+
+  return rec;
+}
+
+Type TestDialect::parseType(DialectAsmParser &parser) const {
+  llvm::SetVector<Type> stack;
+  return parseTestType(parser, stack);
+}
+
+static void printTestType(Type type, DialectAsmPrinter &printer,
+                          llvm::SetVector<Type> &stack) {
+  if (type.isa<TestType>()) {
     printer << "test_type";
+    return;
+  }
+
+  auto rec = type.cast<TestRecursiveType>();
+  printer << "test_rec<" << rec.getName();
+  if (!stack.contains(rec)) {
+    printer << ", ";
+    stack.insert(rec);
+    printTestType(rec.getBody(), printer, stack);
+    stack.pop_back();
+  }
+  printer << ">";
+}
+
+void TestDialect::printType(Type type, DialectAsmPrinter &printer) const {
+  llvm::SetVector<Type> stack;
+  printTestType(type, printer, stack);
 }
 
 LogicalResult TestDialect::verifyOperationAttribute(Operation *op,
                                                     NamedAttribute namedAttr) {
   if (namedAttr.first == "test.invalid_attr")
     return op->emitError() << "invalid to use 'test.invalid_attr'";
   return success();
 }

mlir/test/lib/Dialect/Test/TestTypes.h

static TestType get(MLIRContext *context) {
     return Base::get(context, Type::Kind::FIRST_PRIVATE_EXPERIMENTAL_9_TYPE);
   }
 
   /// Provide a definition for the necessary interface methods.
   void printTypeC(Location loc) const {
     emitRemark(loc) << *this << " - TestC";
   }
 };
+
+/// Storage for simple named recursive types, where the type is identified by
+/// its name and can "contain" another type, including itself.
+struct TestRecursiveTypeStorage : public TypeStorage {
+  using KeyTy = StringRef;
+
+  explicit TestRecursiveTypeStorage(StringRef key) : name(key), body(Type()) {}
+
+  bool operator==(const KeyTy &other) const { return name == other; }
+
+  static TestRecursiveTypeStorage *construct(TypeStorageAllocator &allocator,
+                                             const KeyTy &key) {
+    return new (allocator.allocate<TestRecursiveTypeStorage>())
+        TestRecursiveTypeStorage(allocator.copyInto(key));
+  }
+
+  LogicalResult mutate(TypeStorageAllocator &allocator, Type newBody) {
+    // Cannot set a different body than before.
+    if (body && body != newBody)
+      return failure();
+
+    body = newBody;
+    return success();
+  }
+
+  StringRef name;
+  Type body;
+};
+
+/// Simple recursive type identified by its name and pointing to another named
+/// type, potentially itself. This requires the body to be mutated separately
+/// from type creation.
+class TestRecursiveType
+    : public Type::TypeBase<TestRecursiveType, Type, TestRecursiveTypeStorage> {
+public:
+  using Base::Base;
+
+  static bool kindof(unsigned kind) {
+    return kind == Type::Kind::FIRST_PRIVATE_EXPERIMENTAL_9_TYPE + 1;
+  }
+
+  static TestRecursiveType create(MLIRContext *ctx, StringRef name) {
+    return Base::get(ctx, Type::Kind::FIRST_PRIVATE_EXPERIMENTAL_9_TYPE + 1,
+                     name);
+  }
+
+  /// Body getter and setter.
+  LogicalResult setBody(Type body) { return Base::mutate(body); }
+  Type getBody() { return getImpl()->body; }
+
+  /// Name/key getter.
+  StringRef getName() { return getImpl()->name; }
+};
+
 } // end namespace mlir
 
 #endif // MLIR_TESTTYPES_H

mlir/test/lib/IR/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestIR
   TestFunc.cpp
   TestInterfaces.cpp
   TestMatchers.cpp
   TestSideEffects.cpp
   TestSymbolUses.cpp
+  TestTypes.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRPass
   MLIRTestDialect
   )
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../Dialect/Test)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/../Dialect/Test)

mlir/test/lib/IR/TestTypes.cpp

+//===- TestTypes.cpp - Test passes for MLIR types -------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestTypes.h"
+#include "TestDialect.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+struct TestRecursiveTypesPass
+    : public PassWrapper<TestRecursiveTypesPass, FunctionPass> {
+  void runOnFunction() override {
+    FuncOp func = getFunction();
+
+    // Just make sure recurisve types are printed and parsed.
+    if (func.getName() == "roundtrip")
+      return;
+
+    // Create a recursive type and print it as a part of a dummy op.
+    if (func.getName() == "create") {
+      auto type =
+          TestRecursiveType::create(&getContext(), "some_long_and_unique_name");
+      if (failed(type.setBody(type))) {
+        func.emitError("expected to be able to set the type body");
+        return signalPassFailure();
+      }
+
+      // Setting the same body is fine.
+      if (failed(type.setBody(type))) {
+        func.emitError(
+            "expected to be able to set the type body to the same value");
+        return signalPassFailure();
+      }
+
+      // Setting a different body is not.
+      if (succeeded(type.setBody(IndexType::get(&getContext())))) {
+        func.emitError(
+            "not expected to be able to change function body more than once");
+        return signalPassFailure();
+      }
+
+      // Expecting to get the same type for the same name.
+      auto other =
+          TestRecursiveType::create(&getContext(), "some_long_and_unique_name");
+      if (type != other) {
+        func.emitError("expected type name to be the uniquing key");
+        return signalPassFailure();
+      }
+
+      // Create the op to check how the type is printed.
+      OperationState state(func.getLoc(), "test.dummy_type_test_op");
+      state.addTypes(type);
+      func.getBody().front().push_front(Operation::create(state));
+
+      return;
+    }
+
+    // Unknown key.
+    func.emitOpError() << "unexpected function name";
+    signalPassFailure();
+  }
+};
+} // end namespace
+
+namespace mlir {
+
+void registerTestRecursiveTypesPass() {
+  PassRegistration<TestRecursiveTypesPass> reg(
+      "test-recursive-types", "Test support for recursive types");
+}
+
+} // end namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestLoopFusion();
 void registerTestLoopMappingPass();
 void registerTestLoopUnrollingPass();
 void registerTestMatchers();
 void registerTestMemRefDependenceCheck();
 void registerTestMemRefStrideCalculation();
 void registerTestOpaqueLoc();
 void registerTestPreparationPassWithAllowedMemrefResults();
+void registerTestRecursiveTypesPass();
 void registerTestReducer();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestSCFUtilsPass();
 void registerTestVectorConversions();
 void registerVectorizerTestPass();
 } // namespace mlir
 
 static cl::opt<std::string>
#endif
   registerTestLoopFusion();
   registerTestLoopMappingPass();
   registerTestLoopUnrollingPass();
   registerTestMatchers();
   registerTestMemRefDependenceCheck();
   registerTestMemRefStrideCalculation();
   registerTestOpaqueLoc();
   registerTestPreparationPassWithAllowedMemrefResults();
+  registerTestRecursiveTypesPass();
   registerTestReducer();
   registerTestGpuParallelLoopMappingPass();
   registerTestSCFUtilsPass();
   registerTestVectorConversions();
   registerVectorizerTestPass();
 }
 #endif