Skip to content

[mlir][pybind] export more options on enable_ir_printing() api #65854

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

[mlir][pybind] export more options on enable_ir_printing() api #65854

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions mlir/include/mlir-c/IR.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ DEFINE_C_API_STRUCT(MlirDialectRegistry, void);
DEFINE_C_API_STRUCT(MlirOperation, void);
DEFINE_C_API_STRUCT(MlirOpOperand, void);
DEFINE_C_API_STRUCT(MlirOpPrintingFlags, void);
DEFINE_C_API_STRUCT(MlirIRPrinterConfig, void);
DEFINE_C_API_STRUCT(MlirBlock, void);
DEFINE_C_API_STRUCT(MlirRegion, void);
DEFINE_C_API_STRUCT(MlirSymbolTable, void);
Expand Down Expand Up @@ -450,6 +451,14 @@ mlirOpPrintingFlagsUseLocalScope(MlirOpPrintingFlags flags);
MLIR_CAPI_EXPORTED void
mlirOpPrintingFlagsAssumeVerified(MlirOpPrintingFlags flags);

//===----------------------------------------------------------------------===//
// IR Printing config API.
//===----------------------------------------------------------------------===//

MLIR_CAPI_EXPORTED MlirIRPrinterConfig mlirIRPrinterConfigCreate(void);

MLIR_CAPI_EXPORTED void mlirIRPrinterConfigDestroy(MlirIRPrinterConfig config);

//===----------------------------------------------------------------------===//
// Bytecode printing flags API.
//===----------------------------------------------------------------------===//
Expand Down
8 changes: 7 additions & 1 deletion mlir/include/mlir-c/Pass.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,8 +75,14 @@ MLIR_CAPI_EXPORTED MlirLogicalResult
mlirPassManagerRunOnOp(MlirPassManager passManager, MlirOperation op);

/// Enable mlir-print-ir-after-all.
// MLIR_CAPI_EXPORTED void mlirPassManagerEnableIRPrinting(
// MlirPassManager passManager, bool printBeforePass, bool printAfterPass,
// bool printModuleScope, bool printAfterOnlyOnChange,
// bool printAfterOnlyOnFailure, MlirOpPrintingFlags flags);

MLIR_CAPI_EXPORTED void
mlirPassManagerEnableIRPrinting(MlirPassManager passManager);
mlirPassManagerEnableIRPrinting(MlirPassManager passManager,
MlirIRPrinterConfig config);

/// Enable / disable verify-each.
MLIR_CAPI_EXPORTED void
Expand Down
3 changes: 3 additions & 0 deletions mlir/include/mlir/CAPI/IR.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Operation.h"
#include "mlir/Pass/PassManager.h"

DEFINE_C_API_PTR_METHODS(MlirAsmState, mlir::AsmState)
DEFINE_C_API_PTR_METHODS(MlirBytecodeWriterConfig, mlir::BytecodeWriterConfig)
Expand All @@ -30,6 +31,8 @@ DEFINE_C_API_PTR_METHODS(MlirOperation, mlir::Operation)
DEFINE_C_API_PTR_METHODS(MlirBlock, mlir::Block)
DEFINE_C_API_PTR_METHODS(MlirOpOperand, mlir::OpOperand)
DEFINE_C_API_PTR_METHODS(MlirOpPrintingFlags, mlir::OpPrintingFlags)
DEFINE_C_API_PTR_METHODS(MlirIRPrinterConfig,
mlir::PassManager::IRPrinterConfig)
DEFINE_C_API_PTR_METHODS(MlirRegion, mlir::Region)
DEFINE_C_API_PTR_METHODS(MlirSymbolTable, mlir::SymbolTable)

