5a5c3176ef
Previously, the MLIR's python binding `smt.export_smtlib(...)` always emit `(reset)` to the end of smtlib string as a solver terminator. This PR added an option to suppress this trailing, as downstream users like python z3 module don't need it.
136 lines
4.7 KiB
C++
136 lines
4.7 KiB
C++
//===- DialectSMT.cpp - Pybind module for SMT dialect API support ---------===//
|
|
//
|
|
// 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/Bindings/Python/NanobindUtils.h"
|
|
|
|
#include "mlir-c/Dialect/SMT.h"
|
|
#include "mlir-c/IR.h"
|
|
#include "mlir-c/Support.h"
|
|
#include "mlir-c/Target/ExportSMTLIB.h"
|
|
#include "mlir/Bindings/Python/Diagnostics.h"
|
|
#include "mlir/Bindings/Python/IRCore.h"
|
|
#include "mlir/Bindings/Python/Nanobind.h"
|
|
#include "mlir/Bindings/Python/NanobindAdaptors.h"
|
|
|
|
namespace nb = nanobind;
|
|
|
|
using namespace nanobind::literals;
|
|
using namespace mlir;
|
|
using namespace mlir::python::nanobind_adaptors;
|
|
|
|
namespace mlir {
|
|
namespace python {
|
|
namespace MLIR_BINDINGS_PYTHON_DOMAIN {
|
|
namespace smt {
|
|
struct BoolType : PyConcreteType<BoolType> {
|
|
static constexpr IsAFunctionTy isaFunction = mlirSMTTypeIsABool;
|
|
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
|
|
mlirSMTBoolTypeGetTypeID;
|
|
static constexpr const char *pyClassName = "BoolType";
|
|
static inline const MlirStringRef name = mlirSMTBoolTypeGetName();
|
|
using Base::Base;
|
|
|
|
static void bindDerived(ClassTy &c) {
|
|
c.def_static(
|
|
"get",
|
|
[](DefaultingPyMlirContext context) {
|
|
return BoolType(context->getRef(),
|
|
mlirSMTTypeGetBool(context.get()->get()));
|
|
},
|
|
nb::arg("context").none() = nb::none());
|
|
}
|
|
};
|
|
|
|
struct BitVectorType : PyConcreteType<BitVectorType> {
|
|
static constexpr IsAFunctionTy isaFunction = mlirSMTTypeIsABitVector;
|
|
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
|
|
mlirSMTBitVectorTypeGetTypeID;
|
|
static constexpr const char *pyClassName = "BitVectorType";
|
|
static inline const MlirStringRef name = mlirSMTBitVectorTypeGetName();
|
|
using Base::Base;
|
|
|
|
static void bindDerived(ClassTy &c) {
|
|
c.def_static(
|
|
"get",
|
|
[](int32_t width, DefaultingPyMlirContext context) {
|
|
return BitVectorType(
|
|
context->getRef(),
|
|
mlirSMTTypeGetBitVector(context.get()->get(), width));
|
|
},
|
|
nb::arg("width"), nb::arg("context").none() = nb::none());
|
|
}
|
|
};
|
|
|
|
struct IntType : PyConcreteType<IntType> {
|
|
static constexpr IsAFunctionTy isaFunction = mlirSMTTypeIsAInt;
|
|
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
|
|
mlirSMTIntTypeGetTypeID;
|
|
static constexpr const char *pyClassName = "IntType";
|
|
static inline const MlirStringRef name = mlirSMTIntTypeGetName();
|
|
using Base::Base;
|
|
|
|
static void bindDerived(ClassTy &c) {
|
|
c.def_static(
|
|
"get",
|
|
[](DefaultingPyMlirContext context) {
|
|
return IntType(context->getRef(),
|
|
mlirSMTTypeGetInt(context.get()->get()));
|
|
},
|
|
nb::arg("context").none() = nb::none());
|
|
}
|
|
};
|
|
|
|
static void populateDialectSMTSubmodule(nanobind::module_ &m) {
|
|
BoolType::bind(m);
|
|
BitVectorType::bind(m);
|
|
IntType::bind(m);
|
|
|
|
auto exportSMTLIB = [](MlirOperation module, bool inlineSingleUseValues,
|
|
bool indentLetBody, bool emitReset) {
|
|
CollectDiagnosticsToStringScope scope(mlirOperationGetContext(module));
|
|
PyPrintAccumulator printAccum;
|
|
MlirLogicalResult result = mlirTranslateOperationToSMTLIB(
|
|
module, printAccum.getCallback(), printAccum.getUserData(),
|
|
inlineSingleUseValues, indentLetBody, emitReset);
|
|
if (mlirLogicalResultIsSuccess(result))
|
|
return printAccum.join();
|
|
throw nb::value_error(
|
|
("Failed to export smtlib.\nDiagnostic message " + scope.takeMessage())
|
|
.c_str());
|
|
};
|
|
|
|
m.def(
|
|
"export_smtlib",
|
|
[&exportSMTLIB](const PyOperation &module, bool inlineSingleUseValues,
|
|
bool indentLetBody, bool emitReset) {
|
|
return exportSMTLIB(module, inlineSingleUseValues, indentLetBody,
|
|
emitReset);
|
|
},
|
|
"module"_a, "inline_single_use_values"_a = false,
|
|
"indent_let_body"_a = false, "emit_reset"_a = true);
|
|
m.def(
|
|
"export_smtlib",
|
|
[&exportSMTLIB](PyModule &module, bool inlineSingleUseValues,
|
|
bool indentLetBody, bool emitReset) {
|
|
return exportSMTLIB(mlirModuleGetOperation(module.get()),
|
|
inlineSingleUseValues, indentLetBody, emitReset);
|
|
},
|
|
"module"_a, "inline_single_use_values"_a = false,
|
|
"indent_let_body"_a = false, "emit_reset"_a = true);
|
|
}
|
|
} // namespace smt
|
|
} // namespace MLIR_BINDINGS_PYTHON_DOMAIN
|
|
} // namespace python
|
|
} // namespace mlir
|
|
|
|
NB_MODULE(_mlirDialectsSMT, m) {
|
|
m.doc() = "MLIR SMT Dialect";
|
|
|
|
python::MLIR_BINDINGS_PYTHON_DOMAIN::smt::populateDialectSMTSubmodule(m);
|
|
}
|