Revert "[MLIR] Implement remark emitting policies in MLIR" (#160681)

Reverts #160526

This fails with Sanitizers.

Author: joker-eph

mlir/include/mlir/IR/Remarks.h

@@ -24,8 +24,6 @@
 #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
@@ -146,7 +144,7 @@ class Remark {
 
   llvm::StringRef getCategoryName() const { return categoryName; }
 
-  llvm::StringRef getCombinedCategoryName() const {
+  llvm::StringRef getFullCategoryName() const {
     if (categoryName.empty() && subCategoryName.empty())
       return {};
     if (subCategoryName.empty())
@@ -320,7 +318,7 @@ class InFlightRemark {
 };
 
 //===----------------------------------------------------------------------===//
-// Pluggable Remark Utilities
+// MLIR Remark Streamer
 //===----------------------------------------------------------------------===//
 
 /// Base class for MLIR remark streamers that is used to stream
@@ -340,26 +338,6 @@ 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)
 //===----------------------------------------------------------------------===//
@@ -377,8 +355,6 @@ 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;
 
@@ -416,8 +392,6 @@ 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.
@@ -433,10 +407,8 @@ class RemarkEngine {
   ~RemarkEngine();
 
   /// Setup the remark engine with the given output path and format.
-  LogicalResult
-  initialize(std::unique_ptr<MLIRRemarkStreamerBase> streamer,
-             std::unique_ptr<RemarkEmittingPolicyBase> remarkEmittingPolicy,
-             std::string *errMsg);
+  LogicalResult initialize(std::unique_ptr<MLIRRemarkStreamerBase> streamer,
+                           std::string *errMsg);
 
   /// Report a remark.
   void report(const Remark &&remark);
@@ -474,54 +446,6 @@ 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) {
@@ -581,72 +505,16 @@ 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,7 +45,6 @@ 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,11 +44,6 @@ 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.
@@ -247,8 +242,6 @@ 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.
@@ -272,8 +265,6 @@ 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,8 +278,6 @@ 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 = getCombinedCategoryName();
+  StringRef categoryName = getFullCategoryName();
   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 = getCombinedCategoryName();
+  r.PassName = getFullCategoryName();
   r.FunctionName = getFunction();
   r.Loc = locLambda();
   for (const Remark::Arg &arg : getArgs()) {
@@ -225,7 +225,7 @@ InFlightRemark RemarkEngine::emitOptimizationRemarkAnalysis(Location loc,
 // RemarkEngine
 //===----------------------------------------------------------------------===//
 
-void RemarkEngine::reportImpl(const Remark &remark) {
+void RemarkEngine::report(const Remark &&remark) {
   // Stream the remark
   if (remarkStreamer)
     remarkStreamer->streamOptimizationRemark(remark);
@@ -235,19 +235,19 @@ void RemarkEngine::reportImpl(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,44 +300,19 @@ 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<detail::MLIRRemarkStreamerBase> streamer,
-    std::unique_ptr<detail::RemarkEmittingPolicyBase> remarkEmittingPolicy,
-    const RemarkCategories &cats, bool printAsEmitRemarks) {
+    MLIRContext &ctx,
+    std::unique_ptr<remark::detail::MLIRRemarkStreamerBase> streamer,
+    const remark::RemarkCategories &cats, bool printAsEmitRemarks) {
   auto engine =
-      std::make_unique<detail::RemarkEngine>(printAsEmitRemarks, cats);
+      std::make_unique<remark::detail::RemarkEngine>(printAsEmitRemarks, cats);
 
   std::string errMsg;
-  if (failed(engine->initialize(std::move(streamer),
-                                std::move(remarkEmittingPolicy), &errMsg))) {
+  if (failed(engine->initialize(std::move(streamer), &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,16 +60,14 @@ 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),
-                                           std::move(remarkEmittingPolicy), cat,
+  return remark::enableOptimizationRemarks(ctx, std::move(*sOr), cat,
                                            printAsEmitRemarks);
 }