xref: /llvm-project/mlir/lib/Bindings/Python/DialectLLVM.cpp (revision 91f11611337dde9a8e0a5e19240f6bb4671922c6)
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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir-c/Dialect/LLVM.h"
10 #include "mlir-c/IR.h"
11 #include "mlir-c/Support.h"
12 #include "mlir/Bindings/Python/PybindAdaptors.h"
13 #include <string>
14 
15 namespace py = pybind11;
16 using namespace llvm;
17 using namespace mlir;
18 using namespace mlir::python;
19 using namespace mlir::python::adaptors;
20 
21 void populateDialectLLVMSubmodule(const pybind11::module &m) {
22   auto llvmStructType =
23       mlir_type_subclass(m, "StructType", mlirTypeIsALLVMStructType);
24 
25   llvmStructType.def_classmethod(
26       "get_literal",
27       [](py::object cls, const std::vector<MlirType> &elements, bool packed,
28          MlirLocation loc) {
29         CollectDiagnosticsToStringScope scope(mlirLocationGetContext(loc));
30 
31         MlirType type = mlirLLVMStructTypeLiteralGetChecked(
32             loc, elements.size(), elements.data(), packed);
33         if (mlirTypeIsNull(type)) {
34           throw py::value_error(scope.takeMessage());
35         }
36         return cls(type);
37       },
38       py::arg("cls"), py::arg("elements"), py::kw_only(),
39       py::arg("packed") = false, py::arg("loc") = py::none());
40 
41   llvmStructType.def_classmethod(
42       "get_identified",
43       [](py::object cls, const std::string &name, MlirContext context) {
44         return cls(mlirLLVMStructTypeIdentifiedGet(
45             context, mlirStringRefCreate(name.data(), name.size())));
46       },
47       py::arg("cls"), py::arg("name"), py::kw_only(),
48       py::arg("context") = py::none());
49 
50   llvmStructType.def_classmethod(
51       "get_opaque",
52       [](py::object cls, const std::string &name, MlirContext context) {
53         return cls(mlirLLVMStructTypeOpaqueGet(
54             context, mlirStringRefCreate(name.data(), name.size())));
55       },
56       py::arg("cls"), py::arg("name"), py::arg("context") = py::none());
57 
58   llvmStructType.def(
59       "set_body",
60       [](MlirType self, const std::vector<MlirType> &elements, bool packed) {
61         MlirLogicalResult result = mlirLLVMStructTypeSetBody(
62             self, elements.size(), elements.data(), packed);
63         if (!mlirLogicalResultIsSuccess(result)) {
64           throw py::value_error(
65               "Struct body already set to different content.");
66         }
67       },
68       py::arg("elements"), py::kw_only(), py::arg("packed") = false);
69 
70   llvmStructType.def_classmethod(
71       "new_identified",
72       [](py::object cls, const std::string &name,
73          const std::vector<MlirType> &elements, bool packed, MlirContext ctx) {
74         return cls(mlirLLVMStructTypeIdentifiedNewGet(
75             ctx, mlirStringRefCreate(name.data(), name.length()),
76             elements.size(), elements.data(), packed));
77       },
78       py::arg("cls"), py::arg("name"), py::arg("elements"), py::kw_only(),
79       py::arg("packed") = false, py::arg("context") = py::none());
80 
81   llvmStructType.def_property_readonly(
82       "name", [](MlirType type) -> std::optional<std::string> {
83         if (mlirLLVMStructTypeIsLiteral(type))
84           return std::nullopt;
85 
86         MlirStringRef stringRef = mlirLLVMStructTypeGetIdentifier(type);
87         return StringRef(stringRef.data, stringRef.length).str();
88       });
89 
90   llvmStructType.def_property_readonly("body", [](MlirType type) -> py::object {
91     // Don't crash in absence of a body.
92     if (mlirLLVMStructTypeIsOpaque(type))
93       return py::none();
94 
95     py::list body;
96     for (intptr_t i = 0, e = mlirLLVMStructTypeGetNumElementTypes(type); i < e;
97          ++i) {
98       body.append(mlirLLVMStructTypeGetElementType(type, i));
99     }
100     return body;
101   });
102 
103   llvmStructType.def_property_readonly(
104       "packed", [](MlirType type) { return mlirLLVMStructTypeIsPacked(type); });
105 
106   llvmStructType.def_property_readonly(
107       "opaque", [](MlirType type) { return mlirLLVMStructTypeIsOpaque(type); });
108 }
109 
110 PYBIND11_MODULE(_mlirDialectsLLVM, m) {
111   m.doc() = "MLIR LLVM Dialect";
112 
113   populateDialectLLVMSubmodule(m);
114 }
115