[mlir] Added OpPrintingFlags to AsmState and SSANameState.

[izuk]

Summary

This enables checking the printing flags when formatting names
in SSANameState.

Depends On D105299

[izuk]

Merged some changes that were accidentally in a parent PR.

[izuk]

No change.

[bondhugula]

mlir/lib/IR/AsmPrinter.cpp

836

Field doc comment here please.

1177

Triple /// comment.

2768

Should flags be passed by const ref?

2768–2772

Nit: printerFlags

mlir/lib/Transforms/LocationSnapshot.cpp

25

Should this be passed by const ref?

[izuk]

Applied fixes suggested by Uday.

mlir/include/mlir/IR/AsmState.h

 //===- AsmState.h - State class for AsmPrinter ------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the AsmState class.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_ASMSTATE_H_
 #define MLIR_IR_ASMSTATE_H_
 
+#include "mlir/IR/OperationSupport.h"
 #include "mlir/Support/LLVM.h"
 
 #include <memory>
 
 namespace mlir {
 class Operation;
 
 namespace detail {

 public:
   /// This map represents the raw locations of operations within the output
   /// stream. This maps the original pointer to the operation, to a pair of line
   /// and column in the output stream.
   using LocationMap = DenseMap<Operation *, std::pair<unsigned, unsigned>>;
 
   /// Initialize the asm state at the level of the given operation. A location
   /// map may optionally be provided to be populated when printing.
-  AsmState(Operation *op, LocationMap *locationMap = nullptr);
+  AsmState(Operation *op,
+           const OpPrintingFlags &printerFlags = OpPrintingFlags(),
+           LocationMap *locationMap = nullptr);
   ~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;

mlir/lib/IR/AsmPrinter.cpp

 
 /// This class implements a dummy OpAsmPrinter that doesn't print any output,
 /// and merely collects the attributes and types that *would* be printed in a
 /// normal print invocation so that we can generate proper aliases. This allows
 /// for us to generate aliases only for the attributes and types that would be
 /// in the output, and trims down unnecessary output.
 class DummyAliasOperationPrinter : private OpAsmPrinter {
 public:
-  explicit DummyAliasOperationPrinter(const OpPrintingFlags &flags,
+  explicit DummyAliasOperationPrinter(const OpPrintingFlags &printerFlags,
                                       AliasInitializer &initializer)
-      : printerFlags(flags), initializer(initializer) {}
+      : printerFlags(printerFlags), initializer(initializer) {}
 
   /// Print the given operation.
   void print(Operation *op) {
     // Visit the operation location.
     if (printerFlags.shouldPrintDebugInfo())
       initializer.visit(op->getLoc(), /*canBeDeferred=*/true);
 
     // If requested, always print the generic form.

 
 namespace {
 /// This class manages the state of SSA value names.
 class SSANameState {
 public:
   /// A sentinel value used for values with names set.
   enum : unsigned { NameSentinel = ~0U };
 
-  SSANameState(Operation *op,
+  SSANameState(Operation *op, const OpPrintingFlags &printingFlags,
                DialectInterfaceCollection<OpAsmDialectInterface> &interfaces);
 
   /// Print the SSA identifier for the given value to 'stream'. If
   /// 'printResultNo' is true, it also presents the result number ('#' number)
   /// of this value.
   void printValueID(Value value, bool printResultNo, raw_ostream &stream) const;
 
   /// Return the result indices for each of the result groups registered by this
private:
 
   /// This is the next value ID to assign in numbering.
   unsigned nextValueID = 0;
   /// This is the next ID to assign to a region entry block argument.
   unsigned nextArgumentID = 0;
   /// This is the next ID to assign when a name conflict is detected.
   unsigned nextConflictID = 0;
 
+  OpPrintingFlags printerFlags;
   DialectInterfaceCollection<OpAsmDialectInterface> &interfaces;
 };
 } // end anonymous namespace
 
 SSANameState::SSANameState(
-    Operation *op,
+    Operation *op, const OpPrintingFlags &printerFlags,
     DialectInterfaceCollection<OpAsmDialectInterface> &interfaces)
-    : interfaces(interfaces) {
+    : printerFlags(printerFlags), interfaces(interfaces) {
   llvm::SaveAndRestore<unsigned> valueIDSaver(nextValueID);
   llvm::SaveAndRestore<unsigned> argumentIDSaver(nextArgumentID);
   llvm::SaveAndRestore<unsigned> conflictIDSaver(nextConflictID);
 
   // The naming context includes `nextValueID`, `nextArgumentID`,
   // `nextConflictID` and `usedNames` scoped HashTable. This information is
   // carried from the parent region.
   using UsedNamesScopeTy = llvm::ScopedHashTable<StringRef, char>::ScopeTy;

 //===----------------------------------------------------------------------===//
 // AsmState
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace detail {
 class AsmStateImpl {
 public:
-  explicit AsmStateImpl(Operation *op, AsmState::LocationMap *locationMap)
-      : interfaces(op->getContext()), nameState(op, interfaces),
-        locationMap(locationMap) {}
+  explicit AsmStateImpl(Operation *op, const OpPrintingFlags &printerFlags,
+                        AsmState::LocationMap *locationMap)
+      : interfaces(op->getContext()), nameState(op, printerFlags, interfaces),
+        printerFlags(printerFlags), locationMap(locationMap) {}
 
   /// Initialize the alias state to enable the printing of aliases.
-  void initializeAliases(Operation *op, const OpPrintingFlags &printerFlags) {
+  void initializeAliases(Operation *op) {
     aliasState.initialize(op, printerFlags, interfaces);
   }
 
   /// Get an instance of the OpAsmDialectInterface for the given dialect, or
   /// null if one wasn't registered.
   const OpAsmDialectInterface *getOpAsmInterface(Dialect *dialect) {
     return interfaces.getInterfaceFor(dialect);
   }
private:
   DialectInterfaceCollection<OpAsmDialectInterface> interfaces;
 
   /// The state used for attribute and type aliases.
   AliasState aliasState;
 
   /// The state used for SSA value names.
   SSANameState nameState;
 
+  // Flags that control op output.
+  OpPrintingFlags printerFlags;
+
   /// An optional location map to be populated.
   AsmState::LocationMap *locationMap;
 };
 } // end namespace detail
 } // end namespace mlir
 
-AsmState::AsmState(Operation *op, LocationMap *locationMap)
-    : impl(std::make_unique<AsmStateImpl>(op, locationMap)) {}
+AsmState::AsmState(Operation *op, const OpPrintingFlags &printerFlags,
+                   LocationMap *locationMap)
+    : impl(std::make_unique<AsmStateImpl>(op, printerFlags, locationMap)) {}
 AsmState::~AsmState() {}
 
 //===----------------------------------------------------------------------===//
 // ModulePrinter
 //===----------------------------------------------------------------------===//
 
 namespace {
 class ModulePrinter {
void Value::printAsOperand(raw_ostream &os, AsmState &state) {
   // TODO: 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) {
   // If this is a top level operation, we also print aliases.
   if (!getParent() && !flags.shouldUseLocalScope()) {
-    AsmState state(this);
-    state.getImpl().initializeAliases(this, flags);
+    AsmState state(this, flags);
+    state.getImpl().initializeAliases(this);
     print(os, state, flags);
     return;
   }
 
   // Find the operation to number from based upon the provided flags.
   Operation *op = this;
   bool shouldUseLocalScope = flags.shouldUseLocalScope();
   do {
     // If we are printing local scope, stop at the first operation that is
     // isolated from above.
     if (shouldUseLocalScope && op->hasTrait<OpTrait::IsIsolatedFromAbove>())
       break;
 
     // Otherwise, traverse up to the next parent.
     Operation *parentOp = op->getParentOp();
     if (!parentOp)
       break;
     op = parentOp;
   } while (true);
 
-  AsmState state(op);
+  AsmState state(op, flags);
   print(os, state, flags);
 }
 void Operation::print(raw_ostream &os, AsmState &state, OpPrintingFlags flags) {
   OperationPrinter printer(os, flags, state.getImpl());
   if (!getParent() && !flags.shouldUseLocalScope())
     printer.printTopLevelOperation(this);
   else
     printer.print(this);

mlir/lib/Transforms/LocationSnapshot.cpp

 using namespace mlir;
 
 /// This function generates new locations from the given IR by snapshotting the
 /// IR to the given stream, and using the printed locations within that stream.
 /// If a 'tag' is non-empty, the generated locations are represented as a
 /// NameLoc with the given tag as the name, and then fused with the existing
 /// locations. Otherwise, the existing locations are replaced.
 static void generateLocationsFromIR(raw_ostream &os, StringRef fileName,
                                     Operation *op, OpPrintingFlags flags,
                                     StringRef tag) {
   // Print the IR to the stream, and collect the raw line+column information.
   AsmState::LocationMap opToLineCol;
-  AsmState state(op, &opToLineCol);
+  AsmState state(op, flags, &opToLineCol);
   op->print(os, state, flags);
 
   Builder builder(op->getContext());
   Optional<Identifier> tagIdentifier;
   if (!tag.empty())
     tagIdentifier = builder.getIdentifier(tag);
 
   // Walk and generate new locations for each of the operations.