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

[rriddle]

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

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[rriddle]

Tidy comment in AsmState.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[rriddle]

Rebase.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[jpienaar]

Looks good on first walk through

mlir/include/mlir/IR/AsmState.h

8

File description?

27

What does higher level mean here? It is a bit overloaded term and I think you have a very explicit meaning in mind.

42

So the lifetime of the AsmStateImpl is linked to this object, and unique_ptr is used to enable only forward declaring it above?

mlir/lib/IR/AsmPrinter.cpp

2205

How is guaranteed that getImpl doesn't return null here? (e.g., verify that AsmState has been initialized)

[rriddle]

Address comments.

mlir/include/mlir/IR/AsmState.h

27

Switched to 'parent'.

42

Yes, this is the pImpl idiom.

mlir/lib/IR/AsmPrinter.cpp

2205

Ooops, meant to delete the default constructor. Thanks for the catch.

[merge_guards_bot]

Unit tests: pass. 61744 tests passed, 0 failed and 780 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[jpienaar]

Looks good thanks

[rriddle]

Rebase.

[merge_guards_bot]

Unit tests: pass. 61804 tests passed, 0 failed and 781 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

mlir/include/mlir/IR/AsmState.h

+//===- 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_

mlir/include/mlir/IR/Block.h

public:
   }
 
   /// Returns pointer to member of operation list.
   static OpListType Block::*getSublistAccess(Operation *) {
     return &Block::operations;
   }
 
   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
   /// the operations within this block have a valid ordering.
   llvm::PointerIntPair<Region *, /*IntBits=*/1, bool> parentValidOpOrderPair;
 
   /// This is the list of operations in the block.
   OpListType operations;

mlir/include/mlir/IR/Module.h

public:
   Region &getBodyRegion();
   Block *getBody();
 
   /// Return the name of this module if present.
   Optional<StringRef> getName();
 
   /// 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();
 
   //===--------------------------------------------------------------------===//
   // Body Management.
   //===--------------------------------------------------------------------===//
 
   /// Iteration over the operations in the module.
   using iterator = Block::iterator;

mlir/include/mlir/IR/OpDefinition.h

public:
 
   /// Return the context this operation belongs to.
   MLIRContext *getContext() { return getOperation()->getContext(); }
 
   /// Print the operation to the given stream.
   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(); }
 
   /// The source location the operation was defined or derived from.
   Location getLoc() { return state->getLoc(); }
   void setLoc(Location loc) { state->setLoc(loc); }
 

mlir/include/mlir/IR/Operation.h

public:
   /// Given an operation 'other' that is within the same parent block, return
   /// whether the current operation is before 'other' in the operation list
   /// of the parent block.
   /// Note: This function has an average complexity of O(1), but worst case may
   /// take O(N) where N is the number of operations within the parent block.
   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();
 
   //===--------------------------------------------------------------------===//
   // Operands
   //===--------------------------------------------------------------------===//
 
   /// Returns if the operation has a resizable operation list, i.e. operands can
   /// be added.

mlir/include/mlir/IR/Value.h

 #ifndef MLIR_IR_VALUE_H
 #define MLIR_IR_VALUE_H
 
 #include "mlir/IR/Types.h"
 #include "mlir/IR/UseDefLists.h"
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
+class AsmState;
 class BlockArgument;
 class Operation;
 class OpResult;
 class Region;
 class Value;
 
 namespace detail {
 /// The internal implementation of a BlockArgument.
public:
 
   //===--------------------------------------------------------------------===//
   // Utilities
 
   /// Returns the kind of this value.
   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) {
     Value value;
     value.ownerAndKind.setFromOpaqueValue(const_cast<void *>(pointer));
     return value;
   }
 

mlir/lib/IR/AsmPrinter.cpp

 //===- AsmPrinter.cpp - MLIR Assembly Printer Implementation --------------===//
 //
 // 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 implements the MLIR AsmPrinter class, which is used to implement
 // the various print() methods on the core IR objects.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/AsmState.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Support/STLExtras.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopedHashTable.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/SaveAndRestore.h"
 using namespace mlir;
+using namespace mlir::detail;
 
 void Identifier::print(raw_ostream &os) const { os << str(); }
 
 void Identifier::dump() const { print(llvm::errs()); }
 
 void OperationName::print(raw_ostream &os) const { os << getStringRef(); }
 
 void OperationName::dump() const { print(llvm::errs()); }
if (!usedNames.count(name)) {
     }
   }
 
   usedNames.insert(name, char());
   return name;
 }
 
 //===----------------------------------------------------------------------===//
-// ModuleState
+// AsmState
 //===----------------------------------------------------------------------===//
 
