[mlir] Keep sorted vector of registered operation names for efficient lookup

[ezhulenev]

Summary

I see a lot of array sorting in stack traces of our compiler, canonicalizer traverses this list every time it builds a pattern set, and it gets expensive very quickly.

[rriddle]

Do you have an idea on the cost of this vs. say computing on demand?

[ezhulenev]

In D118937#3294739, @rriddle wrote:

Do you have an idea on the cost of this vs. say computing on demand?

I was looking at the crash, and from 30 threads running MLIR compiler, ~10 were in getRegisteredOperations. For TF/TFRT compiler it takes 300us to get all operations, and it is called ~30 times for a typical module (10ms in total).

Inside google it's in top 7 function under mlir::MLIRContext

mlir::MLIRContext::getOrLoadDialect
mlir::MLIRContext::MLIRContext
mlir::MLIRContext::~MLIRContext
mlir::MLIRContextImpl::~MLIRContextImpl
mlir::MLIRContext::loadAllAvailableDialects
mlir::MLIRContext::getOrLoadDialect()::{lambda()#1}::operator()
mlir::MLIRContext::getRegisteredOperations
mlir::MLIRContext::getRegisteredOperations()::$_1::__invoke

Assuming that this will be called at least once during the MLIRContext lifetime it's cheaper to precompute.

[Hardcode84]

mlir/lib/IR/MLIRContext.cpp

189

Why SmallVector? It's unlikely for context to have that few operations but it will bloat context object size.

[ezhulenev]

mlir/lib/IR/MLIRContext.cpp

189

Just for consistency with every other container, I don't think that this really matters

[mehdi_amini]

Oh wow!
I'm really surprised we're using this API in the canonicalized and in FrozenRewritePatternSet.
This was really supposed to be more of a debugging API only.

It's a bit annoying that we're keeping both a map and a vector in the Context, but that seems justified if this API is actually useful.

Do you have an idea on the cost of this vs. say computing on demand?

Isn't "computing on demand" what to the existing API does? Do you have something else in mind?

mlir/include/mlir/IR/MLIRContext.h

179–181

mlir/lib/IR/MLIRContext.cpp

187

[mehdi_amini]

mlir/lib/IR/MLIRContext.cpp

189

You may use llvm::SmallVector<RegisteredOperationName, 0> to avoid the extra bloat (this is even smaller than std::vector if I remember correctly)

[rriddle]

In D118937#3294837, @mehdi_amini wrote:

Oh wow!
I'm really surprised we're using this API in the canonicalized and in FrozenRewritePatternSet.
This was really supposed to be more of a debugging API only.

It's a bit annoying that we're keeping both a map and a vector in the Context, but that seems justified if this API is actually useful.

Do you have an idea on the cost of this vs. say computing on demand?

Isn't "computing on demand" what to the existing API does? Do you have something else in mind?

No, by "on demand" I mean lazy sorting. Instead of doing a sorted insert for every operation, only sort if the list is requested and not already sorted.

[mehdi_amini]

In D118937#3294850, @rriddle wrote:

Isn't "computing on demand" what to the existing API does? Do you have something else in mind?

No, by "on demand" I mean lazy sorting. Instead of doing a sorted insert for every operation, only sort if the list is requested and not already sorted.

Gotcha!
So you'd keep a bool sorted = falst; next to the vector, and in the getRegisteredOperations() call you would check:

if (!sorted) {
  llvm::sort(ctxImpl.sortedRegisteredOperations),
                        [](auto &lhs, auto &rhs) {
                          return lhs.getIdentifier().compare(
                              rhs.getIdentifier());
                        });
  sorted = true;
}

Makes sense!

[rriddle]

In D118937#3294858, @mehdi_amini wrote:

In D118937#3294850, @rriddle wrote:

Isn't "computing on demand" what to the existing API does? Do you have something else in mind?

No, by "on demand" I mean lazy sorting. Instead of doing a sorted insert for every operation, only sort if the list is requested and not already sorted.

Gotcha!
So you'd keep a bool sorted = falst; next to the vector, and in the getRegisteredOperations() call you would check:

if (!sorted) {
  llvm::sort(ctxImpl.sortedRegisteredOperations),
                        [](auto &lhs, auto &rhs) {
                          return lhs.getIdentifier().compare(
                              rhs.getIdentifier());
                        });
  sorted = true;
}

Makes sense!

You'd also need to lock while actually sorting, but effectively yeah.

