From 73547b08de59d1be0e65d8dbc989e9c186295fcd Mon Sep 17 00:00:00 2001
From: River Riddle <riddleriver@gmail.com>
Date: Wed, 28 Oct 2020 22:00:55 -0700
Subject: [PATCH] [mlir][SymbolTable] Small optimization to walking symbol
 references

* Check region count for unknown symbol tables first, as it is a faster check
* Add an accessor to MutableDictionaryAttr to get the internal dictionary without creating a new one if it is empty. This avoids an otherwise unnecessary lookup of an MLIRContext.
---
 mlir/include/mlir/IR/Attributes.h | 4 ++++
 mlir/lib/IR/SymbolTable.cpp       | 6 +++---
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/IR/Attributes.h b/mlir/include/mlir/IR/Attributes.h
index b4a1473e1479..4a8cc562cb52 100644
--- a/mlir/include/mlir/IR/Attributes.h
+++ b/mlir/include/mlir/IR/Attributes.h
@@ -1621,6 +1621,10 @@ public:
   /// Return the underlying dictionary attribute.
   DictionaryAttr getDictionary(MLIRContext *context) const;
 
+  /// Return the underlying dictionary attribute or null if there are no
+  /// attributes within this dictionary.
+  DictionaryAttr getDictionaryOrNull() const { return attrs; }
+
   /// Return all of the attributes on this operation.
   ArrayRef<NamedAttribute> getAttrs() const;
 
diff --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp
index d4377d6338c0..90b4a9dac4e7 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -17,7 +17,7 @@ using namespace mlir;
 /// Return true if the given operation is unknown and may potentially define a
 /// symbol table.
 static bool isPotentiallyUnknownSymbolTable(Operation *op) {
-  return !op->getDialect() && op->getNumRegions() == 1;
+  return op->getNumRegions() == 1 && !op->getDialect();
 }
 
 /// Returns the string name of the given symbol, or None if this is not a
@@ -442,9 +442,9 @@ static WalkResult walkSymbolRefs(
     Operation *op,
     function_ref<WalkResult(SymbolTable::SymbolUse, ArrayRef<int>)> callback) {
   // Check to see if the operation has any attributes.
-  if (op->getMutableAttrDict().empty())
+  DictionaryAttr attrDict = op->getMutableAttrDict().getDictionaryOrNull();
+  if (!attrDict)
     return WalkResult::advance();
-  DictionaryAttr attrDict = op->getAttrDictionary();
 
   // A worklist of a container attribute and the current index into the held
   // attribute list.