-namespace {
-class ModuleState {
+namespace mlir {
+namespace detail {
+class AsmStateImpl {
 public:
-  explicit ModuleState(Operation *op)
+  explicit AsmStateImpl(Operation *op)
       : interfaces(op->getContext()), nameState(op, interfaces) {}
 
   /// Initialize the alias state to enable the printing of aliases.
   void initializeAliases(Operation *op) {
     aliasState.initialize(op, interfaces);
   }
 
   /// Get an instance of the OpAsmDialectInterface for the given dialect, or
private:
   DialectInterfaceCollection<OpAsmDialectInterface> interfaces;
 
   /// The state used for attribute and type aliases.
   AliasState aliasState;
 
   /// The state used for SSA value names.
   SSANameState nameState;
 };
-} // end anonymous namespace
+} // end namespace detail
+} // end namespace mlir
+
+AsmState::AsmState(Operation *op) : impl(std::make_unique<AsmStateImpl>(op)) {}
+AsmState::~AsmState() {}
 
 //===----------------------------------------------------------------------===//
 // ModulePrinter
 //===----------------------------------------------------------------------===//
 
 namespace {
 class ModulePrinter {
 public:
   ModulePrinter(raw_ostream &os, OpPrintingFlags flags = llvm::None,
-                ModuleState *state = nullptr)
+                AsmStateImpl *state = nullptr)
       : os(os), printerFlags(flags), state(state) {}
   explicit ModulePrinter(ModulePrinter &printer)
       : os(printer.os), printerFlags(printer.printerFlags),
         state(printer.state) {}
 
   /// Returns the output stream of the printer.
   raw_ostream &getStream() { return os; }
 
   template <typename Container, typename UnaryFunctor>
   inline void interleaveComma(const Container &c, UnaryFunctor each_fn) const {
     mlir::interleaveComma(c, os, each_fn);
   }
 
-  void print(ModuleOp module);
-
   /// Print the given attribute. If 'mayElideType' is true, some attributes are
   /// printed without the type when the type matches the default used in the
   /// parser (for example i64 is the default for integer attributes).
   void printAttribute(Attribute attr, bool mayElideType = false);
 
   void printType(Type type);
   void printLocation(LocationAttr loc);
 
protected:
 
   /// The output stream for the printer.
   raw_ostream &os;
 
   /// A set of flags to control the printer's behavior.
   OpPrintingFlags printerFlags;
 
   /// An optional printer state for the module.
-  ModuleState *state;
+  AsmStateImpl *state;
 };
 } // end anonymous namespace
 
 void ModulePrinter::printTrailingLocation(Location loc) {
   // Check to see if we are printing debug information.
   if (!printerFlags.shouldPrintDebugInfo())
     return;
 

 //===----------------------------------------------------------------------===//
 // OperationPrinter
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class contains the logic for printing operations, regions, and blocks.
 class OperationPrinter : public ModulePrinter, private OpAsmPrinter {
 public:
-  explicit OperationPrinter(ModulePrinter &other) : ModulePrinter(other) {
-    assert(state && "expected valid state when printing operation");
-  }
+  explicit OperationPrinter(raw_ostream &os, OpPrintingFlags flags,
+                            AsmStateImpl &state)
+      : ModulePrinter(os, flags, &state) {}
 
+  /// Print the given top-level module.
+  void print(ModuleOp op);
   /// Print the given operation with its indent and location.
   void print(Operation *op);
   /// Print the bare location, not including indentation/location/etc.
   void printOperation(Operation *op);
   /// Print the given operation in the generic form.
   void printGenericOp(Operation *op) override;
 
   /// Print the name of the given block.
private:
   /// The number of spaces used for indenting nested operations.
   const static unsigned indentWidth = 2;
 
   // This is the current indentation level for nested structures.
   unsigned currentIndent = 0;
 };
 } // end anonymous namespace
 
+void OperationPrinter::print(ModuleOp op) {
+  // Output the aliases at the top level.
+  state->getAliasState().printAttributeAliases(os);
+  state->getAliasState().printTypeAliases(os);
+
+  // Print the module.
+  print(op.getOperation());
+}
+
 void OperationPrinter::print(Operation *op) {
   os.indent(currentIndent);
   printOperation(op);
   printTrailingLocation(op->getLoc());
 }
 
 void OperationPrinter::printOperation(Operation *op) {
   if (size_t numResults = op->getNumResults()) {
if (isSymbol)
       os << ')';
   };
 
   interleaveComma(map.getResults(), [&](AffineExpr expr) {
     printAffineExpr(expr, printValueName);
   });
 }
 
-void ModulePrinter::print(ModuleOp module) {
-  assert(state && "expected valid state when printing an operation");
-
-  // Output the aliases at the top level.
-  state->getAliasState().printAttributeAliases(os);
-  state->getAliasState().printTypeAliases(os);
-
-  // Print the module.
-  OperationPrinter(*this).print(module);
-  os << '\n';
-}
-
 //===----------------------------------------------------------------------===//
 // print and dump methods
 //===----------------------------------------------------------------------===//
 
 void Attribute::print(raw_ostream &os) const {
   ModulePrinter(os).printAttribute(*this);
 }
 

 
 void Value::print(raw_ostream &os) {
   if (auto *op = getDefiningOp())
     return op->print(os);
   // TODO: Improve this.
   assert(isa<BlockArgument>());
   os << "<block argument>\n";
 }
