[MLIR] Add OpPrintingFlags to IRPrinterConfig.

[jurahul]

Summary

• This will enable tweaking IR printing options when enabling printing (for ex, tweak elideLargeElementsAttrs to create smaller IR logs)

mlir/include/mlir/Pass/PassManager.h

 //===- PassManager.h - Pass Management Interface ----------------*- 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_PASS_PASSMANAGER_H
 #define MLIR_PASS_PASSMANAGER_H
 
+#include "mlir/IR/OperationSupport.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include <functional>
 #include <vector>
public:
     /// * 'printModuleScope' signals if the top-level module IR should always be
     ///   printed. This should only be set to true when multi-threading is
     ///   disabled, otherwise we may try to print IR that is being modified
     ///   asynchronously.
     /// * 'printAfterOnlyOnChange' signals that when printing the IR after a
     ///   pass, in the case of a non-failure, we should first check if any
     ///   potential mutations were made. This allows for reducing the number of
     ///   logs that don't contain meaningful changes.
-    explicit IRPrinterConfig(bool printModuleScope = false,
-                             bool printAfterOnlyOnChange = false);
+    /// * 'opPrintingFlags' sets up the printing flags to use when printing the
+    ///   IR.
+    explicit IRPrinterConfig(
+        bool printModuleScope = false, bool printAfterOnlyOnChange = false,
+        OpPrintingFlags opPrintingFlags = OpPrintingFlags());
     virtual ~IRPrinterConfig();
 
     /// A hook that may be overridden by a derived config that checks if the IR
     /// of 'operation' should be dumped *before* the pass 'pass' has been
     /// executed. If the IR should be dumped, 'printCallback' should be invoked
     /// with the stream to dump into.
     virtual void printBeforeIfEnabled(Pass *pass, Operation *operation,
                                       PrintCallbackFn printCallback);
 
     /// A hook that may be overridden by a derived config that checks if the IR
     /// of 'operation' should be dumped *after* the pass 'pass' has been
     /// executed. If the IR should be dumped, 'printCallback' should be invoked
     /// with the stream to dump into.
     virtual void printAfterIfEnabled(Pass *pass, Operation *operation,
                                      PrintCallbackFn printCallback);
 
     /// Returns true if the IR should always be printed at the top-level scope.
     bool shouldPrintAtModuleScope() const { return printModuleScope; }
 
     /// Returns true if the IR should only printed after a pass if the IR
     /// "changed".
     bool shouldPrintAfterOnlyOnChange() const { return printAfterOnlyOnChange; }
 
+    /// Returns the printing flags to be used to print the IR.
+    OpPrintingFlags getOpPrintingFlags() const { return opPrintingFlags; }
+
   private:
     /// A flag that indicates if the IR should be printed at module scope.
     bool printModuleScope;
 
     /// A flag that indicates that the IR after a pass should only be printed if
     /// a change is detected.
     bool printAfterOnlyOnChange;
+
+    /// Flags to control printing behavior.
+    OpPrintingFlags opPrintingFlags;
   };
 
   /// Add an instrumentation to print the IR before and after pass execution,
   /// using the provided configuration.
   void enableIRPrinting(std::unique_ptr<IRPrinterConfig> config);
 
   /// Add an instrumentation to print the IR before and after pass execution,
   /// using the provided fields to generate a default configuration:
   /// * 'shouldPrintBeforePass' and 'shouldPrintAfterPass' correspond to filter
   ///   functions that take a 'Pass *' and `Operation *`. These function should
   ///   return true if the IR should be printed or not.
   /// * 'printModuleScope' signals if the module IR should be printed, even
   ///   for non module passes.
   /// * 'printAfterOnlyOnChange' signals that when printing the IR after a
   ///   pass, in the case of a non-failure, we should first check if any
   ///   potential mutations were made.
