89 lines · cpp
1//===- DialectSMT.cpp - Pybind module for SMT dialect API support ---------===//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-exception6//7//===----------------------------------------------------------------------===//8 9#include "NanobindUtils.h"10 11#include "mlir-c/Dialect/SMT.h"12#include "mlir-c/IR.h"13#include "mlir-c/Support.h"14#include "mlir-c/Target/ExportSMTLIB.h"15#include "mlir/Bindings/Python/Diagnostics.h"16#include "mlir/Bindings/Python/Nanobind.h"17#include "mlir/Bindings/Python/NanobindAdaptors.h"18 19namespace nb = nanobind;20 21using namespace nanobind::literals;22 23using namespace mlir;24using namespace mlir::python;25using namespace mlir::python::nanobind_adaptors;26 27static void populateDialectSMTSubmodule(nanobind::module_ &m) {28 29 auto smtBoolType =30 mlir_type_subclass(m, "BoolType", mlirSMTTypeIsABool)31 .def_staticmethod(32 "get",33 [](MlirContext context) { return mlirSMTTypeGetBool(context); },34 "context"_a = nb::none());35 auto smtBitVectorType =36 mlir_type_subclass(m, "BitVectorType", mlirSMTTypeIsABitVector)37 .def_staticmethod(38 "get",39 [](int32_t width, MlirContext context) {40 return mlirSMTTypeGetBitVector(context, width);41 },42 "width"_a, "context"_a = nb::none());43 auto smtIntType =44 mlir_type_subclass(m, "IntType", mlirSMTTypeIsAInt)45 .def_staticmethod(46 "get",47 [](MlirContext context) { return mlirSMTTypeGetInt(context); },48 "context"_a = nb::none());49 50 auto exportSMTLIB = [](MlirOperation module, bool inlineSingleUseValues,51 bool indentLetBody) {52 mlir::python::CollectDiagnosticsToStringScope scope(53 mlirOperationGetContext(module));54 PyPrintAccumulator printAccum;55 MlirLogicalResult result = mlirTranslateOperationToSMTLIB(56 module, printAccum.getCallback(), printAccum.getUserData(),57 inlineSingleUseValues, indentLetBody);58 if (mlirLogicalResultIsSuccess(result))59 return printAccum.join();60 throw nb::value_error(61 ("Failed to export smtlib.\nDiagnostic message " + scope.takeMessage())62 .c_str());63 };64 65 m.def(66 "export_smtlib",67 [&exportSMTLIB](MlirOperation module, bool inlineSingleUseValues,68 bool indentLetBody) {69 return exportSMTLIB(module, inlineSingleUseValues, indentLetBody);70 },71 "module"_a, "inline_single_use_values"_a = false,72 "indent_let_body"_a = false);73 m.def(74 "export_smtlib",75 [&exportSMTLIB](MlirModule module, bool inlineSingleUseValues,76 bool indentLetBody) {77 return exportSMTLIB(mlirModuleGetOperation(module),78 inlineSingleUseValues, indentLetBody);79 },80 "module"_a, "inline_single_use_values"_a = false,81 "indent_let_body"_a = false);82}83 84NB_MODULE(_mlirDialectsSMT, m) {85 m.doc() = "MLIR SMT Dialect";86 87 populateDialectSMTSubmodule(m);88}89