176 lines · cpp
1//===- DialectLLVM.cpp - Pybind module for LLVM 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 <string>10 11#include "mlir-c/Dialect/LLVM.h"12#include "mlir-c/IR.h"13#include "mlir-c/Support.h"14#include "mlir-c/Target/LLVMIR.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 llvm;24using namespace mlir;25using namespace mlir::python;26using namespace mlir::python::nanobind_adaptors;27 28static void populateDialectLLVMSubmodule(nanobind::module_ &m) {29 30 //===--------------------------------------------------------------------===//31 // StructType32 //===--------------------------------------------------------------------===//33 34 auto llvmStructType = mlir_type_subclass(35 m, "StructType", mlirTypeIsALLVMStructType, mlirLLVMStructTypeGetTypeID);36 37 llvmStructType38 .def_classmethod(39 "get_literal",40 [](const nb::object &cls, const std::vector<MlirType> &elements,41 bool packed, MlirLocation loc) {42 CollectDiagnosticsToStringScope scope(mlirLocationGetContext(loc));43 44 MlirType type = mlirLLVMStructTypeLiteralGetChecked(45 loc, elements.size(), elements.data(), packed);46 if (mlirTypeIsNull(type)) {47 throw nb::value_error(scope.takeMessage().c_str());48 }49 return cls(type);50 },51 "cls"_a, "elements"_a, nb::kw_only(), "packed"_a = false,52 "loc"_a = nb::none())53 .def_classmethod(54 "get_literal_unchecked",55 [](const nb::object &cls, const std::vector<MlirType> &elements,56 bool packed, MlirContext context) {57 CollectDiagnosticsToStringScope scope(context);58 59 MlirType type = mlirLLVMStructTypeLiteralGet(60 context, elements.size(), elements.data(), packed);61 if (mlirTypeIsNull(type)) {62 throw nb::value_error(scope.takeMessage().c_str());63 }64 return cls(type);65 },66 "cls"_a, "elements"_a, nb::kw_only(), "packed"_a = false,67 "context"_a = nb::none());68 69 llvmStructType.def_classmethod(70 "get_identified",71 [](const nb::object &cls, const std::string &name, MlirContext context) {72 return cls(mlirLLVMStructTypeIdentifiedGet(73 context, mlirStringRefCreate(name.data(), name.size())));74 },75 "cls"_a, "name"_a, nb::kw_only(), "context"_a = nb::none());76 77 llvmStructType.def_classmethod(78 "get_opaque",79 [](const nb::object &cls, const std::string &name, MlirContext context) {80 return cls(mlirLLVMStructTypeOpaqueGet(81 context, mlirStringRefCreate(name.data(), name.size())));82 },83 "cls"_a, "name"_a, "context"_a = nb::none());84 85 llvmStructType.def(86 "set_body",87 [](MlirType self, const std::vector<MlirType> &elements, bool packed) {88 MlirLogicalResult result = mlirLLVMStructTypeSetBody(89 self, elements.size(), elements.data(), packed);90 if (!mlirLogicalResultIsSuccess(result)) {91 throw nb::value_error(92 "Struct body already set to different content.");93 }94 },95 "elements"_a, nb::kw_only(), "packed"_a = false);96 97 llvmStructType.def_classmethod(98 "new_identified",99 [](const nb::object &cls, const std::string &name,100 const std::vector<MlirType> &elements, bool packed, MlirContext ctx) {101 return cls(mlirLLVMStructTypeIdentifiedNewGet(102 ctx, mlirStringRefCreate(name.data(), name.length()),103 elements.size(), elements.data(), packed));104 },105 "cls"_a, "name"_a, "elements"_a, nb::kw_only(), "packed"_a = false,106 "context"_a = nb::none());107 108 llvmStructType.def_property_readonly(109 "name", [](MlirType type) -> std::optional<std::string> {110 if (mlirLLVMStructTypeIsLiteral(type))111 return std::nullopt;112 113 MlirStringRef stringRef = mlirLLVMStructTypeGetIdentifier(type);114 return StringRef(stringRef.data, stringRef.length).str();115 });116 117 llvmStructType.def_property_readonly("body", [](MlirType type) -> nb::object {118 // Don't crash in absence of a body.119 if (mlirLLVMStructTypeIsOpaque(type))120 return nb::none();121 122 nb::list body;123 for (intptr_t i = 0, e = mlirLLVMStructTypeGetNumElementTypes(type); i < e;124 ++i) {125 body.append(mlirLLVMStructTypeGetElementType(type, i));126 }127 return body;128 });129 130 llvmStructType.def_property_readonly(131 "packed", [](MlirType type) { return mlirLLVMStructTypeIsPacked(type); });132 133 llvmStructType.def_property_readonly(134 "opaque", [](MlirType type) { return mlirLLVMStructTypeIsOpaque(type); });135 136 //===--------------------------------------------------------------------===//137 // PointerType138 //===--------------------------------------------------------------------===//139 140 mlir_type_subclass(m, "PointerType", mlirTypeIsALLVMPointerType,141 mlirLLVMPointerTypeGetTypeID)142 .def_classmethod(143 "get",144 [](const nb::object &cls, std::optional<unsigned> addressSpace,145 MlirContext context) {146 CollectDiagnosticsToStringScope scope(context);147 MlirType type = mlirLLVMPointerTypeGet(148 context, addressSpace.has_value() ? *addressSpace : 0);149 if (mlirTypeIsNull(type)) {150 throw nb::value_error(scope.takeMessage().c_str());151 }152 return cls(type);153 },154 "cls"_a, "address_space"_a = nb::none(), nb::kw_only(),155 "context"_a = nb::none())156 .def_property_readonly("address_space", [](MlirType type) {157 return mlirLLVMPointerTypeGetAddressSpace(type);158 });159 160 m.def(161 "translate_module_to_llvmir",162 [](MlirOperation module) {163 return mlirTranslateModuleToLLVMIRToString(module);164 },165 // clang-format off166 nb::sig("def translate_module_to_llvmir(module: " MAKE_MLIR_PYTHON_QUALNAME("ir.Operation") ") -> str"),167 // clang-format on168 "module"_a, nb::rv_policy::take_ownership);169}170 171NB_MODULE(_mlirDialectsLLVM, m) {172 m.doc() = "MLIR LLVM Dialect";173 174 populateDialectLLVMSubmodule(m);175}176