[mlir] Refactor ModuleState into AsmState and expose it to users.

Summary:
This allows for users to cache printer state, which can be costly to recompute. Each of the IR print methods gain a new overload taking this new state class.

Depends On D72293

Reviewed By: jpienaar

Differential Revision: https://reviews.llvm.org/D72294

6 files changed, 67 insertions, 0 deletions

diff --git a/mlir/include/mlir/IR/AsmState.h b/mlir/include/mlir/IR/AsmState.h
new file mode 100644
index 00000000000..16c764780c3
--- /dev/null
+++ b/mlir/include/mlir/IR/AsmState.h
@@ -0,0 +1,52 @@
+//===- AsmState.h - State class for AsmPrinter ------------------*- C++ -*-===//
+//
+// Part of the MLIR 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the AsmState class.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_ASMSTATE_H_
+#define MLIR_IR_ASMSTATE_H_
+
+#include <memory>
+
+namespace mlir {
+class Operation;
+
+namespace detail {
+class AsmStateImpl;
+} // end namespace detail
+
+/// This class provides management for the lifetime of the state used when
+/// printing the IR. It allows for alleviating the cost of recomputing the
+/// internal state of the asm printer.
+///
+/// The IR should not be mutated in-between invocations using this state, and
+/// the IR being printed must not be an parent of the IR originally used to
+/// initialize this state. This means that if a child operation is provided, a
+/// parent operation cannot reuse this state.
+class AsmState {
+public:
+  /// Initialize the asm state at the level of the given operation.
+  AsmState(Operation *op);
+  ~AsmState();
+
+  /// Return an instance of the internal implementation. Returns nullptr if the
+  /// state has not been initialized.
+  detail::AsmStateImpl &getImpl() { return *impl; }
+
+private:
+  AsmState() = delete;
+
+  /// A pointer to allocated storage for the impl state.
+  std::unique_ptr<detail::AsmStateImpl> impl;
+};
+
+} // end namespace mlir
+
+#endif // MLIR_IR_ASMSTATE_H_
diff --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h
index 2d3eb18d729..eae83ccb634 100644
--- a/mlir/include/mlir/IR/Block.h
+++ b/mlir/include/mlir/IR/Block.h
@@ -312,12 +312,14 @@ public:
   }
 
   void print(raw_ostream &os);
+  void print(raw_ostream &os, AsmState &state);
   void dump();
 
   /// Print out the name of the block without printing its body.
   /// NOTE: The printType argument is ignored.  We keep it for compatibility
   /// with LLVM dominator machinery that expects it to exist.
   void printAsOperand(raw_ostream &os, bool printType = true);
+  void printAsOperand(raw_ostream &os, AsmState &state);
 
 private:
   /// Pair of the parent object that owns this block and a bit that signifies if
diff --git a/mlir/include/mlir/IR/Module.h b/mlir/include/mlir/IR/Module.h
index babc51aad0d..fb7c61f3bac 100644
--- a/mlir/include/mlir/IR/Module.h
+++ b/mlir/include/mlir/IR/Module.h
@@ -57,6 +57,8 @@ public:
 
   /// Print the this module in the custom top-level form.
   void print(raw_ostream &os, OpPrintingFlags flags = llvm::None);
+  void print(raw_ostream &os, AsmState &state,
+             OpPrintingFlags flags = llvm::None);
   void dump();
 
   //===--------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index cc4dfd31190..87cf1a59209 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -122,6 +122,10 @@ public:
   void print(raw_ostream &os, OpPrintingFlags flags = llvm::None) {
     state->print(os, flags);
   }
+  void print(raw_ostream &os, AsmState &asmState,
+             OpPrintingFlags flags = llvm::None) {
+    state->print(os, asmState, flags);
+  }
 
   /// Dump this operation.
   void dump() { state->dump(); }
diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 81fd83b3d25..abb1bf2ac3a 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -187,6 +187,8 @@ public:
   bool isBeforeInBlock(Operation *other);
 
   void print(raw_ostream &os, OpPrintingFlags flags = llvm::None);
+  void print(raw_ostream &os, AsmState &state,
+             OpPrintingFlags flags = llvm::None);
   void dump();
 
   //===--------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/IR/Value.h b/mlir/include/mlir/IR/Value.h
index a2be7bb8a71..6354156248d 100644
--- a/mlir/include/mlir/IR/Value.h
+++ b/mlir/include/mlir/IR/Value.h
@@ -18,6 +18,7 @@
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
+class AsmState;
 class BlockArgument;
 class Operation;
 class OpResult;
@@ -172,8 +173,12 @@ public:
   Kind getKind() const { return ownerAndKind.getInt(); }
 
   void print(raw_ostream &os);
+  void print(raw_ostream &os, AsmState &state);
   void dump();
 
+  /// Print this value as if it were an operand.
+  void printAsOperand(raw_ostream &os, AsmState &state);
+
   /// Methods for supporting PointerLikeTypeTraits.
   void *getAsOpaquePointer() const { return ownerAndKind.getOpaqueValue(); }
   static Value getFromOpaquePointer(const void *pointer) {