diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp index 9deb4fa78761..5395179d7e95 100644 --- a/mlir/lib/IR/MLIRContext.cpp +++ b/mlir/lib/IR/MLIRContext.cpp @@ -494,6 +494,14 @@ public: /// These are identifiers uniqued into this MLIRContext. llvm::StringMap identifiers; + //===--------------------------------------------------------------------===// + // Affine uniquing + //===--------------------------------------------------------------------===// + + // Affine allocator and mutex for thread safety. + llvm::BumpPtrAllocator affineAllocator; + llvm::sys::SmartRWMutex affineMutex; + // Affine map uniquing. using AffineMapSet = DenseSet; AffineMapSet affineMaps; @@ -514,10 +522,14 @@ public: // Uniqui'ing of AffineConstantExprStorage using constant value as key. DenseMap constExprs; - /// Type uniquing. + //===--------------------------------------------------------------------===// + // Type uniquing + //===--------------------------------------------------------------------===// TypeUniquerImpl typeUniquer; - // Attribute uniquing. + //===--------------------------------------------------------------------===// + // Attribute uniquing + //===--------------------------------------------------------------------===// // Attribute allocator and mutex for thread safety. llvm::BumpPtrAllocator attributeAllocator; @@ -1394,28 +1406,21 @@ AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount, assert(rangeSizes.empty() || results.size() == rangeSizes.size()); auto &impl = results[0].getContext()->getImpl(); - - // Check if we already have this affine map. auto key = std::make_tuple(dimCount, symbolCount, results, rangeSizes); - auto existing = impl.affineMaps.insert_as(AffineMap(), key); - // If we already have it, return that value. - if (!existing.second) - return *existing.first; + // Safely get or create an AffineMap instance. + return safeGetOrCreate(impl.affineMaps, key, impl.affineMutex, [&] { + auto *res = impl.affineAllocator.Allocate(); - // On the first use, we allocate them into the bump pointer. - auto *res = impl.allocator.Allocate(); + // Copy the results and range sizes into the bump pointer. + results = copyArrayRefInto(impl.affineAllocator, results); + rangeSizes = copyArrayRefInto(impl.affineAllocator, rangeSizes); - // Copy the results and range sizes into the bump pointer. - results = copyArrayRefInto(impl.allocator, results); - rangeSizes = copyArrayRefInto(impl.allocator, rangeSizes); - - // Initialize the memory using placement new. - new (res) - detail::AffineMapStorage{dimCount, symbolCount, results, rangeSizes}; - - // Cache and return it. - return *existing.first = AffineMap(res); + // Initialize the memory using placement new. + new (res) + detail::AffineMapStorage{dimCount, symbolCount, results, rangeSizes}; + return AffineMap(res); + }); } /// Simplify add expression. Return nullptr if it can't be simplified. @@ -1621,9 +1626,13 @@ AffineExpr AffineBinaryOpExprStorage::get(AffineExprKind kind, AffineExpr lhs, // Check if we already have this affine expression, and return it if we do. auto keyValue = std::make_tuple((unsigned)kind, lhs, rhs); - auto cached = impl.affineExprs.find(keyValue); - if (cached != impl.affineExprs.end()) - return cached->second; + + { // Check for an existing instance in read-only mode. + llvm::sys::SmartScopedReader affineLock(impl.affineMutex); + auto cached = impl.affineExprs.find(keyValue); + if (cached != impl.affineExprs.end()) + return cached->second; + } // Simplify the expression if possible. AffineExpr simplified; @@ -1651,14 +1660,18 @@ AffineExpr AffineBinaryOpExprStorage::get(AffineExprKind kind, AffineExpr lhs, if (simplified) return simplified; - // An expression with these operands will already be in the - // simplified/canonical form. Create and store it. - auto *result = impl.allocator.Allocate(); - // Initialize the memory using placement new. - new (result) AffineBinaryOpExprStorage{{kind, lhs.getContext()}, lhs, rhs}; - bool inserted = impl.affineExprs.insert({keyValue, result}).second; - assert(inserted && "the expression shouldn't already exist in the map"); - (void)inserted; + // Aquire a writer-lock so that we can safely create the new instance. + llvm::sys::SmartScopedWriter affineLock(impl.affineMutex); + + // Check for an existing instance again here, because another writer thread + // may have already created one. + auto &result = impl.affineExprs.insert({keyValue, nullptr}).first->second; + if (!result) { + // An expression with these operands will already be in the + // simplified/canonical form. Create and store it. + result = new (impl.affineAllocator.Allocate()) + AffineBinaryOpExprStorage{{kind, lhs.getContext()}, lhs, rhs}; + } return result; } @@ -1670,15 +1683,26 @@ AffineExpr mlir::getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs, AffineExpr mlir::getAffineDimExpr(unsigned position, MLIRContext *context) { auto &impl = context->getImpl(); + { // Check for an existing instance in read-only mode. + llvm::sys::SmartScopedReader affineLock(impl.affineMutex); + if (impl.dimExprs.size() > position && impl.dimExprs[position]) + return impl.dimExprs[position]; + } + + // Aquire a writer-lock so that we can safely create the new instance. + llvm::sys::SmartScopedWriter affineLock(impl.affineMutex); + // Check if we need to resize. if (position >= impl.dimExprs.size()) impl.dimExprs.resize(position + 1, nullptr); + // Check for an existing instance again here, because another writer thread + // may have already created one. auto *&result = impl.dimExprs[position]; if (result) return result; - result = impl.allocator.Allocate(); + result = impl.affineAllocator.Allocate(); // Initialize the memory using placement new. new (result) AffineDimExprStorage{{AffineExprKind::DimId, context}, position}; return result; @@ -1687,15 +1711,26 @@ AffineExpr mlir::getAffineDimExpr(unsigned position, MLIRContext *context) { AffineExpr mlir::getAffineSymbolExpr(unsigned position, MLIRContext *context) { auto &impl = context->getImpl(); + { // Check for an existing instance in read-only mode. + llvm::sys::SmartScopedReader affineLock(impl.affineMutex); + if (impl.symbolExprs.size() > position && impl.symbolExprs[position]) + return impl.symbolExprs[position]; + } + + // Aquire a writer-lock so that we can safely create the new instance. + llvm::sys::SmartScopedWriter affineLock(impl.affineMutex); + // Check if we need to resize. if (position >= impl.symbolExprs.size()) impl.symbolExprs.resize(position + 1, nullptr); + // Check for an existing instance again here, because another writer thread + // may have already created one. auto *&result = impl.symbolExprs[position]; if (result) return result; - result = impl.allocator.Allocate(); + result = impl.affineAllocator.Allocate(); // Initialize the memory using placement new. new (result) AffineSymbolExprStorage{{AffineExprKind::SymbolId, context}, position}; @@ -1704,16 +1739,13 @@ AffineExpr mlir::getAffineSymbolExpr(unsigned position, MLIRContext *context) { AffineExpr mlir::getAffineConstantExpr(int64_t constant, MLIRContext *context) { auto &impl = context->getImpl(); - auto *&result = impl.constExprs[constant]; - if (result) - return result; - - result = impl.allocator.Allocate(); - // Initialize the memory using placement new. - new (result) - AffineConstantExprStorage{{AffineExprKind::Constant, context}, constant}; - return result; + // Safely get or create an AffineConstantExpr instance. + return safeGetOrCreate(impl.constExprs, constant, impl.affineMutex, [&] { + auto *result = impl.affineAllocator.Allocate(); + return new (result) AffineConstantExprStorage{ + {AffineExprKind::Constant, context}, constant}; + }); } //===----------------------------------------------------------------------===// @@ -1728,35 +1760,32 @@ IntegerSet IntegerSet::get(unsigned dimCount, unsigned symbolCount, assert(!constraints.empty()); assert(constraints.size() == eqFlags.size()); - bool unique = constraints.size() < IntegerSet::kUniquingThreshold; - auto &impl = constraints[0].getContext()->getImpl(); - std::pair::Iterator, bool> existing; - if (unique) { - // Check if we already have this integer set. - auto key = std::make_tuple(dimCount, symbolCount, constraints, eqFlags); - existing = impl.integerSets.insert_as(IntegerSet(nullptr), key); + // A utility function to construct a new IntegerSetStorage instance. + auto constructorFn = [&] { + auto *res = impl.affineAllocator.Allocate(); - // If we already have it, return that value. - if (!existing.second) - return *existing.first; + // Copy the results and equality flags into the bump pointer. + constraints = copyArrayRefInto(impl.affineAllocator, constraints); + eqFlags = copyArrayRefInto(impl.affineAllocator, eqFlags); + + // Initialize the memory using placement new. + new (res) + detail::IntegerSetStorage{dimCount, symbolCount, constraints, eqFlags}; + return IntegerSet(res); + }; + + // If this instance is uniqued, then we handle it separately so that multiple + // threads may simulatenously access existing instances. + if (constraints.size() < IntegerSet::kUniquingThreshold) { + auto key = std::make_tuple(dimCount, symbolCount, constraints, eqFlags); + return safeGetOrCreate(impl.integerSets, key, impl.affineMutex, + constructorFn); } - // On the first use, we allocate them into the bump pointer. - auto *res = impl.allocator.Allocate(); - - // Copy the results and equality flags into the bump pointer. - constraints = copyArrayRefInto(impl.allocator, constraints); - eqFlags = copyArrayRefInto(impl.allocator, eqFlags); - - // Initialize the memory using placement new. - new (res) - detail::IntegerSetStorage{dimCount, symbolCount, constraints, eqFlags}; - - if (unique) - // Cache and return it. - return *existing.first = IntegerSet(res); - - return IntegerSet(res); + // Otherwise, aquire a writer-lock so that we can safely create the new + // instance. + llvm::sys::SmartScopedWriter affineLock(impl.affineMutex); + return constructorFn(); }