MLIR  22.0.0git
Pass.cpp
Go to the documentation of this file.
1 //===- Pass.cpp - Pass Management -----------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "Pass.h"
10 
11 #include "IRModule.h"
12 #include "mlir-c/Pass.h"
13 // clang-format off
15 #include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
16 // clang-format on
17 
18 namespace nb = nanobind;
19 using namespace nb::literals;
20 using namespace mlir;
21 using namespace mlir::python;
22 
23 namespace {
24 
25 /// Owning Wrapper around a PassManager.
26 class PyPassManager {
27 public:
28  PyPassManager(MlirPassManager passManager) : passManager(passManager) {}
29  PyPassManager(PyPassManager &&other) noexcept
30  : passManager(other.passManager) {
31  other.passManager.ptr = nullptr;
32  }
33  ~PyPassManager() {
34  if (!mlirPassManagerIsNull(passManager))
35  mlirPassManagerDestroy(passManager);
36  }
37  MlirPassManager get() { return passManager; }
38 
39  void release() { passManager.ptr = nullptr; }
40  nb::object getCapsule() {
41  return nb::steal<nb::object>(mlirPythonPassManagerToCapsule(get()));
42  }
43 
44  static nb::object createFromCapsule(const nb::object &capsule) {
45  MlirPassManager rawPm = mlirPythonCapsuleToPassManager(capsule.ptr());
46  if (mlirPassManagerIsNull(rawPm))
47  throw nb::python_error();
48  return nb::cast(PyPassManager(rawPm), nb::rv_policy::move);
49  }
50 
51 private:
52  MlirPassManager passManager;
53 };
54 
55 } // namespace
56 
57 /// Create the `mlir.passmanager` here.
58 void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
59  //----------------------------------------------------------------------------
60  // Mapping of MlirExternalPass
61  //----------------------------------------------------------------------------
62  nb::class_<MlirExternalPass>(m, "ExternalPass")
63  .def("signal_pass_failure",
64  [](MlirExternalPass pass) { mlirExternalPassSignalFailure(pass); });
65 
66  //----------------------------------------------------------------------------
67  // Mapping of the top-level PassManager
68  //----------------------------------------------------------------------------
69  nb::class_<PyPassManager>(m, "PassManager")
70  .def(
71  "__init__",
72  [](PyPassManager &self, const std::string &anchorOp,
73  DefaultingPyMlirContext context) {
74  MlirPassManager passManager = mlirPassManagerCreateOnOperation(
75  context->get(),
76  mlirStringRefCreate(anchorOp.data(), anchorOp.size()));
77  new (&self) PyPassManager(passManager);
78  },
79  "anchor_op"_a = nb::str("any"), "context"_a = nb::none(),
80  // clang-format off
81  nb::sig("def __init__(self, anchor_op: str = 'any', context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " | None = None) -> None"),
82  // clang-format on
83  "Create a new PassManager for the current (or provided) Context.")
84  .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyPassManager::getCapsule)
85  .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyPassManager::createFromCapsule)
86  .def("_testing_release", &PyPassManager::release,
87  "Releases (leaks) the backing pass manager (testing)")
88  .def(
89  "enable_ir_printing",
90  [](PyPassManager &passManager, bool printBeforeAll,
91  bool printAfterAll, bool printModuleScope, bool printAfterChange,
92  bool printAfterFailure, std::optional<int64_t> largeElementsLimit,
93  std::optional<int64_t> largeResourceLimit, bool enableDebugInfo,
94  bool printGenericOpForm,
95  std::optional<std::string> optionalTreePrintingPath) {
96  MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
97  if (largeElementsLimit) {
99  *largeElementsLimit);
101  *largeElementsLimit);
102  }
103  if (largeResourceLimit)
105  *largeResourceLimit);
106  if (enableDebugInfo)
107  mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,
108  /*prettyForm=*/false);
109  if (printGenericOpForm)
111  std::string treePrintingPath = "";
112  if (optionalTreePrintingPath.has_value())
113  treePrintingPath = optionalTreePrintingPath.value();
115  passManager.get(), printBeforeAll, printAfterAll,
116  printModuleScope, printAfterChange, printAfterFailure, flags,
117  mlirStringRefCreate(treePrintingPath.data(),
118  treePrintingPath.size()));
120  },
121  "print_before_all"_a = false, "print_after_all"_a = true,
122  "print_module_scope"_a = false, "print_after_change"_a = false,
123  "print_after_failure"_a = false,
124  "large_elements_limit"_a = nb::none(),
125  "large_resource_limit"_a = nb::none(), "enable_debug_info"_a = false,
126  "print_generic_op_form"_a = false,
127  "tree_printing_dir_path"_a = nb::none(),
128  "Enable IR printing, default as mlir-print-ir-after-all.")
129  .def(
130  "enable_verifier",
131  [](PyPassManager &passManager, bool enable) {
132  mlirPassManagerEnableVerifier(passManager.get(), enable);
133  },
134  "enable"_a, "Enable / disable verify-each.")
135  .def(
136  "enable_timing",
137  [](PyPassManager &passManager) {
138  mlirPassManagerEnableTiming(passManager.get());
139  },
140  "Enable pass timing.")
141  .def_static(
142  "parse",
143  [](const std::string &pipeline, DefaultingPyMlirContext context) {
144  MlirPassManager passManager = mlirPassManagerCreate(context->get());
145  PyPrintAccumulator errorMsg;
148  mlirStringRefCreate(pipeline.data(), pipeline.size()),
149  errorMsg.getCallback(), errorMsg.getUserData());
150  if (mlirLogicalResultIsFailure(status))
151  throw nb::value_error(errorMsg.join().c_str());
152  return new PyPassManager(passManager);
153  },
154  "pipeline"_a, "context"_a = nb::none(),
155  // clang-format off
156  nb::sig("def parse(pipeline: str, context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " | None = None) -> PassManager"),
157  // clang-format on
158  "Parse a textual pass-pipeline and return a top-level PassManager "
159  "that can be applied on a Module. Throw a ValueError if the pipeline "
160  "can't be parsed")
161  .def(
162  "add",
163  [](PyPassManager &passManager, const std::string &pipeline) {
164  PyPrintAccumulator errorMsg;
166  mlirPassManagerGetAsOpPassManager(passManager.get()),
167  mlirStringRefCreate(pipeline.data(), pipeline.size()),
168  errorMsg.getCallback(), errorMsg.getUserData());
169  if (mlirLogicalResultIsFailure(status))
170  throw nb::value_error(errorMsg.join().c_str());
171  },
172  "pipeline"_a,
173  "Add textual pipeline elements to the pass manager. Throws a "
174  "ValueError if the pipeline can't be parsed.")
175  .def(
176  "add",
177  [](PyPassManager &passManager, const nb::callable &run,
178  std::optional<std::string> &name, const std::string &argument,
179  const std::string &description, const std::string &opName) {
180  if (!name.has_value()) {
181  name = nb::cast<std::string>(
182  nb::borrow<nb::str>(run.attr("__name__")));
183  }
184  MlirTypeIDAllocator typeIDAllocator = mlirTypeIDAllocatorCreate();
185  MlirTypeID passID =
186  mlirTypeIDAllocatorAllocateTypeID(typeIDAllocator);
187  MlirExternalPassCallbacks callbacks;
188  callbacks.construct = [](void *obj) {
189  (void)nb::handle(static_cast<PyObject *>(obj)).inc_ref();
190  };
191  callbacks.destruct = [](void *obj) {
192  (void)nb::handle(static_cast<PyObject *>(obj)).dec_ref();
193  };
194  callbacks.initialize = nullptr;
195  callbacks.clone = [](void *) -> void * {
196  throw std::runtime_error("Cloning Python passes not supported");
197  };
198  callbacks.run = [](MlirOperation op, MlirExternalPass pass,
199  void *userData) {
200  nb::handle(static_cast<PyObject *>(userData))(op, pass);
201  };
202  auto externalPass = mlirCreateExternalPass(
203  passID, mlirStringRefCreate(name->data(), name->length()),
204  mlirStringRefCreate(argument.data(), argument.length()),
205  mlirStringRefCreate(description.data(), description.length()),
206  mlirStringRefCreate(opName.data(), opName.size()),
207  /*nDependentDialects*/ 0, /*dependentDialects*/ nullptr,
208  callbacks, /*userData*/ run.ptr());
209  mlirPassManagerAddOwnedPass(passManager.get(), externalPass);
210  },
211  "run"_a, "name"_a.none() = nb::none(), "argument"_a.none() = "",
212  "description"_a.none() = "", "op_name"_a.none() = "",
213  "Add a python-defined pass to the pass manager.")
214  .def(
215  "run",
216  [](PyPassManager &passManager, PyOperationBase &op) {
217  // Actually run the pass manager.
220  passManager.get(), op.getOperation().get());
221  if (mlirLogicalResultIsFailure(status))
222  throw MLIRError("Failure while executing pass pipeline",
223  errors.take());
224  },
225  "operation"_a,
226  // clang-format off
227  nb::sig("def run(self, operation: " MAKE_MLIR_PYTHON_QUALNAME("ir._OperationBase") ") -> None"),
228  // clang-format on
229  "Run the pass manager on the provided operation, raising an "
230  "MLIRError on failure.")
231  .def(
232  "__str__",
233  [](PyPassManager &self) {
234  MlirPassManager passManager = self.get();
235  PyPrintAccumulator printAccum;
238  printAccum.getCallback(), printAccum.getUserData());
239  return printAccum.join();
240  },
241  "Print the textual representation for this PassManager, suitable to "
242  "be passed to `parse` for round-tripping.");
243 }
MlirPass mlirCreateExternalPass(MlirTypeID passID, MlirStringRef name, MlirStringRef argument, MlirStringRef description, MlirStringRef opName, intptr_t nDependentDialects, MlirDialectHandle *dependentDialects, MlirExternalPassCallbacks callbacks, void *userData)
Creates an external MlirPass that calls the supplied callbacks using the supplied userData.
Definition: Pass.cpp:204
MlirPassManager mlirPassManagerCreate(MlirContext ctx)
Create a new top-level PassManager with the default anchor.
Definition: Pass.cpp:24
void mlirPassManagerEnableVerifier(MlirPassManager passManager, bool enable)
Enable / disable verify-each.
Definition: Pass.cpp:74
void mlirPassManagerEnableTiming(MlirPassManager passManager)
Enable pass timing.
Definition: Pass.cpp:78
void mlirPassManagerDestroy(MlirPassManager passManager)
Destroy the provided PassManager.
Definition: Pass.cpp:33
MlirLogicalResult mlirParsePassPipeline(MlirOpPassManager passManager, MlirStringRef pipeline, MlirStringCallback callback, void *userData)
Parse a textual MLIR pass pipeline and assign it to the provided OpPassManager.
Definition: Pass.cpp:116
MlirOpPassManager mlirPassManagerGetAsOpPassManager(MlirPassManager passManager)
Cast a top-level PassManager to a generic OpPassManager.
Definition: Pass.cpp:38
MlirLogicalResult mlirPassManagerRunOnOp(MlirPassManager passManager, MlirOperation op)
Run the provided passManager on the given op.
Definition: Pass.cpp:42
void mlirExternalPassSignalFailure(MlirExternalPass pass)
This signals that the pass has failed.
Definition: Pass.cpp:219
void mlirPrintPassPipeline(MlirOpPassManager passManager, MlirStringCallback callback, void *userData)
Print a textual MLIR pass pipeline by sending chunks of the string representation and forwarding user...
Definition: Pass.cpp:110
void mlirPassManagerAddOwnedPass(MlirPassManager passManager, MlirPass pass)
Add a pass and transfer ownership to the provided top-level mlirPassManager.
Definition: Pass.cpp:92
MlirPassManager mlirPassManagerCreateOnOperation(MlirContext ctx, MlirStringRef anchorOp)
Create a new top-level PassManager anchored on anchorOp.
Definition: Pass.cpp:28
MlirLogicalResult mlirOpPassManagerAddPipeline(MlirOpPassManager passManager, MlirStringRef pipelineElements, MlirStringCallback callback, void *userData)
Parse a sequence of textual MLIR pass pipeline elements and add them to the provided OpPassManager.
Definition: Pass.cpp:101
void mlirPassManagerEnableIRPrinting(MlirPassManager passManager, bool printBeforeAll, bool printAfterAll, bool printModuleScope, bool printAfterOnlyOnChange, bool printAfterOnlyOnFailure, MlirOpPrintingFlags flags, MlirStringRef treePrintingPath)
Enable IR printing.
Definition: Pass.cpp:47
#define MLIR_PYTHON_CAPI_PTR_ATTR
Attribute on MLIR Python objects that expose their C-API pointer.
Definition: Interop.h:97
#define MLIR_PYTHON_CAPI_FACTORY_ATTR
Attribute on MLIR Python objects that exposes a factory function for constructing the corresponding P...
Definition: Interop.h:110
static PyObject * mlirPythonPassManagerToCapsule(MlirPassManager pm)
Creates a capsule object encapsulating the raw C-API MlirPassManager.
Definition: Interop.h:311
#define MAKE_MLIR_PYTHON_QUALNAME(local)
Definition: Interop.h:57
static MlirPassManager mlirPythonCapsuleToPassManager(PyObject *capsule)
Extracts an MlirPassManager from a capsule as produced from mlirPythonPassManagerToCapsule.
Definition: Interop.h:320
PyMlirContextRef & getContext()
Accesses the context reference.
Definition: IRModule.h:293
Used in function arguments when None should resolve to the current context manager set instance.
Definition: IRModule.h:273
ReferrentTy * get() const
Definition: NanobindUtils.h:60
Base class for PyOperation and PyOpView which exposes the primary, user visible methods for manipulat...
Definition: IRModule.h:554
virtual PyOperation & getOperation()=0
Each must provide access to the raw Operation.
MlirOperation get() const
Definition: IRModule.h:640
static bool mlirPassManagerIsNull(MlirPassManager passManager)
Checks if a PassManager is null.
Definition: Pass.h:65
MLIR_CAPI_EXPORTED void mlirOpPrintingFlagsElideLargeResourceString(MlirOpPrintingFlags flags, intptr_t largeResourceLimit)
Enables the elision of large resources strings by omitting them from the dialect_resources section.
Definition: IR.cpp:215
MLIR_CAPI_EXPORTED void mlirOpPrintingFlagsPrintGenericOpForm(MlirOpPrintingFlags flags)
Always print operations in the generic form.
Definition: IR.cpp:225
MLIR_CAPI_EXPORTED void mlirOpPrintingFlagsElideLargeElementsAttrs(MlirOpPrintingFlags flags, intptr_t largeElementLimit)
Enables the elision of large elements attributes by printing a lexically valid but otherwise meaningl...
Definition: IR.cpp:210
MLIR_CAPI_EXPORTED void mlirOpPrintingFlagsDestroy(MlirOpPrintingFlags flags)
Destroys printing flags created with mlirOpPrintingFlagsCreate.
Definition: IR.cpp:206
MLIR_CAPI_EXPORTED void mlirOpPrintingFlagsEnableDebugInfo(MlirOpPrintingFlags flags, bool enable, bool prettyForm)
Enable or disable printing of debug information (based on enable).
Definition: IR.cpp:220
MLIR_CAPI_EXPORTED MlirOpPrintingFlags mlirOpPrintingFlagsCreate(void)
Creates new printing flags with defaults, intended for customization.
Definition: IR.cpp:202
static MlirStringRef mlirStringRefCreate(const char *str, size_t length)
Constructs a string reference from the pointer and length.
Definition: Support.h:82
MLIR_CAPI_EXPORTED MlirTypeID mlirTypeIDAllocatorAllocateTypeID(MlirTypeIDAllocator allocator)
Allocates a type id that is valid for the lifetime of the allocator.
Definition: Support.cpp:67
static bool mlirLogicalResultIsFailure(MlirLogicalResult res)
Checks if the given logical result represents a failure.
Definition: Support.h:127
MLIR_CAPI_EXPORTED MlirTypeIDAllocator mlirTypeIDAllocatorCreate(void)
Creates a type id allocator for dynamic type id creation.
Definition: Support.cpp:59
void populatePassManagerSubmodule(nanobind::module_ &m)
Include the generated interface declarations.
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
Structure of external MlirPass callbacks.
Definition: Pass.h:153
void(* run)(MlirOperation op, MlirExternalPass pass, void *userData)
This callback is called when the pass is run.
Definition: Pass.h:174
void *(* clone)(void *userData)
This callback is called when the pass is cloned.
Definition: Pass.h:170
MlirLogicalResult(* initialize)(MlirContext ctx, void *userData)
This callback is optional.
Definition: Pass.h:166
void(* destruct)(void *userData)
This callback is called when the pass is destroyed This is analogous to a C++ pass destructor.
Definition: Pass.h:160
void(* construct)(void *userData)
This callback is called from the pass is created.
Definition: Pass.h:156
A logical result value, essentially a boolean with named states.
Definition: Support.h:116
Accumulates into a python string from a method that accepts an MlirStringCallback.
MlirStringCallback getCallback()
Custom exception that allows access to error diagnostic information.
Definition: IRModule.h:1318
RAII object that captures any error diagnostics emitted to the provided context.
Definition: IRModule.h:409