[mlir] Optimize OperationName construction and usage

[rriddle]

Summary

When constructing an OperationName, the overwhelming majority of
cases are from registered operations. This revision adds a non-locked
lookup into the currently registered operations, which prevents locking
in the common case. This revision also optimizes several uses of
RegisteredOperationName that expect the operation to be registered,
e.g. such as in OpBuilder.

These changes provides a reasonable speedup (5-10%) in some
compilations, especially on platforms where locking is expensive.

[lattner]

This is really great, nice approach to use the registered op information, I didn't realize we had that!

mlir/include/mlir/IR/Builders.h

public:
   /// Insert the given operation at the current insertion point and return it.
   Operation *insert(Operation *op);
 
   /// Creates an operation given the fields represented as an OperationState.
   Operation *createOperation(const OperationState &state);
 
 private:
   /// Helper for sanity checking preconditions for create* methods below.
-  void checkHasRegisteredInfo(const OperationName &name) {
-    if (LLVM_UNLIKELY(!name.isRegistered()))
+  template <typename OpT>
+  RegisteredOperationName getCheckRegisteredInfo(MLIRContext *ctx) {
+    Optional<RegisteredOperationName> opName =
+        RegisteredOperationName::lookup(OpT::getOperationName(), ctx);
+    if (LLVM_UNLIKELY(!opName)) {
       llvm::report_fatal_error(
-          "Building op `" + name.getStringRef() +
+          "Building op `" + OpT::getOperationName() +
           "` but it isn't registered in this MLIRContext: the dialect may not "
           "be loaded or this operation isn't registered by the dialect. See "
           "also https://mlir.llvm.org/getting_started/Faq/"
           "#registered-loaded-dependent-whats-up-with-dialects-management");
     }
+    return *opName;
+  }
 
 public:
   /// Create an operation of specific op type at the current insertion point.
   template <typename OpTy, typename... Args>
   OpTy create(Location location, Args &&...args) {
-    OperationState state(location, OpTy::getOperationName());
-    checkHasRegisteredInfo(state.name);
+    OperationState state(location,
+                         getCheckRegisteredInfo<OpTy>(location.getContext()));
     OpTy::build(*this, state, std::forward<Args>(args)...);
     auto *op = createOperation(state);
     auto result = dyn_cast<OpTy>(op);
     assert(result && "builder didn't return the right type");
     return result;
   }
 
   /// Create an operation of specific op type at the current insertion point,
   /// and immediately try to fold it. This functions populates 'results' with
   /// the results after folding the operation.
   template <typename OpTy, typename... Args>
   void createOrFold(SmallVectorImpl<Value> &results, Location location,
                     Args &&...args) {
     // Create the operation without using 'createOperation' as we don't want to
     // insert it yet.
-    OperationState state(location, OpTy::getOperationName());
-    checkHasRegisteredInfo(state.name);
+    OperationState state(location,
+                         getCheckRegisteredInfo<OpTy>(location.getContext()));
     OpTy::build(*this, state, std::forward<Args>(args)...);
     Operation *op = Operation::create(state);
 
     // Fold the operation. If successful destroy it, otherwise insert it.
     if (succeeded(tryFold(op, results)))
       op->destroy();
     else
       insert(op);

mlir/include/mlir/IR/OperationSupport.h

 /// should only be used by things like the AsmPrinter and other things that need
 /// to be parameterized by generic operation hooks. Most user code should use
 /// the concrete operation types.
 class RegisteredOperationName : public OperationName {
 public:
   /// Lookup the registered operation information for the given operation.
   /// Returns None if the operation isn't registered.
   static Optional<RegisteredOperationName> lookup(StringRef name,
-                                                  MLIRContext *ctx) {
-    return OperationName(name, ctx).getRegisteredInfo();
-  }
+                                                  MLIRContext *ctx);
 
   /// Register a new operation in a Dialect object.
   /// This constructor is used by Dialect objects when they register the list of
   /// operations they contain.
   template <typename T>
   static void insert(Dialect &dialect) {
     insert(T::getOperationName(), dialect, TypeID::get<T>(),
            T::getParseAssemblyFn(), T::getPrintAssemblyFn(),
struct OperationState {
   NamedAttrList attributes;
   /// Successors of this operation and their respective operands.
   SmallVector<Block *, 1> successors;
   /// Regions that the op will hold.
   SmallVector<std::unique_ptr<Region>, 1> regions;
 
 public:
   OperationState(Location location, StringRef name);
-
   OperationState(Location location, OperationName name);
 
+  OperationState(Location location, OperationName name, ValueRange operands,
+                 TypeRange types, ArrayRef<NamedAttribute> attributes,
+                 BlockRange successors = {},
+                 MutableArrayRef<std::unique_ptr<Region>> regions = {});
   OperationState(Location location, StringRef name, ValueRange operands,
                  TypeRange types, ArrayRef<NamedAttribute> attributes,
                  BlockRange successors = {},
                  MutableArrayRef<std::unique_ptr<Region>> regions = {});
 
   void addOperands(ValueRange newOperands);
 
   void addTypes(ArrayRef<Type> newTypes) {

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp

for (Value operand : op->getOperands()) {
     vectorOperands.push_back(vecOperand);
   }
 
   // Create a clone of the op with the proper operands and return types.
   // TODO: The following assumes there is always an op with a fixed
   // name that works both in scalar mode and vector mode.
   // TODO: Is it worth considering an Operation.clone operation which
   // changes the type so we can promote an Operation with less boilerplate?
-  OperationState vecOpState(op->getLoc(), op->getName().getStringRef(),
-                            vectorOperands, vectorTypes, op->getAttrs(),
-                            /*successors=*/{}, /*regions=*/{});
+  OperationState vecOpState(op->getLoc(), op->getName(), vectorOperands,
+                            vectorTypes, op->getAttrs(), /*successors=*/{},
+                            /*regions=*/{});
   Operation *vecOp = state.builder.createOperation(vecOpState);
   state.registerOpVectorReplacement(op, vecOp);
   return vecOp;
 }
 
 /// Vectorizes a yield operation by widening its types. The builder's insertion
 /// point is set after the vectorized parent op to continue vectorizing the
 /// operations after the parent op. When vectorizing a reduction loop a mask may

mlir/lib/Dialect/Vector/VectorUnrollDistribute.cpp

 }
 
 // Clones `op` into a new operations that takes `operands` and returns
 // `resultTypes`.
 static Operation *cloneOpWithOperandsAndTypes(OpBuilder &builder, Location loc,
                                               Operation *op,
                                               ArrayRef<Value> operands,
                                               ArrayRef<Type> resultTypes) {
-  OperationState res(loc, op->getName().getStringRef(), operands, resultTypes,
-                     op->getAttrs());
+  OperationState res(loc, op->getName(), operands, resultTypes, op->getAttrs());
   return builder.createOperation(res);
 }
 
 /// Return the target shape for unrolling for the given `op`. Return llvm::None
 /// if the op shouldn't be or cannot be unrolled.
 static Optional<SmallVector<int64_t, 4>>
 getTargetShape(const vector::UnrollVectorOptions &options, Operation *op) {
   if (options.filterConstraint && failed(options.filterConstraint(op)))

mlir/lib/IR/MLIRContext.cpp

#endif
 
   /// An allocator used for AbstractAttribute and AbstractType objects.
   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.
-  SmallVector<RegisteredOperationName> registeredOperations;
+  llvm::StringMap<RegisteredOperationName> registeredOperations;
 
   /// A mutex used when accessing operation information.
   llvm::sys::SmartRWMutex<true> operationInfoMutex;
 
   //===--------------------------------------------------------------------===//
   // Affine uniquing
   //===--------------------------------------------------------------------===//
 

 
 /// 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.
-  std::vector<RegisteredOperationName> result(
-      impl->registeredOperations.begin(), impl->registeredOperations.end());
+  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;
 }
 
 bool MLIRContext::isOperationRegistered(StringRef name) {
-  return OperationName(name, this).isRegistered();
+  return RegisteredOperationName::lookup(name, this).hasValue();
 }
 
 void Dialect::addType(TypeID typeID, AbstractType &&typeInfo) {
   auto &impl = context->getImpl();
   assert(impl.multiThreadedExecutionContext == 0 &&
          "Registering a new type kind while in a multi-threaded execution "
          "context");
   auto *newInfo =

 //===----------------------------------------------------------------------===//
 
 OperationName::OperationName(StringRef name, MLIRContext *context) {
   MLIRContextImpl &ctxImpl = context->getImpl();
 
   // Check for an existing name in read-only mode.
   bool isMultithreadingEnabled = context->isMultithreadingEnabled();
   if (isMultithreadingEnabled) {
+    // Check the registered info map first. In the overwhelmingly common case,
+    // the entry will be in here and it also removes the need to acquire any
+    // locks.
+    auto registeredIt = ctxImpl.registeredOperations.find(name);
+    if (LLVM_LIKELY(registeredIt != ctxImpl.registeredOperations.end())) {
+      impl = registeredIt->second.impl;
+      return;
+    }
+
     llvm::sys::SmartScopedReader<true> contextLock(ctxImpl.operationInfoMutex);
     auto it = ctxImpl.operations.find(name);
     if (it != ctxImpl.operations.end()) {
       impl = &it->second;
       return;
     }
   }
 
if (Dialect *dialect = getDialect())
     return dialect->getNamespace();
   return getStringRef().split('.').first;
 }
 
 //===----------------------------------------------------------------------===//
 // RegisteredOperationName
 //===----------------------------------------------------------------------===//
 
+Optional<RegisteredOperationName>
+RegisteredOperationName::lookup(StringRef name, MLIRContext *ctx) {
+  auto &impl = ctx->getImpl();
+  auto it = impl.registeredOperations.find(name);
+  if (it != impl.registeredOperations.end())
+    return it->getValue();
+  return llvm::None;
+}
+
 ParseResult
 RegisteredOperationName::parseAssembly(OpAsmParser &parser,
                                        OperationState &result) const {
   return impl->parseAssemblyFn(parser, result);
 }
 
 void RegisteredOperationName::insert(
     StringRef name, Dialect &dialect, TypeID typeID,
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.push_back(RegisteredOperationName(&impl));
+  ctxImpl.registeredOperations.try_emplace(name,
+                                           RegisteredOperationName(&impl));
 
   // 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);

mlir/lib/IR/OperationSupport.cpp

 //===----------------------------------------------------------------------===//
 
 OperationState::OperationState(Location location, StringRef name)
     : location(location), name(name, location->getContext()) {}
 
 OperationState::OperationState(Location location, OperationName name)
     : location(location), name(name) {}
 
-OperationState::OperationState(Location location, StringRef name,
+OperationState::OperationState(Location location, OperationName name,
                                ValueRange operands, TypeRange types,
                                ArrayRef<NamedAttribute> attributes,
                                BlockRange successors,
                                MutableArrayRef<std::unique_ptr<Region>> regions)
-    : location(location), name(name, location->getContext()),
+    : location(location), name(name),
       operands(operands.begin(), operands.end()),
       types(types.begin(), types.end()),
       attributes(attributes.begin(), attributes.end()),
       successors(successors.begin(), successors.end()) {
   for (std::unique_ptr<Region> &r : regions)
     this->regions.push_back(std::move(r));
 }
+OperationState::OperationState(Location location, StringRef name,
+                               ValueRange operands, TypeRange types,
+                               ArrayRef<NamedAttribute> attributes,
+                               BlockRange successors,
+                               MutableArrayRef<std::unique_ptr<Region>> regions)
+    : OperationState(location, OperationName(name, location.getContext()),
+                     operands, types, attributes, successors, regions) {}
 
 void OperationState::addOperands(ValueRange newOperands) {
   operands.append(newOperands.begin(), newOperands.end());
 }
 
 void OperationState::addSuccessors(BlockRange newSuccessors) {
   successors.append(newSuccessors.begin(), newSuccessors.end());
 }