diff --git a/mlir/include/mlir/IR/StorageUniquerSupport.h b/mlir/include/mlir/IR/StorageUniquerSupport.h --- a/mlir/include/mlir/IR/StorageUniquerSupport.h +++ b/mlir/include/mlir/IR/StorageUniquerSupport.h @@ -105,6 +105,14 @@ return UniquerT::template get(loc.getContext(), kind, args...); } + /// Mutate the current storage instance. This will not change the unique key. + /// The arguments are forwarded to 'ConcreteT::mutate'. + template + LogicalResult mutate(Args &&... args) { + return UniquerT::mutate(this->getContext(), getImpl(), + std::forward(args)...); + } + /// Default implementation that just returns success. template static LogicalResult verifyConstructionInvariants(Args... args) { diff --git a/mlir/include/mlir/IR/TypeSupport.h b/mlir/include/mlir/IR/TypeSupport.h --- a/mlir/include/mlir/IR/TypeSupport.h +++ b/mlir/include/mlir/IR/TypeSupport.h @@ -132,6 +132,15 @@ }, kind, std::forward(args)...); } + + /// Change the mutable component of the given type instance in the provided + /// context. + template + static LogicalResult mutate(MLIRContext *ctx, ImplType *impl, + Args &&... args) { + assert(impl && "cannot mutate null type"); + return ctx->getTypeUniquer().mutate(impl, std::forward(args)...); + } }; } // namespace detail diff --git a/mlir/include/mlir/IR/Types.h b/mlir/include/mlir/IR/Types.h --- a/mlir/include/mlir/IR/Types.h +++ b/mlir/include/mlir/IR/Types.h @@ -62,6 +62,7 @@ /// - 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 @@ -75,11 +76,14 @@ /// - 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. diff --git a/mlir/include/mlir/Support/StorageUniquer.h b/mlir/include/mlir/Support/StorageUniquer.h --- a/mlir/include/mlir/Support/StorageUniquer.h +++ b/mlir/include/mlir/Support/StorageUniquer.h @@ -10,6 +10,7 @@ #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" @@ -60,6 +61,20 @@ /// 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(); @@ -166,6 +181,17 @@ return static_cast(getImpl(kind, ctorFn)); } + /// Changes the mutable component of 'storage' by forwarding the trailing + /// arguments to the 'mutate' function of the derived class. + template + LogicalResult mutate(Storage *storage, Args &&... args) { + auto mutationFn = [&](StorageAllocator &allocator) -> LogicalResult { + return static_cast(*storage).mutate( + allocator, std::forward(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 @@ -206,6 +232,10 @@ function_ref isEqual, function_ref cleanupFn); + /// Implementation for mutating an instance of a derived storage. + LogicalResult + mutateImpl(function_ref mutationFn); + /// The internal implementation class. std::unique_ptr impl; diff --git a/mlir/lib/Support/StorageUniquer.cpp b/mlir/lib/Support/StorageUniquer.cpp --- a/mlir/lib/Support/StorageUniquer.cpp +++ b/mlir/lib/Support/StorageUniquer.cpp @@ -124,6 +124,16 @@ storageTypes.erase(existing); } + /// Mutates an instance of a derived storage in a thread-safe way. + LogicalResult + mutate(function_ref mutationFn) { + if (!threadingIsEnabled) + return mutationFn(allocator); + + llvm::sys::SmartScopedWriter lock(mutex); + return mutationFn(allocator); + } + //===--------------------------------------------------------------------===// // Instance Storage //===--------------------------------------------------------------------===// @@ -214,3 +224,9 @@ function_ref cleanupFn) { impl->erase(kind, hashValue, isEqual, cleanupFn); } + +/// Implementation for mutating an instance of a derived storage. +LogicalResult StorageUniquer::mutateImpl( + function_ref mutationFn) { + return impl->mutate(mutationFn); +} diff --git a/mlir/test/IR/recursive-type.mlir b/mlir/test/IR/recursive-type.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/IR/recursive-type.mlir @@ -0,0 +1,16 @@ +// RUN: mlir-opt %s -test-recursive-types | FileCheck %s + +// CHECK-LABEL: @roundtrip +func @roundtrip() { + // CHECK: !test.test_rec> + "test.dummy_op_for_roundtrip"() : () -> !test.test_rec> + // CHECK: !test.test_rec> + "test.dummy_op_for_roundtrip"() : () -> !test.test_rec> + return +} + +// CHECK-LABEL: @create +func @create() { + // CHECK: !test.test_rec> + return +} diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp --- a/mlir/test/lib/Dialect/Test/TestDialect.cpp +++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp @@ -16,6 +16,7 @@ #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; @@ -137,19 +138,73 @@ >(); addInterfaces(); - addTypes(); + addTypes(); allowUnknownOperations(); } -Type TestDialect::parseType(DialectAsmParser &parser) const { - if (failed(parser.parseKeyword("test_type"))) +static Type parseTestType(DialectAsmParser &parser, + llvm::SetVector &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(); + 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; + } + + // 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 TestType::get(getContext()); + + return rec; +} + +Type TestDialect::parseType(DialectAsmParser &parser) const { + llvm::SetVector stack; + return parseTestType(parser, stack); +} + +static void printTestType(Type type, DialectAsmPrinter &printer, + llvm::SetVector &stack) { + if (type.isa()) { + printer << "test_type"; + return; + } + + auto rec = type.cast(); + 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 { - assert(type.isa() && "unexpected type"); - printer << "test_type"; + llvm::SetVector stack; + printTestType(type, printer, stack); } LogicalResult TestDialect::verifyOperationAttribute(Operation *op, diff --git a/mlir/test/lib/Dialect/Test/TestTypes.h b/mlir/test/lib/Dialect/Test/TestTypes.h --- a/mlir/test/lib/Dialect/Test/TestTypes.h +++ b/mlir/test/lib/Dialect/Test/TestTypes.h @@ -39,6 +39,60 @@ 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(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 { +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 diff --git a/mlir/test/lib/IR/CMakeLists.txt b/mlir/test/lib/IR/CMakeLists.txt --- a/mlir/test/lib/IR/CMakeLists.txt +++ b/mlir/test/lib/IR/CMakeLists.txt @@ -5,6 +5,7 @@ TestMatchers.cpp TestSideEffects.cpp TestSymbolUses.cpp + TestTypes.cpp EXCLUDE_FROM_LIBMLIR diff --git a/mlir/test/lib/IR/TestTypes.cpp b/mlir/test/lib/IR/TestTypes.cpp new file mode 100644 --- /dev/null +++ b/mlir/test/lib/IR/TestTypes.cpp @@ -0,0 +1,78 @@ +//===- 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 { + 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 reg( + "test-recursive-types", "Test support for recursive types"); +} + +} // end namespace mlir diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp --- a/mlir/tools/mlir-opt/mlir-opt.cpp +++ b/mlir/tools/mlir-opt/mlir-opt.cpp @@ -63,6 +63,7 @@ void registerTestMemRefStrideCalculation(); void registerTestOpaqueLoc(); void registerTestPreparationPassWithAllowedMemrefResults(); +void registerTestRecursiveTypesPass(); void registerTestReducer(); void registerTestGpuParallelLoopMappingPass(); void registerTestSCFUtilsPass(); @@ -138,6 +139,7 @@ registerTestMemRefStrideCalculation(); registerTestOpaqueLoc(); registerTestPreparationPassWithAllowedMemrefResults(); + registerTestRecursiveTypesPass(); registerTestReducer(); registerTestGpuParallelLoopMappingPass(); registerTestSCFUtilsPass();