+  /// * 'opPrintingFlags' sets up the printing flags to use when printing the
+  ///   IR.
   /// * 'out' corresponds to the stream to output the printed IR to.
   void enableIRPrinting(
       std::function<bool(Pass *, Operation *)> shouldPrintBeforePass =
           [](Pass *, Operation *) { return true; },
       std::function<bool(Pass *, Operation *)> shouldPrintAfterPass =
           [](Pass *, Operation *) { return true; },
       bool printModuleScope = true, bool printAfterOnlyOnChange = true,
-      raw_ostream &out = llvm::errs());
+      raw_ostream &out = llvm::errs(),
+      OpPrintingFlags opPrintingFlags = OpPrintingFlags());
 
   //===--------------------------------------------------------------------===//
   // Pass Timing
 
   /// A configuration struct provided to the pass timing feature.
   class PassTimingConfig {
   public:
     using PrintCallbackFn = function_ref<void(raw_ostream &)>;

mlir/lib/Pass/IRPrinting.cpp

void IRPrinterInstrumentation::runBeforePass(Pass *pass, Operation *op) {
   if (isHiddenPass(pass))
     return;
   // If the config asked to detect changes, record the current fingerprint.
   if (config->shouldPrintAfterOnlyOnChange())
     beforePassFingerPrints.try_emplace(pass, op);
 
   config->printBeforeIfEnabled(pass, op, [&](raw_ostream &out) {
     out << formatv("// *** IR Dump Before {0} ***", pass->getName());
-    printIR(op, config->shouldPrintAtModuleScope(), out, OpPrintingFlags());
+    printIR(op, config->shouldPrintAtModuleScope(), out,
+            config->getOpPrintingFlags());
     out << "\n\n";
   });
 }
 
 void IRPrinterInstrumentation::runAfterPass(Pass *pass, Operation *op) {
   if (isHiddenPass(pass))
     return;
   // If the config asked to detect changes, compare the current fingerprint with
   // the previous.
   if (config->shouldPrintAfterOnlyOnChange()) {
     auto fingerPrintIt = beforePassFingerPrints.find(pass);
     assert(fingerPrintIt != beforePassFingerPrints.end() &&
            "expected valid fingerprint");
     // If the fingerprints are the same, we don't print the IR.
     if (fingerPrintIt->second == OperationFingerPrint(op)) {
       beforePassFingerPrints.erase(fingerPrintIt);
       return;
     }
     beforePassFingerPrints.erase(fingerPrintIt);
   }
 
   config->printAfterIfEnabled(pass, op, [&](raw_ostream &out) {
     out << formatv("// *** IR Dump After {0} ***", pass->getName());
-    printIR(op, config->shouldPrintAtModuleScope(), out, OpPrintingFlags());
+    printIR(op, config->shouldPrintAtModuleScope(), out,
+            config->getOpPrintingFlags());
     out << "\n\n";
   });
 }
 
 void IRPrinterInstrumentation::runAfterPassFailed(Pass *pass, Operation *op) {
   if (isa<OpToOpPassAdaptor>(pass))
     return;
   if (config->shouldPrintAfterOnlyOnChange())
     beforePassFingerPrints.erase(pass);
 
   config->printAfterIfEnabled(pass, op, [&](raw_ostream &out) {
     out << formatv("// *** IR Dump After {0} Failed ***", pass->getName());
     printIR(op, config->shouldPrintAtModuleScope(), out,
             OpPrintingFlags().printGenericOpForm());
     out << "\n\n";
   });
 }
 
 //===----------------------------------------------------------------------===//
 // IRPrinterConfig
 //===----------------------------------------------------------------------===//
 
 /// Initialize the configuration.
 PassManager::IRPrinterConfig::IRPrinterConfig(bool printModuleScope,
-                                              bool printAfterOnlyOnChange)
+                                              bool printAfterOnlyOnChange,
+                                              OpPrintingFlags opPrintingFlags)
     : printModuleScope(printModuleScope),
-      printAfterOnlyOnChange(printAfterOnlyOnChange) {}
+      printAfterOnlyOnChange(printAfterOnlyOnChange),
+      opPrintingFlags(opPrintingFlags) {}
 PassManager::IRPrinterConfig::~IRPrinterConfig() {}
 
 /// A hook that may be overridden by a derived config that checks if the IR
 /// of 'operation' should be dumped *before* the pass 'pass' has been
 /// executed. If the IR should be dumped, 'printCallback' should be invoked
 /// with the stream to dump into.
 void PassManager::IRPrinterConfig::printBeforeIfEnabled(
     Pass *pass, Operation *operation, PrintCallbackFn printCallback) {

 //===----------------------------------------------------------------------===//
 
 namespace {
 /// Simple wrapper config that allows for the simpler interface defined above.
 struct BasicIRPrinterConfig : public PassManager::IRPrinterConfig {
   BasicIRPrinterConfig(
       std::function<bool(Pass *, Operation *)> shouldPrintBeforePass,
       std::function<bool(Pass *, Operation *)> shouldPrintAfterPass,
-      bool printModuleScope, bool printAfterOnlyOnChange, raw_ostream &out)
-      : IRPrinterConfig(printModuleScope, printAfterOnlyOnChange),
+      bool printModuleScope, bool printAfterOnlyOnChange,
+      OpPrintingFlags opPrintingFlags, raw_ostream &out)
+      : IRPrinterConfig(printModuleScope, printAfterOnlyOnChange,
+                        opPrintingFlags),
         shouldPrintBeforePass(shouldPrintBeforePass),
         shouldPrintAfterPass(shouldPrintAfterPass), out(out) {
     assert((shouldPrintBeforePass || shouldPrintAfterPass) &&
            "expected at least one valid filter function");
   }
 
   void printBeforeIfEnabled(Pass *pass, Operation *operation,
                             PrintCallbackFn printCallback) final {
void PassManager::enableIRPrinting(std::unique_ptr<IRPrinterConfig> config) {
   addInstrumentation(
       std::make_unique<IRPrinterInstrumentation>(std::move(config)));
 }
 
 /// Add an instrumentation to print the IR before and after pass execution.
 void PassManager::enableIRPrinting(
     std::function<bool(Pass *, Operation *)> shouldPrintBeforePass,
     std::function<bool(Pass *, Operation *)> shouldPrintAfterPass,
-    bool printModuleScope, bool printAfterOnlyOnChange, raw_ostream &out) {
+    bool printModuleScope, bool printAfterOnlyOnChange, raw_ostream &out,
+    OpPrintingFlags opPrintingFlags) {
   enableIRPrinting(std::make_unique<BasicIRPrinterConfig>(
       std::move(shouldPrintBeforePass), std::move(shouldPrintAfterPass),
-      printModuleScope, printAfterOnlyOnChange, out));
+      printModuleScope, printAfterOnlyOnChange, opPrintingFlags, out));
 }