+void Value::print(raw_ostream &os, AsmState &state) {
+  if (auto *op = getDefiningOp())
+    return op->print(os, state);
+
+  // TODO: Improve this.
+  assert(isa<BlockArgument>());
+  os << "<block argument>\n";
+}
 
 void Value::dump() {
   print(llvm::errs());
   llvm::errs() << "\n";
 }
 
+void Value::printAsOperand(raw_ostream &os, AsmState &state) {
+  // TODO(riverriddle) This doesn't necessarily capture all potential cases.
+  // Currently, region arguments can be shadowed when printing the main
+  // operation. If the IR hasn't been printed, this will produce the old SSA
+  // name and not the shadowed name.
+  state.getImpl().getSSANameState().printValueID(*this, /*printResultNo=*/true,
+                                                 os);
+}
+
 void Operation::print(raw_ostream &os, OpPrintingFlags flags) {
   // Handle top-level operations or local printing.
   if (!getParent() || flags.shouldUseLocalScope()) {
-    ModuleState state(this);
-    ModulePrinter modulePrinter(os, flags, &state);
-    OperationPrinter(modulePrinter).print(this);
+    AsmState state(this);
+    OperationPrinter(os, flags, state.getImpl()).print(this);
     return;
   }
 
   Operation *parentOp = getParentOp();
   if (!parentOp) {
     os << "<<UNLINKED OPERATION>>\n";
     return;
   }
   // Get the top-level op.
   while (auto *nextOp = parentOp->getParentOp())
     parentOp = nextOp;
 
-  ModuleState state(parentOp);
-  ModulePrinter modulePrinter(os, flags, &state);
-  OperationPrinter(modulePrinter).print(this);
+  AsmState state(parentOp);
+  print(os, state, flags);
+}
+void Operation::print(raw_ostream &os, AsmState &state, OpPrintingFlags flags) {
+  OperationPrinter(os, flags, state.getImpl()).print(this);
 }
 
 void Operation::dump() {
   print(llvm::errs(), OpPrintingFlags().useLocalScope());
   llvm::errs() << "\n";
 }
 
 void Block::print(raw_ostream &os) {
   Operation *parentOp = getParentOp();
   if (!parentOp) {
     os << "<<UNLINKED BLOCK>>\n";
     return;
   }
   // Get the top-level op.
   while (auto *nextOp = parentOp->getParentOp())
     parentOp = nextOp;
 
-  ModuleState state(parentOp);
-  ModulePrinter modulePrinter(os, /*flags=*/llvm::None, &state);
-  OperationPrinter(modulePrinter).print(this);
+  AsmState state(parentOp);
+  print(os, state);
+}
+void Block::print(raw_ostream &os, AsmState &state) {
+  OperationPrinter(os, /*flags=*/llvm::None, state.getImpl()).print(this);
 }
 
 void Block::dump() { print(llvm::errs()); }
 
 /// Print out the name of the block without printing its body.
 void Block::printAsOperand(raw_ostream &os, bool printType) {
   Operation *parentOp = getParentOp();
   if (!parentOp) {
     os << "<<UNLINKED BLOCK>>\n";
     return;
   }
   // Get the top-level op.
   while (auto *nextOp = parentOp->getParentOp())
     parentOp = nextOp;
 
-  ModuleState state(parentOp);
-  ModulePrinter modulePrinter(os, /*flags=*/llvm::None, &state);
-  OperationPrinter(modulePrinter).printBlockName(this);
+  AsmState state(parentOp);
+  printAsOperand(os, state);
+}
+void Block::printAsOperand(raw_ostream &os, AsmState &state) {
+  OperationPrinter printer(os, /*flags=*/llvm::None, state.getImpl());
+  printer.printBlockName(this);
 }
 
 void ModuleOp::print(raw_ostream &os, OpPrintingFlags flags) {
-  ModuleState state(*this);
+  AsmState state(*this);
 
   // Don't populate aliases when printing at local scope.
   if (!flags.shouldUseLocalScope())
-    state.initializeAliases(*this);
-  ModulePrinter(os, flags, &state).print(*this);
+    state.getImpl().initializeAliases(*this);
+  print(os, state, flags);
+}
+void ModuleOp::print(raw_ostream &os, AsmState &state, OpPrintingFlags flags) {
+  OperationPrinter(os, flags, state.getImpl()).print(*this);
 }
 
 void ModuleOp::dump() { print(llvm::errs()); }