From 496f9a7d8d4b61b6f5dfa21e2614043a2b27cf67 Mon Sep 17 00:00:00 2001
From: Jeff Niu <jeff@modular.com>
Date: Thu, 8 Dec 2022 12:22:40 -0800
Subject: [PATCH] [mlir][analysis] Add an analysis for preserving symbol tables

This patch adds a `SymbolTableAnalysis` that can be used with the
analysis manager. It contains a symbol table collection. This analysis
allows symbol tables to be preserved across passes so that they do not
need to be recomputed. The analysis assumes it remains valid because
most transformations automatically keep symbol tables up-to-date using
its `insert` and `erase` methods.

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D139666
---
 mlir/include/mlir/Analysis/SymbolTableAnalysis.h | 54 ++++++++++++++++++++++++
 1 file changed, 54 insertions(+)
 create mode 100644 mlir/include/mlir/Analysis/SymbolTableAnalysis.h

diff --git a/mlir/include/mlir/Analysis/SymbolTableAnalysis.h b/mlir/include/mlir/Analysis/SymbolTableAnalysis.h
new file mode 100644
index 0000000..145627e
--- /dev/null
+++ b/mlir/include/mlir/Analysis/SymbolTableAnalysis.h
@@ -0,0 +1,54 @@
+//===- SymbolTableAnalysis.h - Analysis for cached symbol tables --*- C++ -*-=//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_SYMBOLTABLEANALYSIS_H
+#define MLIR_ANALYSIS_SYMBOLTABLEANALYSIS_H
+
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/AnalysisManager.h"
+
+namespace mlir {
+/// This is a simple analysis that contains a symbol table collection and, for
+/// simplicity, a reference to the top-level symbol table. This allows symbol
+/// tables to be preserved across passes. Most often, symbol tables are
+/// automatically kept up-to-date via the `insert` and `erase` functions.
+class SymbolTableAnalysis {
+public:
+  /// Create the symbol table analysis at the provided top-level operation and
+  /// instantiate the symbol table of the top-level operation.
+  SymbolTableAnalysis(Operation *op)
+      : topLevelSymbolTable(symbolTables.getSymbolTable(op)) {}
+
+  /// Get the symbol table collection.
+  SymbolTableCollection &getSymbolTables() { return symbolTables; }
+
+  /// Get the top-level symbol table.
+  SymbolTable &getTopLevelSymbolTable() { return topLevelSymbolTable; }
+
+  /// Get the top-level operation.
+  template <typename OpT>
+  OpT getTopLevelOp() {
+    return cast<OpT>(topLevelSymbolTable.getOp());
+  }
+
+  /// Symbol tables are kept up-to-date by passes. Assume that the analysis
+  /// remains valid.
+  bool isInvalidated(const AnalysisManager::PreservedAnalyses &pa) {
+    return false;
+  }
+
+private:
+  /// The symbol table collection containing cached symbol tables for all nested
+  /// symbol table operations.
+  SymbolTableCollection symbolTables;
+  /// The symbol table of the top-level operation.
+  SymbolTable &topLevelSymbolTable;
+};
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_SYMBOLTABLEANALYSIS_H
-- 
2.7.4