Add capture of "IRUnits" as context for an MLIR Action

[mehdi_amini]

Summary

IRUnit is defined as:

using IRUnit = PointerUnion<Operation *, Region *, Block *, Value>;

The tracing::Action is extended to take an ArrayRef<IRUnit> as context to
describe an Action. It is demonstrated in the "ActionLogging" observer.

[mehdi_amini]

rebase

[Tyker]

The data-flow analysis framework has a concept very similar to IRUnit its called ProgramPoint maybe this should be shared ?

[Mogball]

mlir/include/mlir/IR/Unit.h

18

25

32

43

newline

mlir/lib/IR/AsmPrinter.cpp

681 ↗

(On Diff #500571)

where does the 2 come from?

[rriddle]

mlir/include/mlir/IR/Unit.h

26

It'd be nice not to document this in the context of actions, I'd rather have something be general.

mlir/lib/Debug/Observers/ActionLogging.cpp

10–13

Some of these could be dropped? Like Region.h?

47–48

Can you add a stream operator and simply this?

mlir/lib/IR/AsmPrinter.cpp

228 ↗

(On Diff #500571)

It'd be nice to split this change out, it feels separable. I think we could also use this here: https://github.com/llvm/llvm-project/blob/73058e330134d275070d49dba80e5fda50ec015b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp#L519

[mehdi_amini]

In D144814#4163506, @Tyker wrote:

The data-flow analysis framework has a concept very similar to IRUnit its called ProgramPoint maybe this should be shared ?

I'm not sure how to unify it, it feels that somehow the ProgramPoint could inherit from IRUnit but it is declared as a more extensible framework: public PointerUnion<GenericProgramPoint *, Operation *, Value, Block *> so I'm not sure how to unify.

[mehdi_amini]

address comments

mlir/include/mlir/IR/Unit.h

26

Yeah that's a remainder from a time where this was defined inside the Action.h header.

mlir/lib/IR/AsmPrinter.cpp

228 ↗

(On Diff #500571)

What about updating MLIRServer to use this new flag?

[rriddle]

mlir/include/mlir/IR/OperationSupport.h

867 ↗

(On Diff #504478)

Missing a rebase?

mlir/include/mlir/IR/Unit.h

35

Should be able to drop most of the llvm:: in this file.

mlir/lib/IR/Unit.cpp

42–46

Can you just do std::distance instead (using block->getIterator())?

55–58

Feels cleaner, but feel free to ignore

[mehdi_amini]

mlir/include/mlir/IR/OperationSupport.h

867 ↗

(On Diff #504478)

Yeah haven't updated this patch since I wrote the other one :)

[mehdi_amini]

Rebase & address comments

[mehdi_amini]

One more comment addressed

[mehdi_amini]

In D144814#4187700, @mehdi_amini wrote:

In D144814#4163506, @Tyker wrote:

The data-flow analysis framework has a concept very similar to IRUnit its called ProgramPoint maybe this should be shared ?

I'm not sure how to unify it, it feels that somehow the ProgramPoint could inherit from IRUnit but it is declared as a more extensible framework: public PointerUnion<GenericProgramPoint *, Operation *, Value, Block *> so I'm not sure how to unify.

If anyone gets an idea, we'll refactor then, in the meantime I'll push this as-is!

mlir/include/mlir/Debug/Observers/ActionLogging.h

 
 namespace mlir {
 namespace tracing {
 
 /// This class defines an observer that print Actions before and after execution
 /// on the provided stream.
 struct ActionLogger : public ExecutionContext::Observer {
   ActionLogger(raw_ostream &os, bool printActions = true,
-               bool printBreakpoints = true)
-      : os(os), printActions(printActions), printBreakpoints(printBreakpoints) {
-  }
+               bool printBreakpoints = true, bool printIRUnits = true)
+      : os(os), printActions(printActions), printBreakpoints(printBreakpoints),
+        printIRUnits(printIRUnits) {}
 
   void beforeExecute(const ActionActiveStack *action, Breakpoint *breakpoint,
                      bool willExecute) override;
   void afterExecute(const ActionActiveStack *action) override;
 
 private:
   raw_ostream &os;
   bool printActions;
   bool printBreakpoints;
+  bool printIRUnits;
 };
 
 } // namespace tracing
 } // namespace mlir
 
 #endif // MLIR_TRACING_OBSERVERS_ACTIONLOGGING_H

mlir/include/mlir/IR/Action.h

 // allows for external entities to control certain actions taken by the compiler
 // by registering handler functions.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_ACTION_H
 #define MLIR_IR_ACTION_H
 
+#include "mlir/IR/Unit.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/TypeID.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/TypeName.h"
 #include "llvm/Support/raw_ostream.h"
 #include <functional>
public:
   /// Return a string "tag" which intends to uniquely identify this type of
   /// action. For example "pass-application" or "pattern-rewrite".
   virtual StringRef getTag() const = 0;
 
   virtual void print(raw_ostream &os) const {
     os << "Action \"" << getTag() << "\"";
   }
 
+  /// Return the set of IR units that are associated with this action.
+  virtual ArrayRef<IRUnit> getContextIRUnits() const { return irUnits; }
+
 protected:
-  Action(TypeID actionID) : actionID(actionID) {}
+  Action(TypeID actionID, ArrayRef<IRUnit> irUnits)
+      : actionID(actionID), irUnits(irUnits) {}
 
   /// The type of the derived action class, used for `isa`/`dyn_cast`.
   TypeID actionID;
+
+  /// Set of IR units (operations, regions, blocks, values) that are associated
+  /// with this action.
+  ArrayRef<IRUnit> irUnits;
 };
 
 /// CRTP Implementation of an action. This class provides a base class for
 /// implementing specific actions.
 ///  Derived classes are expected to provide the following:
 ///   * static constexpr StringLiteral tag = "...";
 ///     - This method returns a unique string identifier, similar to a command
 ///       line flag or DEBUG_TYPE.
 template <typename Derived>
 class ActionImpl : public Action {
 public:
-  ActionImpl() : Action(TypeID::get<Derived>()) {}
+  ActionImpl(ArrayRef<IRUnit> irUnits = {})
+      : Action(TypeID::get<Derived>(), irUnits) {}
 
   /// Provide classof to allow casting between action types.
   static bool classof(const Action *action) {
     return action->getActionID() == TypeID::get<Derived>();
   }
 
   /// Forward tag access to the derived class.
   StringRef getTag() const final { return Derived::tag; }
 };
 
 } // namespace tracing
 } // namespace mlir
 
 #endif // MLIR_IR_ACTION_H

mlir/include/mlir/IR/MLIRContext.h

 //===- MLIRContext.h - MLIR Global Context Class ----------------*- 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_IR_MLIRCONTEXT_H
 #define MLIR_IR_MLIRCONTEXT_H
 
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/TypeID.h"
+#include "llvm/ADT/ArrayRef.h"
 #include <functional>
 #include <memory>
 #include <vector>
 
 namespace llvm {
 class ThreadPool;
 } // namespace llvm
 
void executeAction(function_ref<void()> actionFn,
     if (LLVM_UNLIKELY(hasActionHandler()))
       executeActionInternal(actionFn, action);
     else
       actionFn();
   }
 
   /// Dispatch the provided action to the handler if any, or just execute it.
   template <typename ActionTy, typename... Args>
-  void executeAction(function_ref<void()> actionFn, Args &&...args) {
+  void executeAction(function_ref<void()> actionFn, ArrayRef<IRUnit> irUnits,
+                     Args &&...args) {
     if (LLVM_UNLIKELY(hasActionHandler()))
-      executeActionInternal<ActionTy, Args...>(actionFn,
+      executeActionInternal<ActionTy, Args...>(actionFn, irUnits,
                                                std::forward<Args>(args)...);
     else
       actionFn();
   }
 
 private:
   /// Return true if the given dialect is currently loading.
   bool isDialectLoading(StringRef dialectNamespace);
 
   /// Internal helper for the dispatch method.
   void executeActionInternal(function_ref<void()> actionFn,
                              const tracing::Action &action);
 
   /// Internal helper for the dispatch method. We get here after checking that
   /// there is a handler, for the purpose of keeping this code out-of-line. and
   /// avoid calling the ctor for the Action unnecessarily.
   template <typename ActionTy, typename... Args>
   LLVM_ATTRIBUTE_NOINLINE void
-  executeActionInternal(function_ref<void()> actionFn, Args &&...args) {
-    executeActionInternal(actionFn, ActionTy(std::forward<Args>(args)...));
+  executeActionInternal(function_ref<void()> actionFn, ArrayRef<IRUnit> irUnits,
+                        Args &&...args) {
+    executeActionInternal(actionFn,
+                          ActionTy(irUnits, std::forward<Args>(args)...));
   }
 
   const std::unique_ptr<MLIRContextImpl> impl;
 
   MLIRContext(const MLIRContext &) = delete;
   void operator=(const MLIRContext &) = delete;
 };
 

mlir/include/mlir/IR/Unit.h

+//===- Unit.h -  IR Unit definition--------------------*- 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_IR_UNIT_H
+#define MLIR_IR_UNIT_H
+
+#include "mlir/IR/OperationSupport.h"
+#include "llvm/ADT/PointerUnion.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace llvm {
+class raw_ostream;
+} // namespace llvm
+namespace mlir {
+class Operation;
+class Region;
+class Block;
+class Value;
+
+/// IRUnit is a union of the different types of IR objects that consistute the
+/// IR structure (other than Type and Attribute), that is Operation, Region, and
+/// Block.
+class IRUnit : public PointerUnion<Operation *, Region *, Block *, Value> {
+public:
+  using PointerUnion::PointerUnion;
+
+  /// Print the IRUnit to the given stream.
+  void print(raw_ostream &os,
+             OpPrintingFlags flags =
+                 OpPrintingFlags().skipRegions().useLocalScope()) const;
+};
+
+raw_ostream &operator<<(raw_ostream &os, const IRUnit &unit);
+
+} // end namespace mlir
+
+#endif // MLIR_IR_UNIT_H

mlir/lib/Debug/Observers/ActionLogging.cpp

 //===- ActionLogging.cpp -  Logging Actions *- 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Debug/Observers/ActionLogging.h"
+#include "mlir/IR/Action.h"
 #include "llvm/Support/Threading.h"
-#include <sstream>
-#include <thread>
+#include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::tracing;
 
 //===----------------------------------------------------------------------===//
 // ActionLogger
 //===----------------------------------------------------------------------===//
 
 void ActionLogger::beforeExecute(const ActionActiveStack *action,
                                  Breakpoint *breakpoint, bool willExecute) {
   SmallVector<char> name;
   llvm::get_thread_name(name);
+  if (name.empty()) {
+    llvm::raw_svector_ostream os(name);
+    os << llvm::get_threadid();
+  }
   os << "[thread " << name << "] ";
   if (willExecute)
     os << "begins ";
   else
     os << "skipping ";
   if (printBreakpoints) {
     if (breakpoint)
-      os << " (on breakpoint: " << *breakpoint << ") ";
+      os << "(on breakpoint: " << *breakpoint << ") ";
     else
-      os << " (no breakpoint) ";
+      os << "(no breakpoint) ";
   }
   os << "Action ";
   if (printActions)
     action->getAction().print(os);
   else
     os << action->getAction().getTag();
+  if (printIRUnits) {
+    os << " (";
+    interleaveComma(action->getAction().getContextIRUnits(), os);
+    os << ")";
+  }
   os << "`\n";
 }
 
 void ActionLogger::afterExecute(const ActionActiveStack *action) {
   SmallVector<char> name;
   llvm::get_thread_name(name);
+  if (name.empty()) {
+    llvm::raw_svector_ostream os(name);
+    os << llvm::get_threadid();
+  }
   os << "[thread " << name << "] completed `" << action->getAction().getTag()
      << "`\n";
 }

mlir/lib/IR/CMakeLists.txt

add_mlir_library(MLIRIR
   PatternMatch.cpp
   Region.cpp
   RegionKindInterface.cpp
   SymbolTable.cpp
   TensorEncoding.cpp
   Types.cpp
   TypeRange.cpp
   TypeUtilities.cpp
+  Unit.cpp
   Value.cpp
   ValueRange.cpp
   Verifier.cpp
   Visitors.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
 

mlir/lib/IR/Unit.cpp

+//===- Unit.cpp - Support for manipulating IR Unit ------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/Unit.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/Region.h"
+#include "llvm/Support/raw_ostream.h"
+#include <iterator>
+#include <sstream>
+
+using namespace mlir;
+
+static void printOp(llvm::raw_ostream &os, Operation *op,
+                    OpPrintingFlags &flags) {
+  if (!op) {
+    os << "<Operation:nullptr>";
+    return;
+  }
+  op->print(os, flags);
+}
+
+static void printRegion(llvm::raw_ostream &os, Region *region,
+                        OpPrintingFlags &flags) {
+  if (!region) {
+    os << "<Region:nullptr>";
+    return;
+  }
+  os << "Region #" << region->getRegionNumber() << " for op ";
+  printOp(os, region->getParentOp(), flags);
+}
+
+static void printBlock(llvm::raw_ostream &os, Block *block,
+                       OpPrintingFlags &flags) {
+  Region *region = block->getParent();
+  Block *entry = &region->front();
+  int blockId = std::distance(entry->getIterator(), block->getIterator());
+  os << "Block #" << blockId << " for ";
+  bool shouldSkipRegions = flags.shouldSkipRegions();
+  printRegion(os, region, flags.skipRegions());
+  if (!shouldSkipRegions)
+    block->print(os);
+}
+
+void mlir::IRUnit::print(llvm::raw_ostream &os, OpPrintingFlags flags) const {
+  if (auto *op = this->dyn_cast<Operation *>())
+    return printOp(os, op, flags);
+  if (auto *region = this->dyn_cast<Region *>())
+    return printRegion(os, region, flags);
+  if (auto *block = this->dyn_cast<Block *>())
+    return printBlock(os, block, flags);
+  llvm_unreachable("unknown IRUnit");
+}
+
+llvm::raw_ostream &mlir::operator<<(llvm::raw_ostream &os, const IRUnit &unit) {
+  unit.print(os);
+  return os;
+}

mlir/lib/Pass/Pass.cpp

op->getContext()->executeAction<PassExecutionAction>(
       [&]() {
         // Invoke the virtual runOnOperation method.
         if (auto *adaptor = dyn_cast<OpToOpPassAdaptor>(pass))
           adaptor->runOnOperation(verifyPasses);
         else
           pass->runOnOperation();
         passFailed = pass->passState->irAndPassFailed.getInt();
       },
-      *pass, op);
+      {op}, *pass);
 
   // Invalidate any non preserved analyses.
   am.invalidate(pass->passState->preservedAnalyses);
 
   // When verifyPasses is specified, we run the verifier (unless the pass
   // failed).
   if (!passFailed && verifyPasses) {
     bool runVerifierNow = true;

mlir/lib/Pass/PassDetail.h

 //===- PassDetail.h - MLIR Pass details -------------------------*- 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_PASSDETAIL_H_
 #define MLIR_PASS_PASSDETAIL_H_
 
 #include "mlir/IR/Action.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/FormatVariadic.h"
 
 namespace mlir {
 /// Encapsulate the "action" of executing a single pass, used for the MLIR
 /// tracing infrastructure.
 struct PassExecutionAction : public tracing::ActionImpl<PassExecutionAction> {
-  PassExecutionAction(const Pass &pass, Operation *op) : pass(pass), op(op) {}
+  using Base = tracing::ActionImpl<PassExecutionAction>;
+  PassExecutionAction(ArrayRef<IRUnit> irUnits, const Pass &pass)
+      : Base(irUnits), pass(pass) {}
   static constexpr StringLiteral tag = "pass-execution-action";
   void print(raw_ostream &os) const override;
   const Pass &getPass() const { return pass; }
-  Operation *getOp() const { return op; }
+  Operation *getOp() const {
+    ArrayRef<IRUnit> irUnits = getContextIRUnits();
+    return irUnits.empty() ? nullptr : irUnits[0].dyn_cast<Operation *>();
+  }
 
 public:
   const Pass &pass;
   Operation *op;
 };
 
 namespace detail {
 

mlir/test/Pass/action-logging.mlir

 // RUN: mlir-opt %s --log-actions-to=- -canonicalize -test-module-pass | FileCheck %s
 
-// CHECK: [thread {{.*}}] begins (no breakpoint) Action `pass-execution-action` running `Canonicalizer` on Operation `builtin.module`
-// CHECK: [thread {{.*}}] completed `pass-execution-action`
-// CHECK: [thread {{.*}}] begins (no breakpoint) Action `pass-execution-action` running `(anonymous namespace)::TestModulePass` on Operation `builtin.module`
-// CHECK: [thread {{.*}}] completed `pass-execution-action`
+// CHECK: [thread {{.*}}] begins (no breakpoint) Action `pass-execution-action` running `Canonicalizer` on Operation `builtin.module` (module {...}
+// CHECK-NEXT: [thread {{.*}}] completed `pass-execution-action`
+// CHECK-NEXT: [thread {{.*}}] begins (no breakpoint) Action `pass-execution-action`  running `{{.*}}TestModulePass` on Operation `builtin.module` (module {...}
+// CHECK-NEXT: [thread {{.*}}] completed `pass-execution-action`
+// CHECK-NOT: Action