It will lead to slightly slower startup times, but we can always switch to deferred sorting if that is a problem. (It isn't clear how much it matters in practice anyways).

mlir/lib/IR/MLIRContext.cpp

189

[ezhulenev]

It will lead to slightly slower startup times, but we can always switch to deferred sorting if that is a problem. (It isn't clear how much it matters in practice anyways).

I think that inserting into various maps and synchronization is much more expensive, and for simplicity I'd prefer to keep all collections mutations close to each other.

[ezhulenev]

Upd

mlir/include/mlir/IR/MLIRContext.h

public:
   /// Return true if we should attach the current stacktrace to diagnostics when
   /// emitted.
   bool shouldPrintStackTraceOnDiagnostic();
 
   /// Set the flag specifying if we should attach the current stacktrace when
   /// emitting diagnostics.
   void printStackTraceOnDiagnostic(bool enable);
 
-  /// Return information about all registered operations.  This isn't very
-  /// efficient: typically you should ask the operations about their properties
-  /// directly.
-  std::vector<RegisteredOperationName> getRegisteredOperations();
+  /// Return a sorted array containing the information about all registered
+  /// operations.
+  ArrayRef<RegisteredOperationName> getRegisteredOperations();
 
   /// Return true if this operation name is registered in this context.
   bool isOperationRegistered(StringRef name);
 
   // This is effectively private given that only MLIRContext.cpp can see the
   // MLIRContextImpl type.
   MLIRContextImpl &getImpl() { return *impl; }
 

mlir/lib/IR/MLIRContext.cpp

#endif
   llvm::BumpPtrAllocator abstractDialectSymbolAllocator;
 
   /// This is a mapping from operation name to the operation info describing it.
   llvm::StringMap<OperationName::Impl> operations;
 
   /// A vector of operation info specifically for registered operations.
   llvm::StringMap<RegisteredOperationName> registeredOperations;
 
+  /// This is a sorted container of registered operations for a deterministic
+  /// and efficient `getRegisteredOperations` implementation.
+  SmallVector<RegisteredOperationName, 0> sortedRegisteredOperations;
+
   /// A mutex used when accessing operation information.
   llvm::sys::SmartRWMutex<true> operationInfoMutex;
 
   //===--------------------------------------------------------------------===//
   // Affine uniquing
   //===--------------------------------------------------------------------===//
 
   // Affine expression, map and integer set uniquing.

 }
 
 /// Set the flag specifying if we should attach the current stacktrace when
 /// emitting diagnostics.
 void MLIRContext::printStackTraceOnDiagnostic(bool enable) {
   impl->printStackTraceOnDiagnostic = enable;
 }
 
-/// Return information about all registered operations.  This isn't very
-/// efficient, typically you should ask the operations about their properties
-/// directly.
-std::vector<RegisteredOperationName> MLIRContext::getRegisteredOperations() {
-  // We just have the operations in a non-deterministic hash table order. Dump
-  // into a temporary array, then sort it by operation name to get a stable
-  // ordering.
-  auto unwrappedNames = llvm::make_second_range(impl->registeredOperations);
-  std::vector<RegisteredOperationName> result(unwrappedNames.begin(),
-                                              unwrappedNames.end());
-  llvm::array_pod_sort(result.begin(), result.end(),
-                       [](const RegisteredOperationName *lhs,
-                          const RegisteredOperationName *rhs) {
-                         return lhs->getIdentifier().compare(
-                             rhs->getIdentifier());
-                       });
-
-  return result;
+/// Return information about all registered operations.
+ArrayRef<RegisteredOperationName> MLIRContext::getRegisteredOperations() {
+  return impl->sortedRegisteredOperations;
 }
 
 bool MLIRContext::isOperationRegistered(StringRef name) {
   return RegisteredOperationName::lookup(name, this).hasValue();
 }
 
 void Dialect::addType(TypeID typeID, AbstractType &&typeInfo) {
   auto &impl = context->getImpl();
if (it.second)
     it.first->second.name = StringAttr::get(ctx, name);
   OperationName::Impl &impl = it.first->second;
 
   if (impl.isRegistered()) {
     llvm::errs() << "error: operation named '" << name
                  << "' is already registered.\n";
     abort();
   }
-  ctxImpl.registeredOperations.try_emplace(name,
-                                           RegisteredOperationName(&impl));
+  auto emplaced = ctxImpl.registeredOperations.try_emplace(
+      name, RegisteredOperationName(&impl));
+  assert(emplaced.second && "operation name registration must be successful");
+
+  // Add emplaced operation name to the sorted operations container.
+  RegisteredOperationName &value = emplaced.first->getValue();
+  ctxImpl.sortedRegisteredOperations.insert(
+      llvm::upper_bound(ctxImpl.sortedRegisteredOperations, value,
+                        [](auto &lhs, auto &rhs) {
+                          return lhs.getIdentifier().compare(
+                              rhs.getIdentifier());
+                        }),
+      value);
 
   // Update the registered info for this operation.
   impl.dialect = &dialect;
   impl.typeID = typeID;
   impl.interfaceMap = std::move(interfaceMap);
   impl.foldHookFn = std::move(foldHook);
   impl.getCanonicalizationPatternsFn = std::move(getCanonicalizationPatterns);
   impl.hasTraitFn = std::move(hasTrait);