mlir/lib/Bindings/Python/DialectLLVM.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-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
  23   //===--------------------------------------------------------------------===//
  24   // StructType
  25   //===--------------------------------------------------------------------===//
  26
  27   auto llvmStructType =
  28       mlir_type_subclass(m, "StructType", mlirTypeIsALLVMStructType);
  29
  30   llvmStructType.def_classmethod(
  31       "get_literal",
  32       [](py::object cls, const std::vector<MlirType> &elements, bool packed,
  33          MlirLocation loc) {
  34         CollectDiagnosticsToStringScope scope(mlirLocationGetContext(loc));
  35
  36         MlirType type = mlirLLVMStructTypeLiteralGetChecked(
  37             loc, elements.size(), elements.data(), packed);
  38         if (mlirTypeIsNull(type)) {
  39           throw py::value_error(scope.takeMessage());
  40         }
  41         return cls(type);
  42       },
  43       "cls"_a, "elements"_a, py::kw_only(), "packed"_a = false,
  44       "loc"_a = py::none());
  45
  46   llvmStructType.def_classmethod(
  47       "get_identified",
  48       [](py::object cls, const std::string &name, MlirContext context) {
  49         return cls(mlirLLVMStructTypeIdentifiedGet(
  50             context, mlirStringRefCreate(name.data(), name.size())));
  51       },
  52       "cls"_a, "name"_a, py::kw_only(), "context"_a = py::none());
  53
  54   llvmStructType.def_classmethod(
  55       "get_opaque",
  56       [](py::object cls, const std::string &name, MlirContext context) {
  57         return cls(mlirLLVMStructTypeOpaqueGet(
  58             context, mlirStringRefCreate(name.data(), name.size())));
  59       },
  60       "cls"_a, "name"_a, "context"_a = py::none());
  61
  62   llvmStructType.def(
  63       "set_body",
  64       [](MlirType self, const std::vector<MlirType> &elements, bool packed) {
  65         MlirLogicalResult result = mlirLLVMStructTypeSetBody(
  66             self, elements.size(), elements.data(), packed);
  67         if (!mlirLogicalResultIsSuccess(result)) {
  68           throw py::value_error(
  69               "Struct body already set to different content.");
  70         }
  71       },
  72       "elements"_a, py::kw_only(), "packed"_a = false);
  73
  74   llvmStructType.def_classmethod(
  75       "new_identified",
  76       [](py::object cls, const std::string &name,
  77          const std::vector<MlirType> &elements, bool packed, MlirContext ctx) {
  78         return cls(mlirLLVMStructTypeIdentifiedNewGet(
  79             ctx, mlirStringRefCreate(name.data(), name.length()),
  80             elements.size(), elements.data(), packed));
  81       },
  82       "cls"_a, "name"_a, "elements"_a, py::kw_only(), "packed"_a = false,
  83       "context"_a = py::none());
  84
  85   llvmStructType.def_property_readonly(
  86       "name", [](MlirType type) -> std::optional<std::string> {
  87         if (mlirLLVMStructTypeIsLiteral(type))
  88           return std::nullopt;
  89
  90         MlirStringRef stringRef = mlirLLVMStructTypeGetIdentifier(type);
  91         return StringRef(stringRef.data, stringRef.length).str();
  92       });
  93
  94   llvmStructType.def_property_readonly("body", [](MlirType type) -> py::object {
  95     // Don't crash in absence of a body.
  96     if (mlirLLVMStructTypeIsOpaque(type))
  97       return py::none();
  98
  99     py::list body;
 100     for (intptr_t i = 0, e = mlirLLVMStructTypeGetNumElementTypes(type); i < e;
 101          ++i) {
 102       body.append(mlirLLVMStructTypeGetElementType(type, i));
 103     }
 104     return body;
 105   });
 106
 107   llvmStructType.def_property_readonly(
 108       "packed", [](MlirType type) { return mlirLLVMStructTypeIsPacked(type); });
 109
 110   llvmStructType.def_property_readonly(
 111       "opaque", [](MlirType type) { return mlirLLVMStructTypeIsOpaque(type); });
 112
 113   //===--------------------------------------------------------------------===//
 114   // PointerType
 115   //===--------------------------------------------------------------------===//
 116
 117   mlir_type_subclass(m, "PointerType", mlirTypeIsALLVMPointerType)
 118       .def_classmethod(
 119           "get",
 120           [](py::object cls, std::optional<unsigned> addressSpace,
 121              MlirContext context) {
 122             CollectDiagnosticsToStringScope scope(context);
 123             MlirType type = mlirLLVMPointerTypeGet(
 124                 context, addressSpace.has_value() ? *addressSpace : 0);
 125             if (mlirTypeIsNull(type)) {
 126               throw py::value_error(scope.takeMessage());
 127             }
 128             return cls(type);
 129           },
 130           "cls"_a, "address_space"_a = py::none(), py::kw_only(),
 131           "context"_a = py::none())
 132       .def_property_readonly("address_space", [](MlirType type) {
 133         return mlirLLVMPointerTypeGetAddressSpace(type);
 134       });
 135 }
 136
 137 PYBIND11_MODULE(_mlirDialectsLLVM, m) {
 138   m.doc() = "MLIR LLVM Dialect";
 139
 140   populateDialectLLVMSubmodule(m);
 141 }