475 lines · cpp
1//===- IRInterfaces.cpp - MLIR IR interfaces pybind -----------------------===//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 <cstdint>10#include <optional>11#include <string>12#include <utility>13#include <vector>14 15#include "IRModule.h"16#include "mlir-c/BuiltinAttributes.h"17#include "mlir-c/IR.h"18#include "mlir-c/Interfaces.h"19#include "mlir-c/Support.h"20#include "mlir/Bindings/Python/Nanobind.h"21#include "llvm/ADT/STLExtras.h"22#include "llvm/ADT/SmallVector.h"23 24namespace nb = nanobind;25 26namespace mlir {27namespace python {28 29constexpr static const char *constructorDoc =30 R"(Creates an interface from a given operation/opview object or from a31subclass of OpView. Raises ValueError if the operation does not implement the32interface.)";33 34constexpr static const char *operationDoc =35 R"(Returns an Operation for which the interface was constructed.)";36 37constexpr static const char *opviewDoc =38 R"(Returns an OpView subclass _instance_ for which the interface was39constructed)";40 41constexpr static const char *inferReturnTypesDoc =42 R"(Given the arguments required to build an operation, attempts to infer43its return types. Raises ValueError on failure.)";44 45constexpr static const char *inferReturnTypeComponentsDoc =46 R"(Given the arguments required to build an operation, attempts to infer47its return shaped type components. Raises ValueError on failure.)";48 49namespace {50 51/// Takes in an optional ist of operands and converts them into a SmallVector52/// of MlirVlaues. Returns an empty SmallVector if the list is empty.53llvm::SmallVector<MlirValue> wrapOperands(std::optional<nb::list> operandList) {54 llvm::SmallVector<MlirValue> mlirOperands;55 56 if (!operandList || operandList->size() == 0) {57 return mlirOperands;58 }59 60 // Note: as the list may contain other lists this may not be final size.61 mlirOperands.reserve(operandList->size());62 for (const auto &&it : llvm::enumerate(*operandList)) {63 if (it.value().is_none())64 continue;65 66 PyValue *val;67 try {68 val = nb::cast<PyValue *>(it.value());69 if (!val)70 throw nb::cast_error();71 mlirOperands.push_back(val->get());72 continue;73 } catch (nb::cast_error &err) {74 // Intentionally unhandled to try sequence below first.75 (void)err;76 }77 78 try {79 auto vals = nb::cast<nb::sequence>(it.value());80 for (nb::handle v : vals) {81 try {82 val = nb::cast<PyValue *>(v);83 if (!val)84 throw nb::cast_error();85 mlirOperands.push_back(val->get());86 } catch (nb::cast_error &err) {87 throw nb::value_error(88 (llvm::Twine("Operand ") + llvm::Twine(it.index()) +89 " must be a Value or Sequence of Values (" + err.what() + ")")90 .str()91 .c_str());92 }93 }94 continue;95 } catch (nb::cast_error &err) {96 throw nb::value_error((llvm::Twine("Operand ") + llvm::Twine(it.index()) +97 " must be a Value or Sequence of Values (" +98 err.what() + ")")99 .str()100 .c_str());101 }102 103 throw nb::cast_error();104 }105 106 return mlirOperands;107}108 109/// Takes in an optional vector of PyRegions and returns a SmallVector of110/// MlirRegion. Returns an empty SmallVector if the list is empty.111llvm::SmallVector<MlirRegion>112wrapRegions(std::optional<std::vector<PyRegion>> regions) {113 llvm::SmallVector<MlirRegion> mlirRegions;114 115 if (regions) {116 mlirRegions.reserve(regions->size());117 for (PyRegion ®ion : *regions) {118 mlirRegions.push_back(region);119 }120 }121 122 return mlirRegions;123}124 125} // namespace126 127/// CRTP base class for Python classes representing MLIR Op interfaces.128/// Interface hierarchies are flat so no base class is expected here. The129/// derived class is expected to define the following static fields:130/// - `const char *pyClassName` - the name of the Python class to create;131/// - `GetTypeIDFunctionTy getInterfaceID` - the function producing the TypeID132/// of the interface.133/// Derived classes may redefine the `bindDerived(ClassTy &)` method to bind134/// interface-specific methods.135///136/// An interface class may be constructed from either an Operation/OpView object137/// or from a subclass of OpView. In the latter case, only the static interface138/// methods are available, similarly to calling ConcereteOp::staticMethod on the139/// C++ side. Implementations of concrete interfaces can use the `isStatic`140/// method to check whether the interface object was constructed from a class or141/// an operation/opview instance. The `getOpName` always succeeds and returns a142/// canonical name of the operation suitable for lookups.143template <typename ConcreteIface>144class PyConcreteOpInterface {145protected:146 using ClassTy = nb::class_<ConcreteIface>;147 using GetTypeIDFunctionTy = MlirTypeID (*)();148 149public:150 /// Constructs an interface instance from an object that is either an151 /// operation or a subclass of OpView. In the latter case, only the static152 /// methods of the interface are accessible to the caller.153 PyConcreteOpInterface(nb::object object, DefaultingPyMlirContext context)154 : obj(std::move(object)) {155 try {156 operation = &nb::cast<PyOperation &>(obj);157 } catch (nb::cast_error &) {158 // Do nothing.159 }160 161 try {162 operation = &nb::cast<PyOpView &>(obj).getOperation();163 } catch (nb::cast_error &) {164 // Do nothing.165 }166 167 if (operation != nullptr) {168 if (!mlirOperationImplementsInterface(*operation,169 ConcreteIface::getInterfaceID())) {170 std::string msg = "the operation does not implement ";171 throw nb::value_error((msg + ConcreteIface::pyClassName).c_str());172 }173 174 MlirIdentifier identifier = mlirOperationGetName(*operation);175 MlirStringRef stringRef = mlirIdentifierStr(identifier);176 opName = std::string(stringRef.data, stringRef.length);177 } else {178 try {179 opName = nb::cast<std::string>(obj.attr("OPERATION_NAME"));180 } catch (nb::cast_error &) {181 throw nb::type_error(182 "Op interface does not refer to an operation or OpView class");183 }184 185 if (!mlirOperationImplementsInterfaceStatic(186 mlirStringRefCreate(opName.data(), opName.length()),187 context.resolve().get(), ConcreteIface::getInterfaceID())) {188 std::string msg = "the operation does not implement ";189 throw nb::value_error((msg + ConcreteIface::pyClassName).c_str());190 }191 }192 }193 194 /// Creates the Python bindings for this class in the given module.195 static void bind(nb::module_ &m) {196 nb::class_<ConcreteIface> cls(m, ConcreteIface::pyClassName);197 cls.def(nb::init<nb::object, DefaultingPyMlirContext>(), nb::arg("object"),198 nb::arg("context") = nb::none(), constructorDoc)199 .def_prop_ro("operation", &PyConcreteOpInterface::getOperationObject,200 operationDoc)201 .def_prop_ro("opview", &PyConcreteOpInterface::getOpView, opviewDoc);202 ConcreteIface::bindDerived(cls);203 }204 205 /// Hook for derived classes to add class-specific bindings.206 static void bindDerived(ClassTy &cls) {}207 208 /// Returns `true` if this object was constructed from a subclass of OpView209 /// rather than from an operation instance.210 bool isStatic() { return operation == nullptr; }211 212 /// Returns the operation instance from which this object was constructed.213 /// Throws a type error if this object was constructed from a subclass of214 /// OpView.215 nb::typed<nb::object, PyOperation> getOperationObject() {216 if (operation == nullptr)217 throw nb::type_error("Cannot get an operation from a static interface");218 return operation->getRef().releaseObject();219 }220 221 /// Returns the opview of the operation instance from which this object was222 /// constructed. Throws a type error if this object was constructed form a223 /// subclass of OpView.224 nb::typed<nb::object, PyOpView> getOpView() {225 if (operation == nullptr)226 throw nb::type_error("Cannot get an opview from a static interface");227 return operation->createOpView();228 }229 230 /// Returns the canonical name of the operation this interface is constructed231 /// from.232 const std::string &getOpName() { return opName; }233 234private:235 PyOperation *operation = nullptr;236 std::string opName;237 nb::object obj;238};239 240/// Python wrapper for InferTypeOpInterface. This interface has only static241/// methods.242class PyInferTypeOpInterface243 : public PyConcreteOpInterface<PyInferTypeOpInterface> {244public:245 using PyConcreteOpInterface<PyInferTypeOpInterface>::PyConcreteOpInterface;246 247 constexpr static const char *pyClassName = "InferTypeOpInterface";248 constexpr static GetTypeIDFunctionTy getInterfaceID =249 &mlirInferTypeOpInterfaceTypeID;250 251 /// C-style user-data structure for type appending callback.252 struct AppendResultsCallbackData {253 std::vector<PyType> &inferredTypes;254 PyMlirContext &pyMlirContext;255 };256 257 /// Appends the types provided as the two first arguments to the user-data258 /// structure (expects AppendResultsCallbackData).259 static void appendResultsCallback(intptr_t nTypes, MlirType *types,260 void *userData) {261 auto *data = static_cast<AppendResultsCallbackData *>(userData);262 data->inferredTypes.reserve(data->inferredTypes.size() + nTypes);263 for (intptr_t i = 0; i < nTypes; ++i) {264 data->inferredTypes.emplace_back(data->pyMlirContext.getRef(), types[i]);265 }266 }267 268 /// Given the arguments required to build an operation, attempts to infer its269 /// return types. Throws value_error on failure.270 std::vector<PyType>271 inferReturnTypes(std::optional<nb::list> operandList,272 std::optional<PyAttribute> attributes, void *properties,273 std::optional<std::vector<PyRegion>> regions,274 DefaultingPyMlirContext context,275 DefaultingPyLocation location) {276 llvm::SmallVector<MlirValue> mlirOperands =277 wrapOperands(std::move(operandList));278 llvm::SmallVector<MlirRegion> mlirRegions = wrapRegions(std::move(regions));279 280 std::vector<PyType> inferredTypes;281 PyMlirContext &pyContext = context.resolve();282 AppendResultsCallbackData data{inferredTypes, pyContext};283 MlirStringRef opNameRef =284 mlirStringRefCreate(getOpName().data(), getOpName().length());285 MlirAttribute attributeDict =286 attributes ? attributes->get() : mlirAttributeGetNull();287 288 MlirLogicalResult result = mlirInferTypeOpInterfaceInferReturnTypes(289 opNameRef, pyContext.get(), location.resolve(), mlirOperands.size(),290 mlirOperands.data(), attributeDict, properties, mlirRegions.size(),291 mlirRegions.data(), &appendResultsCallback, &data);292 293 if (mlirLogicalResultIsFailure(result)) {294 throw nb::value_error("Failed to infer result types");295 }296 297 return inferredTypes;298 }299 300 static void bindDerived(ClassTy &cls) {301 cls.def("inferReturnTypes", &PyInferTypeOpInterface::inferReturnTypes,302 nb::arg("operands") = nb::none(),303 nb::arg("attributes") = nb::none(),304 nb::arg("properties") = nb::none(), nb::arg("regions") = nb::none(),305 nb::arg("context") = nb::none(), nb::arg("loc") = nb::none(),306 inferReturnTypesDoc);307 }308};309 310/// Wrapper around an shaped type components.311class PyShapedTypeComponents {312public:313 PyShapedTypeComponents(MlirType elementType) : elementType(elementType) {}314 PyShapedTypeComponents(nb::list shape, MlirType elementType)315 : shape(std::move(shape)), elementType(elementType), ranked(true) {}316 PyShapedTypeComponents(nb::list shape, MlirType elementType,317 MlirAttribute attribute)318 : shape(std::move(shape)), elementType(elementType), attribute(attribute),319 ranked(true) {}320 PyShapedTypeComponents(PyShapedTypeComponents &) = delete;321 PyShapedTypeComponents(PyShapedTypeComponents &&other) noexcept322 : shape(other.shape), elementType(other.elementType),323 attribute(other.attribute), ranked(other.ranked) {}324 325 static void bind(nb::module_ &m) {326 nb::class_<PyShapedTypeComponents>(m, "ShapedTypeComponents")327 .def_prop_ro(328 "element_type",329 [](PyShapedTypeComponents &self) { return self.elementType; },330 nb::sig("def element_type(self) -> Type"),331 "Returns the element type of the shaped type components.")332 .def_static(333 "get",334 [](PyType &elementType) {335 return PyShapedTypeComponents(elementType);336 },337 nb::arg("element_type"),338 "Create an shaped type components object with only the element "339 "type.")340 .def_static(341 "get",342 [](nb::list shape, PyType &elementType) {343 return PyShapedTypeComponents(std::move(shape), elementType);344 },345 nb::arg("shape"), nb::arg("element_type"),346 "Create a ranked shaped type components object.")347 .def_static(348 "get",349 [](nb::list shape, PyType &elementType, PyAttribute &attribute) {350 return PyShapedTypeComponents(std::move(shape), elementType,351 attribute);352 },353 nb::arg("shape"), nb::arg("element_type"), nb::arg("attribute"),354 "Create a ranked shaped type components object with attribute.")355 .def_prop_ro(356 "has_rank",357 [](PyShapedTypeComponents &self) -> bool { return self.ranked; },358 "Returns whether the given shaped type component is ranked.")359 .def_prop_ro(360 "rank",361 [](PyShapedTypeComponents &self) -> std::optional<nb::int_> {362 if (!self.ranked)363 return {};364 return nb::int_(self.shape.size());365 },366 "Returns the rank of the given ranked shaped type components. If "367 "the shaped type components does not have a rank, None is "368 "returned.")369 .def_prop_ro(370 "shape",371 [](PyShapedTypeComponents &self) -> std::optional<nb::list> {372 if (!self.ranked)373 return {};374 return nb::list(self.shape);375 },376 "Returns the shape of the ranked shaped type components as a list "377 "of integers. Returns none if the shaped type component does not "378 "have a rank.");379 }380 381 nb::object getCapsule();382 static PyShapedTypeComponents createFromCapsule(nb::object capsule);383 384private:385 nb::list shape;386 MlirType elementType;387 MlirAttribute attribute;388 bool ranked{false};389};390 391/// Python wrapper for InferShapedTypeOpInterface. This interface has only392/// static methods.393class PyInferShapedTypeOpInterface394 : public PyConcreteOpInterface<PyInferShapedTypeOpInterface> {395public:396 using PyConcreteOpInterface<397 PyInferShapedTypeOpInterface>::PyConcreteOpInterface;398 399 constexpr static const char *pyClassName = "InferShapedTypeOpInterface";400 constexpr static GetTypeIDFunctionTy getInterfaceID =401 &mlirInferShapedTypeOpInterfaceTypeID;402 403 /// C-style user-data structure for type appending callback.404 struct AppendResultsCallbackData {405 std::vector<PyShapedTypeComponents> &inferredShapedTypeComponents;406 };407 408 /// Appends the shaped type components provided as unpacked shape, element409 /// type, attribute to the user-data.410 static void appendResultsCallback(bool hasRank, intptr_t rank,411 const int64_t *shape, MlirType elementType,412 MlirAttribute attribute, void *userData) {413 auto *data = static_cast<AppendResultsCallbackData *>(userData);414 if (!hasRank) {415 data->inferredShapedTypeComponents.emplace_back(elementType);416 } else {417 nb::list shapeList;418 for (intptr_t i = 0; i < rank; ++i) {419 shapeList.append(shape[i]);420 }421 data->inferredShapedTypeComponents.emplace_back(shapeList, elementType,422 attribute);423 }424 }425 426 /// Given the arguments required to build an operation, attempts to infer the427 /// shaped type components. Throws value_error on failure.428 std::vector<PyShapedTypeComponents> inferReturnTypeComponents(429 std::optional<nb::list> operandList,430 std::optional<PyAttribute> attributes, void *properties,431 std::optional<std::vector<PyRegion>> regions,432 DefaultingPyMlirContext context, DefaultingPyLocation location) {433 llvm::SmallVector<MlirValue> mlirOperands =434 wrapOperands(std::move(operandList));435 llvm::SmallVector<MlirRegion> mlirRegions = wrapRegions(std::move(regions));436 437 std::vector<PyShapedTypeComponents> inferredShapedTypeComponents;438 PyMlirContext &pyContext = context.resolve();439 AppendResultsCallbackData data{inferredShapedTypeComponents};440 MlirStringRef opNameRef =441 mlirStringRefCreate(getOpName().data(), getOpName().length());442 MlirAttribute attributeDict =443 attributes ? attributes->get() : mlirAttributeGetNull();444 445 MlirLogicalResult result = mlirInferShapedTypeOpInterfaceInferReturnTypes(446 opNameRef, pyContext.get(), location.resolve(), mlirOperands.size(),447 mlirOperands.data(), attributeDict, properties, mlirRegions.size(),448 mlirRegions.data(), &appendResultsCallback, &data);449 450 if (mlirLogicalResultIsFailure(result)) {451 throw nb::value_error("Failed to infer result shape type components");452 }453 454 return inferredShapedTypeComponents;455 }456 457 static void bindDerived(ClassTy &cls) {458 cls.def("inferReturnTypeComponents",459 &PyInferShapedTypeOpInterface::inferReturnTypeComponents,460 nb::arg("operands") = nb::none(),461 nb::arg("attributes") = nb::none(), nb::arg("regions") = nb::none(),462 nb::arg("properties") = nb::none(), nb::arg("context") = nb::none(),463 nb::arg("loc") = nb::none(), inferReturnTypeComponentsDoc);464 }465};466 467void populateIRInterfaces(nb::module_ &m) {468 PyInferTypeOpInterface::bind(m);469 PyShapedTypeComponents::bind(m);470 PyInferShapedTypeOpInterface::bind(m);471}472 473} // namespace python474} // namespace mlir475