[MLIR] Introduce a OpWithState class to act as a stream modifier for Operations (NFC)

[joker-eph]

On the model of OpWithFlags, this modifier allows to stream an operation using a custom AsmPrinter.

[llvmbot]

@llvm/pr-subscribers-mlir-core

Author: Mehdi Amini (joker-eph)

Changes

On the model of OpWithFlags, this modifier allows to stream an operation using a custom AsmPrinter.

---

Full diff: https://github.com/llvm/llvm-project/pull/151547.diff

2 Files Affected:

• (modified) mlir/include/mlir/IR/Operation.h (+21)
• (modified) mlir/lib/Tools/mlir-opt/MlirOptMain.cpp (+1-2)

diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index edc8ab489e12d..f4bbc8e6c984a 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -1125,6 +1125,27 @@ inline raw_ostream &operator<<(raw_ostream &os,
   return os;
 }
 
+/// A wrapper class that allows for printing an operation with a custom
+/// AsmState, useful to act as a "stream modifier" to customize printing an
+/// operation with a stream using the operator<< overload, e.g.:
+///   llvm::dbgs() << OpWithState(op, OpPrintingFlags().skipRegions());
+class OpWithState {
+public:
+  OpWithState(Operation *op, AsmState &state) : op(op), theState(state) {}
+  AsmState &state() { return theState; }
+  const AsmState &state() const { return theState; }
+
+private:
+  Operation *op;
+  AsmState &theState;
+  friend raw_ostream &operator<<(raw_ostream &os, OpWithState &op);
+};
+
+inline raw_ostream &operator<<(raw_ostream &os, OpWithState &opWithState) {
+  opWithState.op->print(os, opWithState.state());
+  return os;
+}
+
 } // namespace mlir
 
 namespace llvm {
diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
index bdcdaa407e616..de714d8b740af 100644
--- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
+++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
@@ -501,8 +501,7 @@ performActions(raw_ostream &os,
            << "bytecode version while not emitting bytecode";
   AsmState asmState(op.get(), OpPrintingFlags(), /*locationMap=*/nullptr,
                     &fallbackResourceMap);
-  op.get()->print(os, asmState);
-  os << '\n';
+  os << OpWithState(op.get(), asmState) << '\n';
   return success();
 }
 

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

Changes

On the model of OpWithFlags, this modifier allows to stream an operation using a custom AsmPrinter.

---

Full diff: https://github.com/llvm/llvm-project/pull/151547.diff

2 Files Affected:

• (modified) mlir/include/mlir/IR/Operation.h (+21)
• (modified) mlir/lib/Tools/mlir-opt/MlirOptMain.cpp (+1-2)

diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index edc8ab489e12d..f4bbc8e6c984a 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -1125,6 +1125,27 @@ inline raw_ostream &operator<<(raw_ostream &os,
   return os;
 }
 
+/// A wrapper class that allows for printing an operation with a custom
+/// AsmState, useful to act as a "stream modifier" to customize printing an
+/// operation with a stream using the operator<< overload, e.g.:
+///   llvm::dbgs() << OpWithState(op, OpPrintingFlags().skipRegions());
+class OpWithState {
+public:
+  OpWithState(Operation *op, AsmState &state) : op(op), theState(state) {}
+  AsmState &state() { return theState; }
+  const AsmState &state() const { return theState; }
+
+private:
+  Operation *op;
+  AsmState &theState;
+  friend raw_ostream &operator<<(raw_ostream &os, OpWithState &op);
+};
+
+inline raw_ostream &operator<<(raw_ostream &os, OpWithState &opWithState) {
+  opWithState.op->print(os, opWithState.state());
+  return os;
+}
+
 } // namespace mlir
 
 namespace llvm {
diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
index bdcdaa407e616..de714d8b740af 100644
--- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
+++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
@@ -501,8 +501,7 @@ performActions(raw_ostream &os,
            << "bytecode version while not emitting bytecode";
   AsmState asmState(op.get(), OpPrintingFlags(), /*locationMap=*/nullptr,
                     &fallbackResourceMap);
-  op.get()->print(os, asmState);
-  os << '\n';
+  os << OpWithState(op.get(), asmState) << '\n';
   return success();
 }
 

[Copilot]

Pull Request Overview

This PR introduces a new OpWithState class that serves as a stream modifier for printing MLIR operations with custom AsmState configurations. The class follows the pattern established by OpWithFlags to enable more flexible operation printing.

Key changes:

• Introduces OpWithState wrapper class in Operation.h for custom AsmState printing
• Updates MlirOptMain.cpp to use the new streaming interface instead of direct print calls

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
mlir/include/mlir/IR/Operation.h	Adds OpWithState class definition with streaming operator overload
mlir/lib/Tools/mlir-opt/MlirOptMain.cpp	Replaces direct print call with OpWithState streaming usage

Comments suppressed due to low confidence (1)

mlir/include/mlir/IR/Operation.h:1139

• The parameter name 'op' in the friend declaration is misleading since it refers to an OpWithState object, not an Operation. Consider renaming to 'opWithState' for clarity.

  friend raw_ostream &operator<<(raw_ostream &os, const OpWithState &op);