Expand Down
11 changes: 10 additions & 1 deletion mlir/include/mlir/Pass/PassManager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -307,7 +307,8 @@ class PassManager : public OpPassManager {
/// IR.
explicit IRPrinterConfig(
bool printModuleScope = false, bool printAfterOnlyOnChange = false,
bool printAfterOnlyOnFailure = false,
bool printAfterOnlyOnFailure = false, bool printBeforePass = false,
bool printAfterPass = false,
Copy link
Contributor Author

@qingyunqu qingyunqu Apr 14, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@joker-eph I try to create C struct of PassManager::IRPrinterConfig, but this will change the member variable. I don't know whether it is a good practice.

Copy link
Contributor Author

@qingyunqu qingyunqu Apr 14, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

BTW, this PR is only a draft now, don't care about the detail code.

OpPrintingFlags opPrintingFlags = OpPrintingFlags());
virtual ~IRPrinterConfig();

Expand Down Expand Up @@ -338,6 +339,10 @@ class PassManager : public OpPassManager {
return printAfterOnlyOnFailure;
}

bool shouldPrintBeforePass() const { return printBeforePass; }

bool shouldPrintAfterPass() const { return printAfterPass; }

/// Returns the printing flags to be used to print the IR.
OpPrintingFlags getOpPrintingFlags() const { return opPrintingFlags; }

Expand All @@ -353,6 +358,10 @@ class PassManager : public OpPassManager {
/// the pass failed.
bool printAfterOnlyOnFailure;

bool printBeforePass;

bool printAfterPass;

/// Flags to control printing behavior.
OpPrintingFlags opPrintingFlags;
};
Expand Down
33 changes: 31 additions & 2 deletions mlir/lib/Bindings/Python/Pass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,9 +74,38 @@ void mlir::python::populatePassManagerSubmodule(py::module &m) {
"Releases (leaks) the backing pass manager (testing)")
.def(
"enable_ir_printing",
[](PyPassManager &passManager) {
mlirPassManagerEnableIRPrinting(passManager.get());
[](PyPassManager &passManager, bool printBeforePass,
bool printAfterPass, bool printModuleScope,
bool printAfterOnlyOnChange, bool printAfterOnlyOnFailure,
std::optional<int64_t> largeElementsLimit, bool enableDebugInfo,
bool printGenericOpForm) {
MlirIRPrinterConfig config = mlirIRPrinterConfigCreate();
mlirPassManagerEnableIRPrinting(passManager.get(), config);
mlirIRPrinterConfigDestroy(config);
// MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
// if (largeElementsLimit)
// mlirOpPrintingFlagsElideLargeElementsAttrs(flags,
// *largeElementsLimit);
// if (enableDebugInfo)
// mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,
// /*prettyForm=*/false);
// if (printGenericOpForm)
// mlirOpPrintingFlagsPrintGenericOpForm(flags);
// mlirPassManagerEnableIRPrinting(passManager.get(),
// printBeforePass,
// printAfterPass, printModuleScope,
// printAfterOnlyOnChange,
// printAfterOnlyOnFailure, flags);
// mlirOpPrintingFlagsDestroy(flags);
},
py::arg("print_before_pass") = true,
py::arg("print_after_pass") = true,
py::arg("print_module_scope") = true,
py::arg("print_after_only_on_change") = true,
py::arg("print_after_only_on_failure") = false,
py::arg("large_elements_limit") = py::none(),
py::arg("enable_debug_info") = false,
py::arg("print_generic_op_form") = false,
"Enable mlir-print-ir-after-all.")
.def(
"enable_verifier",
Expand Down
8 changes: 8 additions & 0 deletions mlir/lib/CAPI/IR/IR.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,14 @@ void mlirOpPrintingFlagsAssumeVerified(MlirOpPrintingFlags flags) {
unwrap(flags)->assumeVerified();
}

MlirIRPrinterConfig mlirIRPrinterConfigCreate() {
return wrap(new PassManager::IRPrinterConfig());
}

void mlirIRPrinterConfigDestroy(MlirIRPrinterConfig config) {
delete unwrap(config);
}

//===----------------------------------------------------------------------===//
// Bytecode printing flags API.
//===----------------------------------------------------------------------===//
Expand Down
18 changes: 16 additions & 2 deletions mlir/lib/CAPI/IR/Pass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
#include "mlir/CAPI/Support.h"
#include "mlir/CAPI/Utils.h"
#include "mlir/Pass/PassManager.h"
#include <functional>
#include <optional>

using namespace mlir;
Expand Down Expand Up @@ -44,8 +45,21 @@ MlirLogicalResult mlirPassManagerRunOnOp(MlirPassManager passManager,
return wrap(unwrap(passManager)->run(unwrap(op)));
}

void mlirPassManagerEnableIRPrinting(MlirPassManager passManager) {
return unwrap(passManager)->enableIRPrinting();
void mlirPassManagerEnableIRPrinting(MlirPassManager passManager,
MlirIRPrinterConfig config) {
std::function<bool(Pass *, Operation *)> shouldPrintBeforePass = nullptr;
std::function<bool(Pass *, Operation *)> shouldPrintAfterPass = nullptr;
if (unwrap(config)->shouldPrintBeforePass())
shouldPrintBeforePass = [](Pass *, Operation *) { return true; };
if (unwrap(config)->shouldPrintAfterPass())
shouldPrintAfterPass = [](Pass *, Operation *) { return true; };
return unwrap(passManager)
->enableIRPrinting(shouldPrintBeforePass, shouldPrintAfterPass,
unwrap(config)->shouldPrintAtModuleScope(),
unwrap(config)->shouldPrintAfterOnlyOnChange(),
unwrap(config)->shouldPrintAfterOnlyOnFailure(),
/*out=*/llvm::errs(),
unwrap(config)->getOpPrintingFlags());
}

void mlirPassManagerEnableVerifier(MlirPassManager passManager, bool enable) {
Expand Down
5 changes: 4 additions & 1 deletion mlir/lib/Pass/IRPrinting.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,10 +133,13 @@ void IRPrinterInstrumentation::runAfterPassFailed(Pass *pass, Operation *op) {
PassManager::IRPrinterConfig::IRPrinterConfig(bool printModuleScope,
bool printAfterOnlyOnChange,
bool printAfterOnlyOnFailure,
bool printBeforePass,
bool printAfterPass,
OpPrintingFlags opPrintingFlags)
: printModuleScope(printModuleScope),
printAfterOnlyOnChange(printAfterOnlyOnChange),
printAfterOnlyOnFailure(printAfterOnlyOnFailure),
printBeforePass(printBeforePass), printAfterPass(printAfterPass),
opPrintingFlags(opPrintingFlags) {}
PassManager::IRPrinterConfig::~IRPrinterConfig() = default;

Expand Down Expand Up @@ -172,7 +175,7 @@ struct BasicIRPrinterConfig : public PassManager::IRPrinterConfig {
bool printAfterOnlyOnFailure, OpPrintingFlags opPrintingFlags,
raw_ostream &out)
: IRPrinterConfig(printModuleScope, printAfterOnlyOnChange,
printAfterOnlyOnFailure, opPrintingFlags),
printAfterOnlyOnFailure, false, false, opPrintingFlags),
shouldPrintBeforePass(std::move(shouldPrintBeforePass)),
shouldPrintAfterPass(std::move(shouldPrintAfterPass)), out(out) {
assert((this->shouldPrintBeforePass || this->shouldPrintAfterPass) &&
Expand Down