[MLIR] Implement remark emitting policies in MLIR (#160526)

This update introduces two new remark emitting policies: 1. `RemarkEmittingPolicyAll`, which emits all remarks, 2. `RemarkEmittingPolicyFinal`, which only emits final remarks after processing. The `RemarkEngine` is modified to support these policies, allowing for more flexible remark handling based on user configuration. PR also adds flag to `mlir-opt` ``` --remark-policy=<value> - Specify the policy for remark output. =all - Print all remarks =final - Print final remarks ```

Author: grypp

mlir/include/mlir/IR/Remarks.h

@@ -24,6 +24,8 @@
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Value.h"
 
+#include <functional>
+
 namespace mlir::remark {
 
 /// Define an the set of categories to accept. By default none are, the provided
@@ -144,7 +146,7 @@ class Remark {
 
  llvm::StringRef getCategoryName() const { return categoryName; }
 
- llvm::StringRef getFullCategoryName() const {
+ llvm::StringRef getCombinedCategoryName() const {
  if (categoryName.empty() && subCategoryName.empty())
  return {};
  if (subCategoryName.empty())
@@ -318,7 +320,7 @@ class InFlightRemark {
 };
 
 //===----------------------------------------------------------------------===//
-// MLIR Remark Streamer
+// Pluggable Remark Utilities
 //===----------------------------------------------------------------------===//
 
 /// Base class for MLIR remark streamers that is used to stream
@@ -338,6 +340,26 @@ class MLIRRemarkStreamerBase {
  virtual void finalize() {} // optional
 };
 
+using ReportFn = llvm::unique_function<void(const Remark &)>;
+
+/// Base class for MLIR remark emitting policies that is used to emit
+/// optimization remarks to the underlying remark streamer. The derived classes
+/// should implement the `reportRemark` method to provide the actual emitting
+/// implementation.
+class RemarkEmittingPolicyBase {
+protected:
+ ReportFn reportImpl;
+
+public:
+ RemarkEmittingPolicyBase() = default;
+ virtual ~RemarkEmittingPolicyBase() = default;
+
+ void initialize(ReportFn fn) { reportImpl = std::move(fn); }
+
+ virtual void reportRemark(const Remark &remark) = 0;
+ virtual void finalize() = 0;
+};
+
 //===----------------------------------------------------------------------===//
 // Remark Engine (MLIR Context will own this class)
 //===----------------------------------------------------------------------===//
@@ -355,6 +377,8 @@ class RemarkEngine {
  std::optional<llvm::Regex> failedFilter;
  /// The MLIR remark streamer that will be used to emit the remarks.
  std::unique_ptr<MLIRRemarkStreamerBase> remarkStreamer;
+ /// The MLIR remark policy that will be used to emit the remarks.
+ std::unique_ptr<RemarkEmittingPolicyBase> remarkEmittingPolicy;
  /// When is enabled, engine also prints remarks as mlir::emitRemarks.
  bool printAsEmitRemarks = false;
 
@@ -392,6 +416,8 @@ class RemarkEngine {
  InFlightRemark emitIfEnabled(Location loc, RemarkOpts opts,
  bool (RemarkEngine::*isEnabled)(StringRef)
  const);
+ /// Report a remark.
+ void reportImpl(const Remark &remark);
 
 public:
  /// Default constructor is deleted, use the other constructor.
@@ -407,8 +433,10 @@ class RemarkEngine {
  ~RemarkEngine();
 
  /// Setup the remark engine with the given output path and format.
- LogicalResult initialize(std::unique_ptr<MLIRRemarkStreamerBase> streamer,
- std::string *errMsg);
+ LogicalResult
+ initialize(std::unique_ptr<MLIRRemarkStreamerBase> streamer,
+ std::unique_ptr<RemarkEmittingPolicyBase> remarkEmittingPolicy,
+ std::string *errMsg);
 
  /// Report a remark.
  void report(const Remark &&remark);
@@ -446,6 +474,54 @@ inline InFlightRemark withEngine(Fn fn, Location loc, Args &&...args) {
 
 namespace mlir::remark {
 
+//===----------------------------------------------------------------------===//
+// Remark Emitting Policies
+//===----------------------------------------------------------------------===//
+
+/// Policy that emits all remarks.
+class RemarkEmittingPolicyAll : public detail::RemarkEmittingPolicyBase {
+public:
+ RemarkEmittingPolicyAll();
+
+ void reportRemark(const detail::Remark &remark) override {
+ reportImpl(remark);
+ }
+ void finalize() override {}
+};
+
+/// Policy that emits final remarks.
+class RemarkEmittingPolicyFinal : public detail::RemarkEmittingPolicyBase {
+private:
+ /// user can intercept them for custom processing via a registered callback,
+ /// otherwise they will be reported on engine destruction.
+ llvm::DenseSet<detail::Remark> postponedRemarks;
+ /// Optional user callback for intercepting postponed remarks.
+ std::function<void(const detail::Remark &)> postponedRemarkCallback;
+
+public:
+ RemarkEmittingPolicyFinal();
+
+ /// Register a callback to intercept postponed remarks before they are
+ /// reported. The callback will be invoked for each postponed remark in
+ /// finalize().
+ void
+ setPostponedRemarkCallback(std::function<void(const detail::Remark &)> cb) {
+ postponedRemarkCallback = std::move(cb);
+ }
+
+ void reportRemark(const detail::Remark &remark) override {
+ postponedRemarks.erase(remark);
+ postponedRemarks.insert(remark);
+ }
+ void finalize() override {
+ for (auto &remark : postponedRemarks) {
+ if (postponedRemarkCallback)
+ postponedRemarkCallback(remark);
+ reportImpl(remark);
+ }
+ }
+};
+
 /// Create a Reason with llvm::formatv formatting.
 template <class... Ts>
 inline detail::LazyTextBuild reason(const char *fmt, Ts &&...ts) {
@@ -505,16 +581,72 @@ inline detail::InFlightRemark analysis(Location loc, RemarkOpts opts) {
 
 /// Setup remarks for the context. This function will enable the remark engine
 /// and set the streamer to be used for optimization remarks. The remark
-/// categories are used to filter the remarks that will be emitted by the remark
-/// engine. If a category is not specified, it will not be emitted. If
+/// categories are used to filter the remarks that will be emitted by the
+/// remark engine. If a category is not specified, it will not be emitted. If
 /// `printAsEmitRemarks` is true, the remarks will be printed as
 /// mlir::emitRemarks. 'streamer' must inherit from MLIRRemarkStreamerBase and
 /// will be used to stream the remarks.
 LogicalResult enableOptimizationRemarks(
  MLIRContext &ctx,
  std::unique_ptr<remark::detail::MLIRRemarkStreamerBase> streamer,
+ std::unique_ptr<remark::detail::RemarkEmittingPolicyBase>
+ remarkEmittingPolicy,
  const remark::RemarkCategories &cats, bool printAsEmitRemarks = false);
 
 } // namespace mlir::remark
 
+// DenseMapInfo specialization for Remark
+namespace llvm {
+template <>
+struct DenseMapInfo<mlir::remark::detail::Remark> {
+ static constexpr StringRef kEmptyKey = "<EMPTY_KEY>";
+ static constexpr StringRef kTombstoneKey = "<TOMBSTONE_KEY>";
+
+ /// Helper to provide a static dummy context for sentinel keys.
+ static mlir::MLIRContext *getStaticDummyContext() {
+ static mlir::MLIRContext dummyContext;
+ return &dummyContext;
+ }
+
+ /// Create an empty remark
+ static inline mlir::remark::detail::Remark getEmptyKey() {
+ return mlir::remark::detail::Remark(
+ mlir::remark::RemarkKind::RemarkUnknown, mlir::DiagnosticSeverity::Note,
+ mlir::UnknownLoc::get(getStaticDummyContext()),
+ mlir::remark::RemarkOpts::name(kEmptyKey));
+ }
+
+ /// Create a dead remark
+ static inline mlir::remark::detail::Remark getTombstoneKey() {
+ return mlir::remark::detail::Remark(
+ mlir::remark::RemarkKind::RemarkUnknown, mlir::DiagnosticSeverity::Note,
+ mlir::UnknownLoc::get(getStaticDummyContext()),
+ mlir::remark::RemarkOpts::name(kTombstoneKey));
+ }
+
+ /// Compute the hash value of the remark
+ static unsigned getHashValue(const mlir::remark::detail::Remark &remark) {
+ return llvm::hash_combine(
+ remark.getLocation().getAsOpaquePointer(),
+ llvm::hash_value(remark.getRemarkName()),
+ llvm::hash_value(remark.getCombinedCategoryName()));
+ }
+
+ static bool isEqual(const mlir::remark::detail::Remark &lhs,
+ const mlir::remark::detail::Remark &rhs) {
+ // Check for empty/tombstone keys first
+ if (lhs.getRemarkName() == kEmptyKey ||
+ lhs.getRemarkName() == kTombstoneKey ||
+ rhs.getRemarkName() == kEmptyKey ||
+ rhs.getRemarkName() == kTombstoneKey) {
+ return lhs.getRemarkName() == rhs.getRemarkName();
+ }
+
+ // For regular remarks, compare key identifying fields
+ return lhs.getLocation() == rhs.getLocation() &&
+ lhs.getRemarkName() == rhs.getRemarkName() &&
+ lhs.getCombinedCategoryName() == rhs.getCombinedCategoryName();
+ }
+};
+} // namespace llvm
 #endif // MLIR_IR_REMARKS_H

mlir/include/mlir/Remark/RemarkStreamer.h

@@ -45,6 +45,7 @@ namespace mlir::remark {
 /// mlir::emitRemarks.
 LogicalResult enableOptimizationRemarksWithLLVMStreamer(
  MLIRContext &ctx, StringRef filePath, llvm::remarks::Format fmt,
+ std::unique_ptr<detail::RemarkEmittingPolicyBase> remarkEmittingPolicy,
  const RemarkCategories &cat, bool printAsEmitRemarks = false);
 
 } // namespace mlir::remark

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

@@ -44,6 +44,11 @@ enum class RemarkFormat {
  REMARK_FORMAT_BITSTREAM,
 };
 
+enum class RemarkPolicy {
+ REMARK_POLICY_ALL,
+ REMARK_POLICY_FINAL,
+};
+
 /// Configuration options for the mlir-opt tool.
 /// This is intended to help building tools like mlir-opt by collecting the
 /// supported options.
@@ -242,6 +247,8 @@ class MlirOptMainConfig {
 
  /// Set the reproducer output filename
  RemarkFormat getRemarkFormat() const { return remarkFormatFlag; }
+ /// Set the remark policy to use.
+ RemarkPolicy getRemarkPolicy() const { return remarkPolicyFlag; }
  /// Set the remark format to use.
  std::string getRemarksAllFilter() const { return remarksAllFilterFlag; }
  /// Set the remark output file.
@@ -265,6 +272,8 @@ class MlirOptMainConfig {
 
  /// Remark format
  RemarkFormat remarkFormatFlag = RemarkFormat::REMARK_FORMAT_STDOUT;
+ /// Remark policy
+ RemarkPolicy remarkPolicyFlag = RemarkPolicy::REMARK_POLICY_ALL;
  /// Remark file to output to
  std::string remarksOutputFileFlag = "";
  /// Remark filters

mlir/lib/IR/MLIRContext.cpp

@@ -278,6 +278,8 @@ class MLIRContextImpl {
  }
  }
  ~MLIRContextImpl() {
+ // finalize remark engine before destroying anything else.
+ remarkEngine.reset();
  for (auto typeMapping : registeredTypes)
  typeMapping.second->~AbstractType();
  for (auto attrMapping : registeredAttributes)

mlir/lib/IR/Remarks.cpp

@@ -16,7 +16,7 @@
 #include "llvm/ADT/StringRef.h"
 
 using namespace mlir::remark::detail;
-
+using namespace mlir::remark;
 //------------------------------------------------------------------------------
 // Remark
 //------------------------------------------------------------------------------
@@ -70,7 +70,7 @@ static void printArgs(llvm::raw_ostream &os, llvm::ArrayRef<Remark::Arg> args) {
 void Remark::print(llvm::raw_ostream &os, bool printLocation) const {
  // Header: [Type] pass:remarkName
  StringRef type = getRemarkTypeString();
- StringRef categoryName = getFullCategoryName();
+ StringRef categoryName = getCombinedCategoryName();
  StringRef name = remarkName;
 
  os << '[' << type << "] ";
@@ -140,7 +140,7 @@ llvm::remarks::Remark Remark::generateRemark() const {
  r.RemarkType = getRemarkType();
  r.RemarkName = getRemarkName();
  // MLIR does not use passes; instead, it has categories and sub-categories.
- r.PassName = getFullCategoryName();
+ r.PassName = getCombinedCategoryName();
  r.FunctionName = getFunction();
  r.Loc = locLambda();
  for (const Remark::Arg &arg : getArgs()) {
@@ -225,7 +225,7 @@ InFlightRemark RemarkEngine::emitOptimizationRemarkAnalysis(Location loc,
 // RemarkEngine
 //===----------------------------------------------------------------------===//
 
-void RemarkEngine::report(const Remark &&remark) {
+void RemarkEngine::reportImpl(const Remark &remark) {
  // Stream the remark
  if (remarkStreamer)
  remarkStreamer->streamOptimizationRemark(remark);
@@ -235,19 +235,19 @@ void RemarkEngine::report(const Remark &&remark) {
  emitRemark(remark.getLocation(), remark.getMsg());
 }
 
+void RemarkEngine::report(const Remark &&remark) {
+ if (remarkEmittingPolicy)
+ remarkEmittingPolicy->reportRemark(remark);
+}
+
 RemarkEngine::~RemarkEngine() {
+ if (remarkEmittingPolicy)
+ remarkEmittingPolicy->finalize();
+
  if (remarkStreamer)
  remarkStreamer->finalize();
 }
 
-llvm::LogicalResult
-RemarkEngine::initialize(std::unique_ptr<MLIRRemarkStreamerBase> streamer,
- std::string *errMsg) {
- // If you need to validate categories/filters, do so here and set errMsg.
- remarkStreamer = std::move(streamer);
- return success();
-}
-
 /// Returns true if filter is already anchored like ^...$
 static bool isAnchored(llvm::StringRef s) {
  s = s.trim();
@@ -300,19 +300,44 @@ RemarkEngine::RemarkEngine(bool printAsEmitRemarks,
  failedFilter = buildFilter(cats, cats.failed);
 }
 
+llvm::LogicalResult RemarkEngine::initialize(
+ std::unique_ptr<MLIRRemarkStreamerBase> streamer,
+ std::unique_ptr<RemarkEmittingPolicyBase> remarkEmittingPolicy,
+ std::string *errMsg) {
+
+ remarkStreamer = std::move(streamer);
+
+ // Capture `this`. Ensure RemarkEngine is not moved after this.
+ auto reportFunc = [this](const Remark &r) { this->reportImpl(r); };
+ remarkEmittingPolicy->initialize(ReportFn(std::move(reportFunc)));
+
+ this->remarkEmittingPolicy = std::move(remarkEmittingPolicy);
+ return success();
+}
+
 llvm::LogicalResult mlir::remark::enableOptimizationRemarks(
- MLIRContext &ctx,
- std::unique_ptr<remark::detail::MLIRRemarkStreamerBase> streamer,
- const remark::RemarkCategories &cats, bool printAsEmitRemarks) {
+ MLIRContext &ctx, std::unique_ptr<detail::MLIRRemarkStreamerBase> streamer,
+ std::unique_ptr<detail::RemarkEmittingPolicyBase> remarkEmittingPolicy,
+ const RemarkCategories &cats, bool printAsEmitRemarks) {
  auto engine =
- std::make_unique<remark::detail::RemarkEngine>(printAsEmitRemarks, cats);
+ std::make_unique<detail::RemarkEngine>(printAsEmitRemarks, cats);
 
  std::string errMsg;
- if (failed(engine->initialize(std::move(streamer), &errMsg))) {
+ if (failed(engine->initialize(std::move(streamer),
+ std::move(remarkEmittingPolicy), &errMsg))) {
  llvm::report_fatal_error(
  llvm::Twine("Failed to initialize remark engine. Error: ") + errMsg);
  }
  ctx.setRemarkEngine(std::move(engine));
 
  return success();
 }
+
+//===----------------------------------------------------------------------===//
+// Remark emitting policies
+//===----------------------------------------------------------------------===//
+
+namespace mlir::remark {
+RemarkEmittingPolicyAll::RemarkEmittingPolicyAll() = default;
+RemarkEmittingPolicyFinal::RemarkEmittingPolicyFinal() = default;
+} // namespace mlir::remark

mlir/lib/Remark/RemarkStreamer.cpp

@@ -60,14 +60,16 @@ void LLVMRemarkStreamer::finalize() {
 namespace mlir::remark {
 LogicalResult enableOptimizationRemarksWithLLVMStreamer(
  MLIRContext &ctx, StringRef path, llvm::remarks::Format fmt,
+ std::unique_ptr<detail::RemarkEmittingPolicyBase> remarkEmittingPolicy,
  const RemarkCategories &cat, bool printAsEmitRemarks) {
 
  FailureOr<std::unique_ptr<detail::MLIRRemarkStreamerBase>> sOr =
  detail::LLVMRemarkStreamer::createToFile(path, fmt);
  if (failed(sOr))
  return failure();
 
- return remark::enableOptimizationRemarks(ctx, std::move(*sOr), cat,
+ return remark::enableOptimizationRemarks(ctx, std::move(*sOr),
+ std::move(remarkEmittingPolicy), cat,
  printAsEmitRemarks);
 }