brintos / llvm-project-archived public Read only

1//===- IRModules.cpp - IR Submodules of pybind module ---------------------===//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 "Globals.h"10#include "IRModule.h"11#include "NanobindUtils.h"12#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.13#include "mlir-c/BuiltinAttributes.h"14#include "mlir-c/Debug.h"15#include "mlir-c/Diagnostics.h"16#include "mlir-c/IR.h"17#include "mlir-c/Support.h"18#include "mlir/Bindings/Python/Nanobind.h"19#include "mlir/Bindings/Python/NanobindAdaptors.h"20#include "nanobind/nanobind.h"21#include "nanobind/typing.h"22#include "llvm/ADT/ArrayRef.h"23#include "llvm/ADT/SmallVector.h"24 25#include <optional>26 27namespace nb = nanobind;28using namespace nb::literals;29using namespace mlir;30using namespace mlir::python;31 32using llvm::SmallVector;33using llvm::StringRef;34using llvm::Twine;35 36static const char kModuleParseDocstring[] =37    R"(Parses a module's assembly format from a string.38 39Returns a new MlirModule or raises an MLIRError if the parsing fails.40 41See also: https://mlir.llvm.org/docs/LangRef/42)";43 44static const char kDumpDocstring[] =45    "Dumps a debug representation of the object to stderr.";46 47static const char kValueReplaceAllUsesExceptDocstring[] =48    R"(Replace all uses of this value with the `with` value, except for those49in `exceptions`. `exceptions` can be either a single operation or a list of50operations.51)";52 53//------------------------------------------------------------------------------54// Utilities.55//------------------------------------------------------------------------------56 57/// Helper for creating an @classmethod.58template <class Func, typename... Args>59static nb::object classmethod(Func f, Args... args) {60  nb::object cf = nb::cpp_function(f, args...);61  return nb::borrow<nb::object>((PyClassMethod_New(cf.ptr())));62}63 64static nb::object65createCustomDialectWrapper(const std::string &dialectNamespace,66                           nb::object dialectDescriptor) {67  auto dialectClass = PyGlobals::get().lookupDialectClass(dialectNamespace);68  if (!dialectClass) {69    // Use the base class.70    return nb::cast(PyDialect(std::move(dialectDescriptor)));71  }72 73  // Create the custom implementation.74  return (*dialectClass)(std::move(dialectDescriptor));75}76 77static MlirStringRef toMlirStringRef(const std::string &s) {78  return mlirStringRefCreate(s.data(), s.size());79}80 81static MlirStringRef toMlirStringRef(std::string_view s) {82  return mlirStringRefCreate(s.data(), s.size());83}84 85static MlirStringRef toMlirStringRef(const nb::bytes &s) {86  return mlirStringRefCreate(static_cast<const char *>(s.data()), s.size());87}88 89/// Create a block, using the current location context if no locations are90/// specified.91static MlirBlock createBlock(const nb::sequence &pyArgTypes,92                             const std::optional<nb::sequence> &pyArgLocs) {93  SmallVector<MlirType> argTypes;94  argTypes.reserve(nb::len(pyArgTypes));95  for (const auto &pyType : pyArgTypes)96    argTypes.push_back(nb::cast<PyType &>(pyType));97 98  SmallVector<MlirLocation> argLocs;99  if (pyArgLocs) {100    argLocs.reserve(nb::len(*pyArgLocs));101    for (const auto &pyLoc : *pyArgLocs)102      argLocs.push_back(nb::cast<PyLocation &>(pyLoc));103  } else if (!argTypes.empty()) {104    argLocs.assign(argTypes.size(), DefaultingPyLocation::resolve());105  }106 107  if (argTypes.size() != argLocs.size())108    throw nb::value_error(("Expected " + Twine(argTypes.size()) +109                           " locations, got: " + Twine(argLocs.size()))110                              .str()111                              .c_str());112  return mlirBlockCreate(argTypes.size(), argTypes.data(), argLocs.data());113}114 115/// Wrapper for the global LLVM debugging flag.116struct PyGlobalDebugFlag {117  static void set(nb::object &o, bool enable) {118    nb::ft_lock_guard lock(mutex);119    mlirEnableGlobalDebug(enable);120  }121 122  static bool get(const nb::object &) {123    nb::ft_lock_guard lock(mutex);124    return mlirIsGlobalDebugEnabled();125  }126 127  static void bind(nb::module_ &m) {128    // Debug flags.129    nb::class_<PyGlobalDebugFlag>(m, "_GlobalDebug")130        .def_prop_rw_static("flag", &PyGlobalDebugFlag::get,131                            &PyGlobalDebugFlag::set, "LLVM-wide debug flag.")132        .def_static(133            "set_types",134            [](const std::string &type) {135              nb::ft_lock_guard lock(mutex);136              mlirSetGlobalDebugType(type.c_str());137            },138            "types"_a, "Sets specific debug types to be produced by LLVM.")139        .def_static(140            "set_types",141            [](const std::vector<std::string> &types) {142              std::vector<const char *> pointers;143              pointers.reserve(types.size());144              for (const std::string &str : types)145                pointers.push_back(str.c_str());146              nb::ft_lock_guard lock(mutex);147              mlirSetGlobalDebugTypes(pointers.data(), pointers.size());148            },149            "types"_a,150            "Sets multiple specific debug types to be produced by LLVM.");151  }152 153private:154  static nb::ft_mutex mutex;155};156 157nb::ft_mutex PyGlobalDebugFlag::mutex;158 159struct PyAttrBuilderMap {160  static bool dunderContains(const std::string &attributeKind) {161    return PyGlobals::get().lookupAttributeBuilder(attributeKind).has_value();162  }163  static nb::callable dunderGetItemNamed(const std::string &attributeKind) {164    auto builder = PyGlobals::get().lookupAttributeBuilder(attributeKind);165    if (!builder)166      throw nb::key_error(attributeKind.c_str());167    return *builder;168  }169  static void dunderSetItemNamed(const std::string &attributeKind,170                                 nb::callable func, bool replace) {171    PyGlobals::get().registerAttributeBuilder(attributeKind, std::move(func),172                                              replace);173  }174 175  static void bind(nb::module_ &m) {176    nb::class_<PyAttrBuilderMap>(m, "AttrBuilder")177        .def_static("contains", &PyAttrBuilderMap::dunderContains,178                    "attribute_kind"_a,179                    "Checks whether an attribute builder is registered for the "180                    "given attribute kind.")181        .def_static("get", &PyAttrBuilderMap::dunderGetItemNamed,182                    "attribute_kind"_a,183                    "Gets the registered attribute builder for the given "184                    "attribute kind.")185        .def_static("insert", &PyAttrBuilderMap::dunderSetItemNamed,186                    "attribute_kind"_a, "attr_builder"_a, "replace"_a = false,187                    "Register an attribute builder for building MLIR "188                    "attributes from Python values.");189  }190};191 192//------------------------------------------------------------------------------193// PyBlock194//------------------------------------------------------------------------------195 196nb::object PyBlock::getCapsule() {197  return nb::steal<nb::object>(mlirPythonBlockToCapsule(get()));198}199 200//------------------------------------------------------------------------------201// Collections.202//------------------------------------------------------------------------------203 204namespace {205 206class PyRegionIterator {207public:208  PyRegionIterator(PyOperationRef operation, int nextIndex)209      : operation(std::move(operation)), nextIndex(nextIndex) {}210 211  PyRegionIterator &dunderIter() { return *this; }212 213  PyRegion dunderNext() {214    operation->checkValid();215    if (nextIndex >= mlirOperationGetNumRegions(operation->get())) {216      throw nb::stop_iteration();217    }218    MlirRegion region = mlirOperationGetRegion(operation->get(), nextIndex++);219    return PyRegion(operation, region);220  }221 222  static void bind(nb::module_ &m) {223    nb::class_<PyRegionIterator>(m, "RegionIterator")224        .def("__iter__", &PyRegionIterator::dunderIter,225             "Returns an iterator over the regions in the operation.")226        .def("__next__", &PyRegionIterator::dunderNext,227             "Returns the next region in the iteration.");228  }229 230private:231  PyOperationRef operation;232  intptr_t nextIndex = 0;233};234 235/// Regions of an op are fixed length and indexed numerically so are represented236/// with a sequence-like container.237class PyRegionList : public Sliceable<PyRegionList, PyRegion> {238public:239  static constexpr const char *pyClassName = "RegionSequence";240 241  PyRegionList(PyOperationRef operation, intptr_t startIndex = 0,242               intptr_t length = -1, intptr_t step = 1)243      : Sliceable(startIndex,244                  length == -1 ? mlirOperationGetNumRegions(operation->get())245                               : length,246                  step),247        operation(std::move(operation)) {}248 249  PyRegionIterator dunderIter() {250    operation->checkValid();251    return PyRegionIterator(operation, startIndex);252  }253 254  static void bindDerived(ClassTy &c) {255    c.def("__iter__", &PyRegionList::dunderIter,256          "Returns an iterator over the regions in the sequence.");257  }258 259private:260  /// Give the parent CRTP class access to hook implementations below.261  friend class Sliceable<PyRegionList, PyRegion>;262 263  intptr_t getRawNumElements() {264    operation->checkValid();265    return mlirOperationGetNumRegions(operation->get());266  }267 268  PyRegion getRawElement(intptr_t pos) {269    operation->checkValid();270    return PyRegion(operation, mlirOperationGetRegion(operation->get(), pos));271  }272 273  PyRegionList slice(intptr_t startIndex, intptr_t length, intptr_t step) {274    return PyRegionList(operation, startIndex, length, step);275  }276 277  PyOperationRef operation;278};279 280class PyBlockIterator {281public:282  PyBlockIterator(PyOperationRef operation, MlirBlock next)283      : operation(std::move(operation)), next(next) {}284 285  PyBlockIterator &dunderIter() { return *this; }286 287  PyBlock dunderNext() {288    operation->checkValid();289    if (mlirBlockIsNull(next)) {290      throw nb::stop_iteration();291    }292 293    PyBlock returnBlock(operation, next);294    next = mlirBlockGetNextInRegion(next);295    return returnBlock;296  }297 298  static void bind(nb::module_ &m) {299    nb::class_<PyBlockIterator>(m, "BlockIterator")300        .def("__iter__", &PyBlockIterator::dunderIter,301             "Returns an iterator over the blocks in the operation's region.")302        .def("__next__", &PyBlockIterator::dunderNext,303             "Returns the next block in the iteration.");304  }305 306private:307  PyOperationRef operation;308  MlirBlock next;309};310 311/// Blocks are exposed by the C-API as a forward-only linked list. In Python,312/// we present them as a more full-featured list-like container but optimize313/// it for forward iteration. Blocks are always owned by a region.314class PyBlockList {315public:316  PyBlockList(PyOperationRef operation, MlirRegion region)317      : operation(std::move(operation)), region(region) {}318 319  PyBlockIterator dunderIter() {320    operation->checkValid();321    return PyBlockIterator(operation, mlirRegionGetFirstBlock(region));322  }323 324  intptr_t dunderLen() {325    operation->checkValid();326    intptr_t count = 0;327    MlirBlock block = mlirRegionGetFirstBlock(region);328    while (!mlirBlockIsNull(block)) {329      count += 1;330      block = mlirBlockGetNextInRegion(block);331    }332    return count;333  }334 335  PyBlock dunderGetItem(intptr_t index) {336    operation->checkValid();337    if (index < 0) {338      index += dunderLen();339    }340    if (index < 0) {341      throw nb::index_error("attempt to access out of bounds block");342    }343    MlirBlock block = mlirRegionGetFirstBlock(region);344    while (!mlirBlockIsNull(block)) {345      if (index == 0) {346        return PyBlock(operation, block);347      }348      block = mlirBlockGetNextInRegion(block);349      index -= 1;350    }351    throw nb::index_error("attempt to access out of bounds block");352  }353 354  PyBlock appendBlock(const nb::args &pyArgTypes,355                      const std::optional<nb::sequence> &pyArgLocs) {356    operation->checkValid();357    MlirBlock block =358        createBlock(nb::cast<nb::sequence>(pyArgTypes), pyArgLocs);359    mlirRegionAppendOwnedBlock(region, block);360    return PyBlock(operation, block);361  }362 363  static void bind(nb::module_ &m) {364    nb::class_<PyBlockList>(m, "BlockList")365        .def("__getitem__", &PyBlockList::dunderGetItem,366             "Returns the block at the specified index.")367        .def("__iter__", &PyBlockList::dunderIter,368             "Returns an iterator over blocks in the operation's region.")369        .def("__len__", &PyBlockList::dunderLen,370             "Returns the number of blocks in the operation's region.")371        .def("append", &PyBlockList::appendBlock,372             R"(373              Appends a new block, with argument types as positional args.374 375              Returns:376                The created block.377             )",378             nb::arg("args"), nb::kw_only(),379             nb::arg("arg_locs") = std::nullopt);380  }381 382private:383  PyOperationRef operation;384  MlirRegion region;385};386 387class PyOperationIterator {388public:389  PyOperationIterator(PyOperationRef parentOperation, MlirOperation next)390      : parentOperation(std::move(parentOperation)), next(next) {}391 392  PyOperationIterator &dunderIter() { return *this; }393 394  nb::typed<nb::object, PyOpView> dunderNext() {395    parentOperation->checkValid();396    if (mlirOperationIsNull(next)) {397      throw nb::stop_iteration();398    }399 400    PyOperationRef returnOperation =401        PyOperation::forOperation(parentOperation->getContext(), next);402    next = mlirOperationGetNextInBlock(next);403    return returnOperation->createOpView();404  }405 406  static void bind(nb::module_ &m) {407    nb::class_<PyOperationIterator>(m, "OperationIterator")408        .def("__iter__", &PyOperationIterator::dunderIter,409             "Returns an iterator over the operations in an operation's block.")410        .def("__next__", &PyOperationIterator::dunderNext,411             "Returns the next operation in the iteration.");412  }413 414private:415  PyOperationRef parentOperation;416  MlirOperation next;417};418 419/// Operations are exposed by the C-API as a forward-only linked list. In420/// Python, we present them as a more full-featured list-like container but421/// optimize it for forward iteration. Iterable operations are always owned422/// by a block.423class PyOperationList {424public:425  PyOperationList(PyOperationRef parentOperation, MlirBlock block)426      : parentOperation(std::move(parentOperation)), block(block) {}427 428  PyOperationIterator dunderIter() {429    parentOperation->checkValid();430    return PyOperationIterator(parentOperation,431                               mlirBlockGetFirstOperation(block));432  }433 434  intptr_t dunderLen() {435    parentOperation->checkValid();436    intptr_t count = 0;437    MlirOperation childOp = mlirBlockGetFirstOperation(block);438    while (!mlirOperationIsNull(childOp)) {439      count += 1;440      childOp = mlirOperationGetNextInBlock(childOp);441    }442    return count;443  }444 445  nb::typed<nb::object, PyOpView> dunderGetItem(intptr_t index) {446    parentOperation->checkValid();447    if (index < 0) {448      index += dunderLen();449    }450    if (index < 0) {451      throw nb::index_error("attempt to access out of bounds operation");452    }453    MlirOperation childOp = mlirBlockGetFirstOperation(block);454    while (!mlirOperationIsNull(childOp)) {455      if (index == 0) {456        return PyOperation::forOperation(parentOperation->getContext(), childOp)457            ->createOpView();458      }459      childOp = mlirOperationGetNextInBlock(childOp);460      index -= 1;461    }462    throw nb::index_error("attempt to access out of bounds operation");463  }464 465  static void bind(nb::module_ &m) {466    nb::class_<PyOperationList>(m, "OperationList")467        .def("__getitem__", &PyOperationList::dunderGetItem,468             "Returns the operation at the specified index.")469        .def("__iter__", &PyOperationList::dunderIter,470             "Returns an iterator over operations in the list.")471        .def("__len__", &PyOperationList::dunderLen,472             "Returns the number of operations in the list.");473  }474 475private:476  PyOperationRef parentOperation;477  MlirBlock block;478};479 480class PyOpOperand {481public:482  PyOpOperand(MlirOpOperand opOperand) : opOperand(opOperand) {}483 484  nb::typed<nb::object, PyOpView> getOwner() {485    MlirOperation owner = mlirOpOperandGetOwner(opOperand);486    PyMlirContextRef context =487        PyMlirContext::forContext(mlirOperationGetContext(owner));488    return PyOperation::forOperation(context, owner)->createOpView();489  }490 491  size_t getOperandNumber() { return mlirOpOperandGetOperandNumber(opOperand); }492 493  static void bind(nb::module_ &m) {494    nb::class_<PyOpOperand>(m, "OpOperand")495        .def_prop_ro("owner", &PyOpOperand::getOwner,496                     "Returns the operation that owns this operand.")497        .def_prop_ro("operand_number", &PyOpOperand::getOperandNumber,498                     "Returns the operand number in the owning operation.");499  }500 501private:502  MlirOpOperand opOperand;503};504 505class PyOpOperandIterator {506public:507  PyOpOperandIterator(MlirOpOperand opOperand) : opOperand(opOperand) {}508 509  PyOpOperandIterator &dunderIter() { return *this; }510 511  PyOpOperand dunderNext() {512    if (mlirOpOperandIsNull(opOperand))513      throw nb::stop_iteration();514 515    PyOpOperand returnOpOperand(opOperand);516    opOperand = mlirOpOperandGetNextUse(opOperand);517    return returnOpOperand;518  }519 520  static void bind(nb::module_ &m) {521    nb::class_<PyOpOperandIterator>(m, "OpOperandIterator")522        .def("__iter__", &PyOpOperandIterator::dunderIter,523             "Returns an iterator over operands.")524        .def("__next__", &PyOpOperandIterator::dunderNext,525             "Returns the next operand in the iteration.");526  }527 528private:529  MlirOpOperand opOperand;530};531 532} // namespace533 534//------------------------------------------------------------------------------535// PyMlirContext536//------------------------------------------------------------------------------537 538PyMlirContext::PyMlirContext(MlirContext context) : context(context) {539  nb::gil_scoped_acquire acquire;540  nb::ft_lock_guard lock(live_contexts_mutex);541  auto &liveContexts = getLiveContexts();542  liveContexts[context.ptr] = this;543}544 545PyMlirContext::~PyMlirContext() {546  // Note that the only public way to construct an instance is via the547  // forContext method, which always puts the associated handle into548  // liveContexts.549  nb::gil_scoped_acquire acquire;550  {551    nb::ft_lock_guard lock(live_contexts_mutex);552    getLiveContexts().erase(context.ptr);553  }554  mlirContextDestroy(context);555}556 557nb::object PyMlirContext::getCapsule() {558  return nb::steal<nb::object>(mlirPythonContextToCapsule(get()));559}560 561nb::object PyMlirContext::createFromCapsule(nb::object capsule) {562  MlirContext rawContext = mlirPythonCapsuleToContext(capsule.ptr());563  if (mlirContextIsNull(rawContext))564    throw nb::python_error();565  return forContext(rawContext).releaseObject();566}567 568PyMlirContextRef PyMlirContext::forContext(MlirContext context) {569  nb::gil_scoped_acquire acquire;570  nb::ft_lock_guard lock(live_contexts_mutex);571  auto &liveContexts = getLiveContexts();572  auto it = liveContexts.find(context.ptr);573  if (it == liveContexts.end()) {574    // Create.575    PyMlirContext *unownedContextWrapper = new PyMlirContext(context);576    nb::object pyRef = nb::cast(unownedContextWrapper);577    assert(pyRef && "cast to nb::object failed");578    liveContexts[context.ptr] = unownedContextWrapper;579    return PyMlirContextRef(unownedContextWrapper, std::move(pyRef));580  }581  // Use existing.582  nb::object pyRef = nb::cast(it->second);583  return PyMlirContextRef(it->second, std::move(pyRef));584}585 586nb::ft_mutex PyMlirContext::live_contexts_mutex;587 588PyMlirContext::LiveContextMap &PyMlirContext::getLiveContexts() {589  static LiveContextMap liveContexts;590  return liveContexts;591}592 593size_t PyMlirContext::getLiveCount() {594  nb::ft_lock_guard lock(live_contexts_mutex);595  return getLiveContexts().size();596}597 598nb::object PyMlirContext::contextEnter(nb::object context) {599  return PyThreadContextEntry::pushContext(context);600}601 602void PyMlirContext::contextExit(const nb::object &excType,603                                const nb::object &excVal,604                                const nb::object &excTb) {605  PyThreadContextEntry::popContext(*this);606}607 608nb::object PyMlirContext::attachDiagnosticHandler(nb::object callback) {609  // Note that ownership is transferred to the delete callback below by way of610  // an explicit inc_ref (borrow).611  PyDiagnosticHandler *pyHandler =612      new PyDiagnosticHandler(get(), std::move(callback));613  nb::object pyHandlerObject =614      nb::cast(pyHandler, nb::rv_policy::take_ownership);615  (void)pyHandlerObject.inc_ref();616 617  // In these C callbacks, the userData is a PyDiagnosticHandler* that is618  // guaranteed to be known to pybind.619  auto handlerCallback =620      +[](MlirDiagnostic diagnostic, void *userData) -> MlirLogicalResult {621    PyDiagnostic *pyDiagnostic = new PyDiagnostic(diagnostic);622    nb::object pyDiagnosticObject =623        nb::cast(pyDiagnostic, nb::rv_policy::take_ownership);624 625    auto *pyHandler = static_cast<PyDiagnosticHandler *>(userData);626    bool result = false;627    {628      // Since this can be called from arbitrary C++ contexts, always get the629      // gil.630      nb::gil_scoped_acquire gil;631      try {632        result = nb::cast<bool>(pyHandler->callback(pyDiagnostic));633      } catch (std::exception &e) {634        fprintf(stderr, "MLIR Python Diagnostic handler raised exception: %s\n",635                e.what());636        pyHandler->hadError = true;637      }638    }639 640    pyDiagnostic->invalidate();641    return result ? mlirLogicalResultSuccess() : mlirLogicalResultFailure();642  };643  auto deleteCallback = +[](void *userData) {644    auto *pyHandler = static_cast<PyDiagnosticHandler *>(userData);645    assert(pyHandler->registeredID && "handler is not registered");646    pyHandler->registeredID.reset();647 648    // Decrement reference, balancing the inc_ref() above.649    nb::object pyHandlerObject = nb::cast(pyHandler, nb::rv_policy::reference);650    pyHandlerObject.dec_ref();651  };652 653  pyHandler->registeredID = mlirContextAttachDiagnosticHandler(654      get(), handlerCallback, static_cast<void *>(pyHandler), deleteCallback);655  return pyHandlerObject;656}657 658MlirLogicalResult PyMlirContext::ErrorCapture::handler(MlirDiagnostic diag,659                                                       void *userData) {660  auto *self = static_cast<ErrorCapture *>(userData);661  // Check if the context requested we emit errors instead of capturing them.662  if (self->ctx->emitErrorDiagnostics)663    return mlirLogicalResultFailure();664 665  if (mlirDiagnosticGetSeverity(diag) != MlirDiagnosticError)666    return mlirLogicalResultFailure();667 668  self->errors.emplace_back(PyDiagnostic(diag).getInfo());669  return mlirLogicalResultSuccess();670}671 672PyMlirContext &DefaultingPyMlirContext::resolve() {673  PyMlirContext *context = PyThreadContextEntry::getDefaultContext();674  if (!context) {675    throw std::runtime_error(676        "An MLIR function requires a Context but none was provided in the call "677        "or from the surrounding environment. Either pass to the function with "678        "a 'context=' argument or establish a default using 'with Context():'");679  }680  return *context;681}682 683//------------------------------------------------------------------------------684// PyThreadContextEntry management685//------------------------------------------------------------------------------686 687std::vector<PyThreadContextEntry> &PyThreadContextEntry::getStack() {688  static thread_local std::vector<PyThreadContextEntry> stack;689  return stack;690}691 692PyThreadContextEntry *PyThreadContextEntry::getTopOfStack() {693  auto &stack = getStack();694  if (stack.empty())695    return nullptr;696  return &stack.back();697}698 699void PyThreadContextEntry::push(FrameKind frameKind, nb::object context,700                                nb::object insertionPoint,701                                nb::object location) {702  auto &stack = getStack();703  stack.emplace_back(frameKind, std::move(context), std::move(insertionPoint),704                     std::move(location));705  // If the new stack has more than one entry and the context of the new top706  // entry matches the previous, copy the insertionPoint and location from the707  // previous entry if missing from the new top entry.708  if (stack.size() > 1) {709    auto &prev = *(stack.rbegin() + 1);710    auto &current = stack.back();711    if (current.context.is(prev.context)) {712      // Default non-context objects from the previous entry.713      if (!current.insertionPoint)714        current.insertionPoint = prev.insertionPoint;715      if (!current.location)716        current.location = prev.location;717    }718  }719}720 721PyMlirContext *PyThreadContextEntry::getContext() {722  if (!context)723    return nullptr;724  return nb::cast<PyMlirContext *>(context);725}726 727PyInsertionPoint *PyThreadContextEntry::getInsertionPoint() {728  if (!insertionPoint)729    return nullptr;730  return nb::cast<PyInsertionPoint *>(insertionPoint);731}732 733PyLocation *PyThreadContextEntry::getLocation() {734  if (!location)735    return nullptr;736  return nb::cast<PyLocation *>(location);737}738 739PyMlirContext *PyThreadContextEntry::getDefaultContext() {740  auto *tos = getTopOfStack();741  return tos ? tos->getContext() : nullptr;742}743 744PyInsertionPoint *PyThreadContextEntry::getDefaultInsertionPoint() {745  auto *tos = getTopOfStack();746  return tos ? tos->getInsertionPoint() : nullptr;747}748 749PyLocation *PyThreadContextEntry::getDefaultLocation() {750  auto *tos = getTopOfStack();751  return tos ? tos->getLocation() : nullptr;752}753 754nb::object PyThreadContextEntry::pushContext(nb::object context) {755  push(FrameKind::Context, /*context=*/context,756       /*insertionPoint=*/nb::object(),757       /*location=*/nb::object());758  return context;759}760 761void PyThreadContextEntry::popContext(PyMlirContext &context) {762  auto &stack = getStack();763  if (stack.empty())764    throw std::runtime_error("Unbalanced Context enter/exit");765  auto &tos = stack.back();766  if (tos.frameKind != FrameKind::Context && tos.getContext() != &context)767    throw std::runtime_error("Unbalanced Context enter/exit");768  stack.pop_back();769}770 771nb::object772PyThreadContextEntry::pushInsertionPoint(nb::object insertionPointObj) {773  PyInsertionPoint &insertionPoint =774      nb::cast<PyInsertionPoint &>(insertionPointObj);775  nb::object contextObj =776      insertionPoint.getBlock().getParentOperation()->getContext().getObject();777  push(FrameKind::InsertionPoint,778       /*context=*/contextObj,779       /*insertionPoint=*/insertionPointObj,780       /*location=*/nb::object());781  return insertionPointObj;782}783 784void PyThreadContextEntry::popInsertionPoint(PyInsertionPoint &insertionPoint) {785  auto &stack = getStack();786  if (stack.empty())787    throw std::runtime_error("Unbalanced InsertionPoint enter/exit");788  auto &tos = stack.back();789  if (tos.frameKind != FrameKind::InsertionPoint &&790      tos.getInsertionPoint() != &insertionPoint)791    throw std::runtime_error("Unbalanced InsertionPoint enter/exit");792  stack.pop_back();793}794 795nb::object PyThreadContextEntry::pushLocation(nb::object locationObj) {796  PyLocation &location = nb::cast<PyLocation &>(locationObj);797  nb::object contextObj = location.getContext().getObject();798  push(FrameKind::Location, /*context=*/contextObj,799       /*insertionPoint=*/nb::object(),800       /*location=*/locationObj);801  return locationObj;802}803 804void PyThreadContextEntry::popLocation(PyLocation &location) {805  auto &stack = getStack();806  if (stack.empty())807    throw std::runtime_error("Unbalanced Location enter/exit");808  auto &tos = stack.back();809  if (tos.frameKind != FrameKind::Location && tos.getLocation() != &location)810    throw std::runtime_error("Unbalanced Location enter/exit");811  stack.pop_back();812}813 814//------------------------------------------------------------------------------815// PyDiagnostic*816//------------------------------------------------------------------------------817 818void PyDiagnostic::invalidate() {819  valid = false;820  if (materializedNotes) {821    for (nb::handle noteObject : *materializedNotes) {822      PyDiagnostic *note = nb::cast<PyDiagnostic *>(noteObject);823      note->invalidate();824    }825  }826}827 828PyDiagnosticHandler::PyDiagnosticHandler(MlirContext context,829                                         nb::object callback)830    : context(context), callback(std::move(callback)) {}831 832PyDiagnosticHandler::~PyDiagnosticHandler() = default;833 834void PyDiagnosticHandler::detach() {835  if (!registeredID)836    return;837  MlirDiagnosticHandlerID localID = *registeredID;838  mlirContextDetachDiagnosticHandler(context, localID);839  assert(!registeredID && "should have unregistered");840  // Not strictly necessary but keeps stale pointers from being around to cause841  // issues.842  context = {nullptr};843}844 845void PyDiagnostic::checkValid() {846  if (!valid) {847    throw std::invalid_argument(848        "Diagnostic is invalid (used outside of callback)");849  }850}851 852MlirDiagnosticSeverity PyDiagnostic::getSeverity() {853  checkValid();854  return mlirDiagnosticGetSeverity(diagnostic);855}856 857PyLocation PyDiagnostic::getLocation() {858  checkValid();859  MlirLocation loc = mlirDiagnosticGetLocation(diagnostic);860  MlirContext context = mlirLocationGetContext(loc);861  return PyLocation(PyMlirContext::forContext(context), loc);862}863 864nb::str PyDiagnostic::getMessage() {865  checkValid();866  nb::object fileObject = nb::module_::import_("io").attr("StringIO")();867  PyFileAccumulator accum(fileObject, /*binary=*/false);868  mlirDiagnosticPrint(diagnostic, accum.getCallback(), accum.getUserData());869  return nb::cast<nb::str>(fileObject.attr("getvalue")());870}871 872nb::tuple PyDiagnostic::getNotes() {873  checkValid();874  if (materializedNotes)875    return *materializedNotes;876  intptr_t numNotes = mlirDiagnosticGetNumNotes(diagnostic);877  nb::tuple notes = nb::steal<nb::tuple>(PyTuple_New(numNotes));878  for (intptr_t i = 0; i < numNotes; ++i) {879    MlirDiagnostic noteDiag = mlirDiagnosticGetNote(diagnostic, i);880    nb::object diagnostic = nb::cast(PyDiagnostic(noteDiag));881    PyTuple_SET_ITEM(notes.ptr(), i, diagnostic.release().ptr());882  }883  materializedNotes = std::move(notes);884 885  return *materializedNotes;886}887 888PyDiagnostic::DiagnosticInfo PyDiagnostic::getInfo() {889  std::vector<DiagnosticInfo> notes;890  for (nb::handle n : getNotes())891    notes.emplace_back(nb::cast<PyDiagnostic>(n).getInfo());892  return {getSeverity(), getLocation(), nb::cast<std::string>(getMessage()),893          std::move(notes)};894}895 896//------------------------------------------------------------------------------897// PyDialect, PyDialectDescriptor, PyDialects, PyDialectRegistry898//------------------------------------------------------------------------------899 900MlirDialect PyDialects::getDialectForKey(const std::string &key,901                                         bool attrError) {902  MlirDialect dialect = mlirContextGetOrLoadDialect(getContext()->get(),903                                                    {key.data(), key.size()});904  if (mlirDialectIsNull(dialect)) {905    std::string msg = (Twine("Dialect '") + key + "' not found").str();906    if (attrError)907      throw nb::attribute_error(msg.c_str());908    throw nb::index_error(msg.c_str());909  }910  return dialect;911}912 913nb::object PyDialectRegistry::getCapsule() {914  return nb::steal<nb::object>(mlirPythonDialectRegistryToCapsule(*this));915}916 917PyDialectRegistry PyDialectRegistry::createFromCapsule(nb::object capsule) {918  MlirDialectRegistry rawRegistry =919      mlirPythonCapsuleToDialectRegistry(capsule.ptr());920  if (mlirDialectRegistryIsNull(rawRegistry))921    throw nb::python_error();922  return PyDialectRegistry(rawRegistry);923}924 925//------------------------------------------------------------------------------926// PyLocation927//------------------------------------------------------------------------------928 929nb::object PyLocation::getCapsule() {930  return nb::steal<nb::object>(mlirPythonLocationToCapsule(*this));931}932 933PyLocation PyLocation::createFromCapsule(nb::object capsule) {934  MlirLocation rawLoc = mlirPythonCapsuleToLocation(capsule.ptr());935  if (mlirLocationIsNull(rawLoc))936    throw nb::python_error();937  return PyLocation(PyMlirContext::forContext(mlirLocationGetContext(rawLoc)),938                    rawLoc);939}940 941nb::object PyLocation::contextEnter(nb::object locationObj) {942  return PyThreadContextEntry::pushLocation(locationObj);943}944 945void PyLocation::contextExit(const nb::object &excType,946                             const nb::object &excVal,947                             const nb::object &excTb) {948  PyThreadContextEntry::popLocation(*this);949}950 951PyLocation &DefaultingPyLocation::resolve() {952  auto *location = PyThreadContextEntry::getDefaultLocation();953  if (!location) {954    throw std::runtime_error(955        "An MLIR function requires a Location but none was provided in the "956        "call or from the surrounding environment. Either pass to the function "957        "with a 'loc=' argument or establish a default using 'with loc:'");958  }959  return *location;960}961 962//------------------------------------------------------------------------------963// PyModule964//------------------------------------------------------------------------------965 966PyModule::PyModule(PyMlirContextRef contextRef, MlirModule module)967    : BaseContextObject(std::move(contextRef)), module(module) {}968 969PyModule::~PyModule() {970  nb::gil_scoped_acquire acquire;971  auto &liveModules = getContext()->liveModules;972  assert(liveModules.count(module.ptr) == 1 &&973         "destroying module not in live map");974  liveModules.erase(module.ptr);975  mlirModuleDestroy(module);976}977 978PyModuleRef PyModule::forModule(MlirModule module) {979  MlirContext context = mlirModuleGetContext(module);980  PyMlirContextRef contextRef = PyMlirContext::forContext(context);981 982  nb::gil_scoped_acquire acquire;983  auto &liveModules = contextRef->liveModules;984  auto it = liveModules.find(module.ptr);985  if (it == liveModules.end()) {986    // Create.987    PyModule *unownedModule = new PyModule(std::move(contextRef), module);988    // Note that the default return value policy on cast is automatic_reference,989    // which does not take ownership (delete will not be called).990    // Just be explicit.991    nb::object pyRef = nb::cast(unownedModule, nb::rv_policy::take_ownership);992    unownedModule->handle = pyRef;993    liveModules[module.ptr] =994        std::make_pair(unownedModule->handle, unownedModule);995    return PyModuleRef(unownedModule, std::move(pyRef));996  }997  // Use existing.998  PyModule *existing = it->second.second;999  nb::object pyRef = nb::borrow<nb::object>(it->second.first);1000  return PyModuleRef(existing, std::move(pyRef));1001}1002 1003nb::object PyModule::createFromCapsule(nb::object capsule) {1004  MlirModule rawModule = mlirPythonCapsuleToModule(capsule.ptr());1005  if (mlirModuleIsNull(rawModule))1006    throw nb::python_error();1007  return forModule(rawModule).releaseObject();1008}1009 1010nb::object PyModule::getCapsule() {1011  return nb::steal<nb::object>(mlirPythonModuleToCapsule(get()));1012}1013 1014//------------------------------------------------------------------------------1015// PyOperation1016//------------------------------------------------------------------------------1017 1018PyOperation::PyOperation(PyMlirContextRef contextRef, MlirOperation operation)1019    : BaseContextObject(std::move(contextRef)), operation(operation) {}1020 1021PyOperation::~PyOperation() {1022  // If the operation has already been invalidated there is nothing to do.1023  if (!valid)1024    return;1025  // Otherwise, invalidate the operation when it is attached.1026  if (isAttached())1027    setInvalid();1028  else {1029    // And destroy it when it is detached, i.e. owned by Python.1030    erase();1031  }1032}1033 1034namespace {1035 1036// Constructs a new object of type T in-place on the Python heap, returning a1037// PyObjectRef to it, loosely analogous to std::make_shared<T>().1038template <typename T, class... Args>1039PyObjectRef<T> makeObjectRef(Args &&...args) {1040  nb::handle type = nb::type<T>();1041  nb::object instance = nb::inst_alloc(type);1042  T *ptr = nb::inst_ptr<T>(instance);1043  new (ptr) T(std::forward<Args>(args)...);1044  nb::inst_mark_ready(instance);1045  return PyObjectRef<T>(ptr, std::move(instance));1046}1047 1048} // namespace1049 1050PyOperationRef PyOperation::createInstance(PyMlirContextRef contextRef,1051                                           MlirOperation operation,1052                                           nb::object parentKeepAlive) {1053  // Create.1054  PyOperationRef unownedOperation =1055      makeObjectRef<PyOperation>(std::move(contextRef), operation);1056  unownedOperation->handle = unownedOperation.getObject();1057  if (parentKeepAlive) {1058    unownedOperation->parentKeepAlive = std::move(parentKeepAlive);1059  }1060  return unownedOperation;1061}1062 1063PyOperationRef PyOperation::forOperation(PyMlirContextRef contextRef,1064                                         MlirOperation operation,1065                                         nb::object parentKeepAlive) {1066  return createInstance(std::move(contextRef), operation,1067                        std::move(parentKeepAlive));1068}1069 1070PyOperationRef PyOperation::createDetached(PyMlirContextRef contextRef,1071                                           MlirOperation operation,1072                                           nb::object parentKeepAlive) {1073  PyOperationRef created = createInstance(std::move(contextRef), operation,1074                                          std::move(parentKeepAlive));1075  created->attached = false;1076  return created;1077}1078 1079PyOperationRef PyOperation::parse(PyMlirContextRef contextRef,1080                                  const std::string &sourceStr,1081                                  const std::string &sourceName) {1082  PyMlirContext::ErrorCapture errors(contextRef);1083  MlirOperation op =1084      mlirOperationCreateParse(contextRef->get(), toMlirStringRef(sourceStr),1085                               toMlirStringRef(sourceName));1086  if (mlirOperationIsNull(op))1087    throw MLIRError("Unable to parse operation assembly", errors.take());1088  return PyOperation::createDetached(std::move(contextRef), op);1089}1090 1091void PyOperation::checkValid() const {1092  if (!valid) {1093    throw std::runtime_error("the operation has been invalidated");1094  }1095}1096 1097void PyOperationBase::print(std::optional<int64_t> largeElementsLimit,1098                            std::optional<int64_t> largeResourceLimit,1099                            bool enableDebugInfo, bool prettyDebugInfo,1100                            bool printGenericOpForm, bool useLocalScope,1101                            bool useNameLocAsPrefix, bool assumeVerified,1102                            nb::object fileObject, bool binary,1103                            bool skipRegions) {1104  PyOperation &operation = getOperation();1105  operation.checkValid();1106  if (fileObject.is_none())1107    fileObject = nb::module_::import_("sys").attr("stdout");1108 1109  MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();1110  if (largeElementsLimit)1111    mlirOpPrintingFlagsElideLargeElementsAttrs(flags, *largeElementsLimit);1112  if (largeResourceLimit)1113    mlirOpPrintingFlagsElideLargeResourceString(flags, *largeResourceLimit);1114  if (enableDebugInfo)1115    mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,1116                                       /*prettyForm=*/prettyDebugInfo);1117  if (printGenericOpForm)1118    mlirOpPrintingFlagsPrintGenericOpForm(flags);1119  if (useLocalScope)1120    mlirOpPrintingFlagsUseLocalScope(flags);1121  if (assumeVerified)1122    mlirOpPrintingFlagsAssumeVerified(flags);1123  if (skipRegions)1124    mlirOpPrintingFlagsSkipRegions(flags);1125  if (useNameLocAsPrefix)1126    mlirOpPrintingFlagsPrintNameLocAsPrefix(flags);1127 1128  PyFileAccumulator accum(fileObject, binary);1129  mlirOperationPrintWithFlags(operation, flags, accum.getCallback(),1130                              accum.getUserData());1131  mlirOpPrintingFlagsDestroy(flags);1132}1133 1134void PyOperationBase::print(PyAsmState &state, nb::object fileObject,1135                            bool binary) {1136  PyOperation &operation = getOperation();1137  operation.checkValid();1138  if (fileObject.is_none())1139    fileObject = nb::module_::import_("sys").attr("stdout");1140  PyFileAccumulator accum(fileObject, binary);1141  mlirOperationPrintWithState(operation, state.get(), accum.getCallback(),1142                              accum.getUserData());1143}1144 1145void PyOperationBase::writeBytecode(const nb::object &fileOrStringObject,1146                                    std::optional<int64_t> bytecodeVersion) {1147  PyOperation &operation = getOperation();1148  operation.checkValid();1149  PyFileAccumulator accum(fileOrStringObject, /*binary=*/true);1150 1151  if (!bytecodeVersion.has_value())1152    return mlirOperationWriteBytecode(operation, accum.getCallback(),1153                                      accum.getUserData());1154 1155  MlirBytecodeWriterConfig config = mlirBytecodeWriterConfigCreate();1156  mlirBytecodeWriterConfigDesiredEmitVersion(config, *bytecodeVersion);1157  MlirLogicalResult res = mlirOperationWriteBytecodeWithConfig(1158      operation, config, accum.getCallback(), accum.getUserData());1159  mlirBytecodeWriterConfigDestroy(config);1160  if (mlirLogicalResultIsFailure(res))1161    throw nb::value_error((Twine("Unable to honor desired bytecode version ") +1162                           Twine(*bytecodeVersion))1163                              .str()1164                              .c_str());1165}1166 1167void PyOperationBase::walk(1168    std::function<MlirWalkResult(MlirOperation)> callback,1169    MlirWalkOrder walkOrder) {1170  PyOperation &operation = getOperation();1171  operation.checkValid();1172  struct UserData {1173    std::function<MlirWalkResult(MlirOperation)> callback;1174    bool gotException;1175    std::string exceptionWhat;1176    nb::object exceptionType;1177  };1178  UserData userData{callback, false, {}, {}};1179  MlirOperationWalkCallback walkCallback = [](MlirOperation op,1180                                              void *userData) {1181    UserData *calleeUserData = static_cast<UserData *>(userData);1182    try {1183      return (calleeUserData->callback)(op);1184    } catch (nb::python_error &e) {1185      calleeUserData->gotException = true;1186      calleeUserData->exceptionWhat = std::string(e.what());1187      calleeUserData->exceptionType = nb::borrow(e.type());1188      return MlirWalkResult::MlirWalkResultInterrupt;1189    }1190  };1191  mlirOperationWalk(operation, walkCallback, &userData, walkOrder);1192  if (userData.gotException) {1193    std::string message("Exception raised in callback: ");1194    message.append(userData.exceptionWhat);1195    throw std::runtime_error(message);1196  }1197}1198 1199nb::object PyOperationBase::getAsm(bool binary,1200                                   std::optional<int64_t> largeElementsLimit,1201                                   std::optional<int64_t> largeResourceLimit,1202                                   bool enableDebugInfo, bool prettyDebugInfo,1203                                   bool printGenericOpForm, bool useLocalScope,1204                                   bool useNameLocAsPrefix, bool assumeVerified,1205                                   bool skipRegions) {1206  nb::object fileObject;1207  if (binary) {1208    fileObject = nb::module_::import_("io").attr("BytesIO")();1209  } else {1210    fileObject = nb::module_::import_("io").attr("StringIO")();1211  }1212  print(/*largeElementsLimit=*/largeElementsLimit,1213        /*largeResourceLimit=*/largeResourceLimit,1214        /*enableDebugInfo=*/enableDebugInfo,1215        /*prettyDebugInfo=*/prettyDebugInfo,1216        /*printGenericOpForm=*/printGenericOpForm,1217        /*useLocalScope=*/useLocalScope,1218        /*useNameLocAsPrefix=*/useNameLocAsPrefix,1219        /*assumeVerified=*/assumeVerified,1220        /*fileObject=*/fileObject,1221        /*binary=*/binary,1222        /*skipRegions=*/skipRegions);1223 1224  return fileObject.attr("getvalue")();1225}1226 1227void PyOperationBase::moveAfter(PyOperationBase &other) {1228  PyOperation &operation = getOperation();1229  PyOperation &otherOp = other.getOperation();1230  operation.checkValid();1231  otherOp.checkValid();1232  mlirOperationMoveAfter(operation, otherOp);1233  operation.parentKeepAlive = otherOp.parentKeepAlive;1234}1235 1236void PyOperationBase::moveBefore(PyOperationBase &other) {1237  PyOperation &operation = getOperation();1238  PyOperation &otherOp = other.getOperation();1239  operation.checkValid();1240  otherOp.checkValid();1241  mlirOperationMoveBefore(operation, otherOp);1242  operation.parentKeepAlive = otherOp.parentKeepAlive;1243}1244 1245bool PyOperationBase::isBeforeInBlock(PyOperationBase &other) {1246  PyOperation &operation = getOperation();1247  PyOperation &otherOp = other.getOperation();1248  operation.checkValid();1249  otherOp.checkValid();1250  return mlirOperationIsBeforeInBlock(operation, otherOp);1251}1252 1253bool PyOperationBase::verify() {1254  PyOperation &op = getOperation();1255  PyMlirContext::ErrorCapture errors(op.getContext());1256  if (!mlirOperationVerify(op.get()))1257    throw MLIRError("Verification failed", errors.take());1258  return true;1259}1260 1261std::optional<PyOperationRef> PyOperation::getParentOperation() {1262  checkValid();1263  if (!isAttached())1264    throw nb::value_error("Detached operations have no parent");1265  MlirOperation operation = mlirOperationGetParentOperation(get());1266  if (mlirOperationIsNull(operation))1267    return {};1268  return PyOperation::forOperation(getContext(), operation);1269}1270 1271PyBlock PyOperation::getBlock() {1272  checkValid();1273  std::optional<PyOperationRef> parentOperation = getParentOperation();1274  MlirBlock block = mlirOperationGetBlock(get());1275  assert(!mlirBlockIsNull(block) && "Attached operation has null parent");1276  assert(parentOperation && "Operation has no parent");1277  return PyBlock{std::move(*parentOperation), block};1278}1279 1280nb::object PyOperation::getCapsule() {1281  checkValid();1282  return nb::steal<nb::object>(mlirPythonOperationToCapsule(get()));1283}1284 1285nb::object PyOperation::createFromCapsule(const nb::object &capsule) {1286  MlirOperation rawOperation = mlirPythonCapsuleToOperation(capsule.ptr());1287  if (mlirOperationIsNull(rawOperation))1288    throw nb::python_error();1289  MlirContext rawCtxt = mlirOperationGetContext(rawOperation);1290  return forOperation(PyMlirContext::forContext(rawCtxt), rawOperation)1291      .releaseObject();1292}1293 1294static void maybeInsertOperation(PyOperationRef &op,1295                                 const nb::object &maybeIp) {1296  // InsertPoint active?1297  if (!maybeIp.is(nb::cast(false))) {1298    PyInsertionPoint *ip;1299    if (maybeIp.is_none()) {1300      ip = PyThreadContextEntry::getDefaultInsertionPoint();1301    } else {1302      ip = nb::cast<PyInsertionPoint *>(maybeIp);1303    }1304    if (ip)1305      ip->insert(*op.get());1306  }1307}1308 1309nb::object PyOperation::create(std::string_view name,1310                               std::optional<std::vector<PyType *>> results,1311                               llvm::ArrayRef<MlirValue> operands,1312                               std::optional<nb::dict> attributes,1313                               std::optional<std::vector<PyBlock *>> successors,1314                               int regions, PyLocation &location,1315                               const nb::object &maybeIp, bool inferType) {1316  llvm::SmallVector<MlirType, 4> mlirResults;1317  llvm::SmallVector<MlirBlock, 4> mlirSuccessors;1318  llvm::SmallVector<std::pair<std::string, MlirAttribute>, 4> mlirAttributes;1319 1320  // General parameter validation.1321  if (regions < 0)1322    throw nb::value_error("number of regions must be >= 0");1323 1324  // Unpack/validate results.1325  if (results) {1326    mlirResults.reserve(results->size());1327    for (PyType *result : *results) {1328      // TODO: Verify result type originate from the same context.1329      if (!result)1330        throw nb::value_error("result type cannot be None");1331      mlirResults.push_back(*result);1332    }1333  }1334  // Unpack/validate attributes.1335  if (attributes) {1336    mlirAttributes.reserve(attributes->size());1337    for (std::pair<nb::handle, nb::handle> it : *attributes) {1338      std::string key;1339      try {1340        key = nb::cast<std::string>(it.first);1341      } catch (nb::cast_error &err) {1342        std::string msg = "Invalid attribute key (not a string) when "1343                          "attempting to create the operation \"" +1344                          std::string(name) + "\" (" + err.what() + ")";1345        throw nb::type_error(msg.c_str());1346      }1347      try {1348        auto &attribute = nb::cast<PyAttribute &>(it.second);1349        // TODO: Verify attribute originates from the same context.1350        mlirAttributes.emplace_back(std::move(key), attribute);1351      } catch (nb::cast_error &err) {1352        std::string msg = "Invalid attribute value for the key \"" + key +1353                          "\" when attempting to create the operation \"" +1354                          std::string(name) + "\" (" + err.what() + ")";1355        throw nb::type_error(msg.c_str());1356      } catch (std::runtime_error &) {1357        // This exception seems thrown when the value is "None".1358        std::string msg =1359            "Found an invalid (`None`?) attribute value for the key \"" + key +1360            "\" when attempting to create the operation \"" +1361            std::string(name) + "\"";1362        throw std::runtime_error(msg);1363      }1364    }1365  }1366  // Unpack/validate successors.1367  if (successors) {1368    mlirSuccessors.reserve(successors->size());1369    for (auto *successor : *successors) {1370      // TODO: Verify successor originate from the same context.1371      if (!successor)1372        throw nb::value_error("successor block cannot be None");1373      mlirSuccessors.push_back(successor->get());1374    }1375  }1376 1377  // Apply unpacked/validated to the operation state. Beyond this1378  // point, exceptions cannot be thrown or else the state will leak.1379  MlirOperationState state =1380      mlirOperationStateGet(toMlirStringRef(name), location);1381  if (!operands.empty())1382    mlirOperationStateAddOperands(&state, operands.size(), operands.data());1383  state.enableResultTypeInference = inferType;1384  if (!mlirResults.empty())1385    mlirOperationStateAddResults(&state, mlirResults.size(),1386                                 mlirResults.data());1387  if (!mlirAttributes.empty()) {1388    // Note that the attribute names directly reference bytes in1389    // mlirAttributes, so that vector must not be changed from here1390    // on.1391    llvm::SmallVector<MlirNamedAttribute, 4> mlirNamedAttributes;1392    mlirNamedAttributes.reserve(mlirAttributes.size());1393    for (auto &it : mlirAttributes)1394      mlirNamedAttributes.push_back(mlirNamedAttributeGet(1395          mlirIdentifierGet(mlirAttributeGetContext(it.second),1396                            toMlirStringRef(it.first)),1397          it.second));1398    mlirOperationStateAddAttributes(&state, mlirNamedAttributes.size(),1399                                    mlirNamedAttributes.data());1400  }1401  if (!mlirSuccessors.empty())1402    mlirOperationStateAddSuccessors(&state, mlirSuccessors.size(),1403                                    mlirSuccessors.data());1404  if (regions) {1405    llvm::SmallVector<MlirRegion, 4> mlirRegions;1406    mlirRegions.resize(regions);1407    for (int i = 0; i < regions; ++i)1408      mlirRegions[i] = mlirRegionCreate();1409    mlirOperationStateAddOwnedRegions(&state, mlirRegions.size(),1410                                      mlirRegions.data());1411  }1412 1413  // Construct the operation.1414  PyMlirContext::ErrorCapture errors(location.getContext());1415  MlirOperation operation = mlirOperationCreate(&state);1416  if (!operation.ptr)1417    throw MLIRError("Operation creation failed", errors.take());1418  PyOperationRef created =1419      PyOperation::createDetached(location.getContext(), operation);1420  maybeInsertOperation(created, maybeIp);1421 1422  return created.getObject();1423}1424 1425nb::object PyOperation::clone(const nb::object &maybeIp) {1426  MlirOperation clonedOperation = mlirOperationClone(operation);1427  PyOperationRef cloned =1428      PyOperation::createDetached(getContext(), clonedOperation);1429  maybeInsertOperation(cloned, maybeIp);1430 1431  return cloned->createOpView();1432}1433 1434nb::object PyOperation::createOpView() {1435  checkValid();1436  MlirIdentifier ident = mlirOperationGetName(get());1437  MlirStringRef identStr = mlirIdentifierStr(ident);1438  auto operationCls = PyGlobals::get().lookupOperationClass(1439      StringRef(identStr.data, identStr.length));1440  if (operationCls)1441    return PyOpView::constructDerived(*operationCls, getRef().getObject());1442  return nb::cast(PyOpView(getRef().getObject()));1443}1444 1445void PyOperation::erase() {1446  checkValid();1447  setInvalid();1448  mlirOperationDestroy(operation);1449}1450 1451namespace {1452/// CRTP base class for Python MLIR values that subclass Value and should be1453/// castable from it. The value hierarchy is one level deep and is not supposed1454/// to accommodate other levels unless core MLIR changes.1455template <typename DerivedTy>1456class PyConcreteValue : public PyValue {1457public:1458  // Derived classes must define statics for:1459  //   IsAFunctionTy isaFunction1460  //   const char *pyClassName1461  // and redefine bindDerived.1462  using ClassTy = nb::class_<DerivedTy, PyValue>;1463  using IsAFunctionTy = bool (*)(MlirValue);1464 1465  PyConcreteValue() = default;1466  PyConcreteValue(PyOperationRef operationRef, MlirValue value)1467      : PyValue(operationRef, value) {}1468  PyConcreteValue(PyValue &orig)1469      : PyConcreteValue(orig.getParentOperation(), castFrom(orig)) {}1470 1471  /// Attempts to cast the original value to the derived type and throws on1472  /// type mismatches.1473  static MlirValue castFrom(PyValue &orig) {1474    if (!DerivedTy::isaFunction(orig.get())) {1475      auto origRepr = nb::cast<std::string>(nb::repr(nb::cast(orig)));1476      throw nb::value_error((Twine("Cannot cast value to ") +1477                             DerivedTy::pyClassName + " (from " + origRepr +1478                             ")")1479                                .str()1480                                .c_str());1481    }1482    return orig.get();1483  }1484 1485  /// Binds the Python module objects to functions of this class.1486  static void bind(nb::module_ &m) {1487    auto cls = ClassTy(1488        m, DerivedTy::pyClassName, nb::is_generic(),1489        nb::sig((Twine("class ") + DerivedTy::pyClassName + "(Value[_T])")1490                    .str()1491                    .c_str()));1492    cls.def(nb::init<PyValue &>(), nb::keep_alive<0, 1>(), nb::arg("value"));1493    cls.def_static(1494        "isinstance",1495        [](PyValue &otherValue) -> bool {1496          return DerivedTy::isaFunction(otherValue);1497        },1498        nb::arg("other_value"));1499    cls.def(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR,1500            [](DerivedTy &self) -> nb::typed<nb::object, DerivedTy> {1501              return self.maybeDownCast();1502            });1503    DerivedTy::bindDerived(cls);1504  }1505 1506  /// Implemented by derived classes to add methods to the Python subclass.1507  static void bindDerived(ClassTy &m) {}1508};1509 1510} // namespace1511 1512/// Python wrapper for MlirOpResult.1513class PyOpResult : public PyConcreteValue<PyOpResult> {1514public:1515  static constexpr IsAFunctionTy isaFunction = mlirValueIsAOpResult;1516  static constexpr const char *pyClassName = "OpResult";1517  using PyConcreteValue::PyConcreteValue;1518 1519  static void bindDerived(ClassTy &c) {1520    c.def_prop_ro(1521        "owner",1522        [](PyOpResult &self) -> nb::typed<nb::object, PyOperation> {1523          assert(mlirOperationEqual(self.getParentOperation()->get(),1524                                    mlirOpResultGetOwner(self.get())) &&1525                 "expected the owner of the value in Python to match that in "1526                 "the IR");1527          return self.getParentOperation().getObject();1528        },1529        "Returns the operation that produces this result.");1530    c.def_prop_ro(1531        "result_number",1532        [](PyOpResult &self) {1533          return mlirOpResultGetResultNumber(self.get());1534        },1535        "Returns the position of this result in the operation's result list.");1536  }1537};1538 1539/// Returns the list of types of the values held by container.1540template <typename Container>1541static std::vector<nb::typed<nb::object, PyType>>1542getValueTypes(Container &container, PyMlirContextRef &context) {1543  std::vector<nb::typed<nb::object, PyType>> result;1544  result.reserve(container.size());1545  for (int i = 0, e = container.size(); i < e; ++i) {1546    result.push_back(PyType(context->getRef(),1547                            mlirValueGetType(container.getElement(i).get()))1548                         .maybeDownCast());1549  }1550  return result;1551}1552 1553/// A list of operation results. Internally, these are stored as consecutive1554/// elements, random access is cheap. The (returned) result list is associated1555/// with the operation whose results these are, and thus extends the lifetime of1556/// this operation.1557class PyOpResultList : public Sliceable<PyOpResultList, PyOpResult> {1558public:1559  static constexpr const char *pyClassName = "OpResultList";1560  using SliceableT = Sliceable<PyOpResultList, PyOpResult>;1561 1562  PyOpResultList(PyOperationRef operation, intptr_t startIndex = 0,1563                 intptr_t length = -1, intptr_t step = 1)1564      : Sliceable(startIndex,1565                  length == -1 ? mlirOperationGetNumResults(operation->get())1566                               : length,1567                  step),1568        operation(std::move(operation)) {}1569 1570  static void bindDerived(ClassTy &c) {1571    c.def_prop_ro(1572        "types",1573        [](PyOpResultList &self) {1574          return getValueTypes(self, self.operation->getContext());1575        },1576        "Returns a list of types for all results in this result list.");1577    c.def_prop_ro(1578        "owner",1579        [](PyOpResultList &self) -> nb::typed<nb::object, PyOpView> {1580          return self.operation->createOpView();1581        },1582        "Returns the operation that owns this result list.");1583  }1584 1585  PyOperationRef &getOperation() { return operation; }1586 1587private:1588  /// Give the parent CRTP class access to hook implementations below.1589  friend class Sliceable<PyOpResultList, PyOpResult>;1590 1591  intptr_t getRawNumElements() {1592    operation->checkValid();1593    return mlirOperationGetNumResults(operation->get());1594  }1595 1596  PyOpResult getRawElement(intptr_t index) {1597    PyValue value(operation, mlirOperationGetResult(operation->get(), index));1598    return PyOpResult(value);1599  }1600 1601  PyOpResultList slice(intptr_t startIndex, intptr_t length, intptr_t step) {1602    return PyOpResultList(operation, startIndex, length, step);1603  }1604 1605  PyOperationRef operation;1606};1607 1608//------------------------------------------------------------------------------1609// PyOpView1610//------------------------------------------------------------------------------1611 1612static void populateResultTypes(StringRef name, nb::list resultTypeList,1613                                const nb::object &resultSegmentSpecObj,1614                                std::vector<int32_t> &resultSegmentLengths,1615                                std::vector<PyType *> &resultTypes) {1616  resultTypes.reserve(resultTypeList.size());1617  if (resultSegmentSpecObj.is_none()) {1618    // Non-variadic result unpacking.1619    for (const auto &it : llvm::enumerate(resultTypeList)) {1620      try {1621        resultTypes.push_back(nb::cast<PyType *>(it.value()));1622        if (!resultTypes.back())1623          throw nb::cast_error();1624      } catch (nb::cast_error &err) {1625        throw nb::value_error((llvm::Twine("Result ") +1626                               llvm::Twine(it.index()) + " of operation \"" +1627                               name + "\" must be a Type (" + err.what() + ")")1628                                  .str()1629                                  .c_str());1630      }1631    }1632  } else {1633    // Sized result unpacking.1634    auto resultSegmentSpec = nb::cast<std::vector<int>>(resultSegmentSpecObj);1635    if (resultSegmentSpec.size() != resultTypeList.size()) {1636      throw nb::value_error((llvm::Twine("Operation \"") + name +1637                             "\" requires " +1638                             llvm::Twine(resultSegmentSpec.size()) +1639                             " result segments but was provided " +1640                             llvm::Twine(resultTypeList.size()))1641                                .str()1642                                .c_str());1643    }1644    resultSegmentLengths.reserve(resultTypeList.size());1645    for (const auto &it :1646         llvm::enumerate(llvm::zip(resultTypeList, resultSegmentSpec))) {1647      int segmentSpec = std::get<1>(it.value());1648      if (segmentSpec == 1 || segmentSpec == 0) {1649        // Unpack unary element.1650        try {1651          auto *resultType = nb::cast<PyType *>(std::get<0>(it.value()));1652          if (resultType) {1653            resultTypes.push_back(resultType);1654            resultSegmentLengths.push_back(1);1655          } else if (segmentSpec == 0) {1656            // Allowed to be optional.1657            resultSegmentLengths.push_back(0);1658          } else {1659            throw nb::value_error(1660                (llvm::Twine("Result ") + llvm::Twine(it.index()) +1661                 " of operation \"" + name +1662                 "\" must be a Type (was None and result is not optional)")1663                    .str()1664                    .c_str());1665          }1666        } catch (nb::cast_error &err) {1667          throw nb::value_error((llvm::Twine("Result ") +1668                                 llvm::Twine(it.index()) + " of operation \"" +1669                                 name + "\" must be a Type (" + err.what() +1670                                 ")")1671                                    .str()1672                                    .c_str());1673        }1674      } else if (segmentSpec == -1) {1675        // Unpack sequence by appending.1676        try {1677          if (std::get<0>(it.value()).is_none()) {1678            // Treat it as an empty list.1679            resultSegmentLengths.push_back(0);1680          } else {1681            // Unpack the list.1682            auto segment = nb::cast<nb::sequence>(std::get<0>(it.value()));1683            for (nb::handle segmentItem : segment) {1684              resultTypes.push_back(nb::cast<PyType *>(segmentItem));1685              if (!resultTypes.back()) {1686                throw nb::type_error("contained a None item");1687              }1688            }1689            resultSegmentLengths.push_back(nb::len(segment));1690          }1691        } catch (std::exception &err) {1692          // NOTE: Sloppy to be using a catch-all here, but there are at least1693          // three different unrelated exceptions that can be thrown in the1694          // above "casts". Just keep the scope above small and catch them all.1695          throw nb::value_error((llvm::Twine("Result ") +1696                                 llvm::Twine(it.index()) + " of operation \"" +1697                                 name + "\" must be a Sequence of Types (" +1698                                 err.what() + ")")1699                                    .str()1700                                    .c_str());1701        }1702      } else {1703        throw nb::value_error("Unexpected segment spec");1704      }1705    }1706  }1707}1708 1709static MlirValue getUniqueResult(MlirOperation operation) {1710  auto numResults = mlirOperationGetNumResults(operation);1711  if (numResults != 1) {1712    auto name = mlirIdentifierStr(mlirOperationGetName(operation));1713    throw nb::value_error((Twine("Cannot call .result on operation ") +1714                           StringRef(name.data, name.length) + " which has " +1715                           Twine(numResults) +1716                           " results (it is only valid for operations with a "1717                           "single result)")1718                              .str()1719                              .c_str());1720  }1721  return mlirOperationGetResult(operation, 0);1722}1723 1724static MlirValue getOpResultOrValue(nb::handle operand) {1725  if (operand.is_none()) {1726    throw nb::value_error("contained a None item");1727  }1728  PyOperationBase *op;1729  if (nb::try_cast<PyOperationBase *>(operand, op)) {1730    return getUniqueResult(op->getOperation());1731  }1732  PyOpResultList *opResultList;1733  if (nb::try_cast<PyOpResultList *>(operand, opResultList)) {1734    return getUniqueResult(opResultList->getOperation()->get());1735  }1736  PyValue *value;1737  if (nb::try_cast<PyValue *>(operand, value)) {1738    return value->get();1739  }1740  throw nb::value_error("is not a Value");1741}1742 1743nb::object PyOpView::buildGeneric(1744    std::string_view name, std::tuple<int, bool> opRegionSpec,1745    nb::object operandSegmentSpecObj, nb::object resultSegmentSpecObj,1746    std::optional<nb::list> resultTypeList, nb::list operandList,1747    std::optional<nb::dict> attributes,1748    std::optional<std::vector<PyBlock *>> successors,1749    std::optional<int> regions, PyLocation &location,1750    const nb::object &maybeIp) {1751  PyMlirContextRef context = location.getContext();1752 1753  // Class level operation construction metadata.1754  // Operand and result segment specs are either none, which does no1755  // variadic unpacking, or a list of ints with segment sizes, where each1756  // element is either a positive number (typically 1 for a scalar) or -1 to1757  // indicate that it is derived from the length of the same-indexed operand1758  // or result (implying that it is a list at that position).1759  std::vector<int32_t> operandSegmentLengths;1760  std::vector<int32_t> resultSegmentLengths;1761 1762  // Validate/determine region count.1763  int opMinRegionCount = std::get<0>(opRegionSpec);1764  bool opHasNoVariadicRegions = std::get<1>(opRegionSpec);1765  if (!regions) {1766    regions = opMinRegionCount;1767  }1768  if (*regions < opMinRegionCount) {1769    throw nb::value_error(1770        (llvm::Twine("Operation \"") + name + "\" requires a minimum of " +1771         llvm::Twine(opMinRegionCount) +1772         " regions but was built with regions=" + llvm::Twine(*regions))1773            .str()1774            .c_str());1775  }1776  if (opHasNoVariadicRegions && *regions > opMinRegionCount) {1777    throw nb::value_error(1778        (llvm::Twine("Operation \"") + name + "\" requires a maximum of " +1779         llvm::Twine(opMinRegionCount) +1780         " regions but was built with regions=" + llvm::Twine(*regions))1781            .str()1782            .c_str());1783  }1784 1785  // Unpack results.1786  std::vector<PyType *> resultTypes;1787  if (resultTypeList.has_value()) {1788    populateResultTypes(name, *resultTypeList, resultSegmentSpecObj,1789                        resultSegmentLengths, resultTypes);1790  }1791 1792  // Unpack operands.1793  llvm::SmallVector<MlirValue, 4> operands;1794  operands.reserve(operands.size());1795  if (operandSegmentSpecObj.is_none()) {1796    // Non-sized operand unpacking.1797    for (const auto &it : llvm::enumerate(operandList)) {1798      try {1799        operands.push_back(getOpResultOrValue(it.value()));1800      } catch (nb::builtin_exception &err) {1801        throw nb::value_error((llvm::Twine("Operand ") +1802                               llvm::Twine(it.index()) + " of operation \"" +1803                               name + "\" must be a Value (" + err.what() + ")")1804                                  .str()1805                                  .c_str());1806      }1807    }1808  } else {1809    // Sized operand unpacking.1810    auto operandSegmentSpec = nb::cast<std::vector<int>>(operandSegmentSpecObj);1811    if (operandSegmentSpec.size() != operandList.size()) {1812      throw nb::value_error((llvm::Twine("Operation \"") + name +1813                             "\" requires " +1814                             llvm::Twine(operandSegmentSpec.size()) +1815                             "operand segments but was provided " +1816                             llvm::Twine(operandList.size()))1817                                .str()1818                                .c_str());1819    }1820    operandSegmentLengths.reserve(operandList.size());1821    for (const auto &it :1822         llvm::enumerate(llvm::zip(operandList, operandSegmentSpec))) {1823      int segmentSpec = std::get<1>(it.value());1824      if (segmentSpec == 1 || segmentSpec == 0) {1825        // Unpack unary element.1826        auto &operand = std::get<0>(it.value());1827        if (!operand.is_none()) {1828          try {1829 1830            operands.push_back(getOpResultOrValue(operand));1831          } catch (nb::builtin_exception &err) {1832            throw nb::value_error((llvm::Twine("Operand ") +1833                                   llvm::Twine(it.index()) +1834                                   " of operation \"" + name +1835                                   "\" must be a Value (" + err.what() + ")")1836                                      .str()1837                                      .c_str());1838          }1839 1840          operandSegmentLengths.push_back(1);1841        } else if (segmentSpec == 0) {1842          // Allowed to be optional.1843          operandSegmentLengths.push_back(0);1844        } else {1845          throw nb::value_error(1846              (llvm::Twine("Operand ") + llvm::Twine(it.index()) +1847               " of operation \"" + name +1848               "\" must be a Value (was None and operand is not optional)")1849                  .str()1850                  .c_str());1851        }1852      } else if (segmentSpec == -1) {1853        // Unpack sequence by appending.1854        try {1855          if (std::get<0>(it.value()).is_none()) {1856            // Treat it as an empty list.1857            operandSegmentLengths.push_back(0);1858          } else {1859            // Unpack the list.1860            auto segment = nb::cast<nb::sequence>(std::get<0>(it.value()));1861            for (nb::handle segmentItem : segment) {1862              operands.push_back(getOpResultOrValue(segmentItem));1863            }1864            operandSegmentLengths.push_back(nb::len(segment));1865          }1866        } catch (std::exception &err) {1867          // NOTE: Sloppy to be using a catch-all here, but there are at least1868          // three different unrelated exceptions that can be thrown in the1869          // above "casts". Just keep the scope above small and catch them all.1870          throw nb::value_error((llvm::Twine("Operand ") +1871                                 llvm::Twine(it.index()) + " of operation \"" +1872                                 name + "\" must be a Sequence of Values (" +1873                                 err.what() + ")")1874                                    .str()1875                                    .c_str());1876        }1877      } else {1878        throw nb::value_error("Unexpected segment spec");1879      }1880    }1881  }1882 1883  // Merge operand/result segment lengths into attributes if needed.1884  if (!operandSegmentLengths.empty() || !resultSegmentLengths.empty()) {1885    // Dup.1886    if (attributes) {1887      attributes = nb::dict(*attributes);1888    } else {1889      attributes = nb::dict();1890    }1891    if (attributes->contains("resultSegmentSizes") ||1892        attributes->contains("operandSegmentSizes")) {1893      throw nb::value_error("Manually setting a 'resultSegmentSizes' or "1894                            "'operandSegmentSizes' attribute is unsupported. "1895                            "Use Operation.create for such low-level access.");1896    }1897 1898    // Add resultSegmentSizes attribute.1899    if (!resultSegmentLengths.empty()) {1900      MlirAttribute segmentLengthAttr =1901          mlirDenseI32ArrayGet(context->get(), resultSegmentLengths.size(),1902                               resultSegmentLengths.data());1903      (*attributes)["resultSegmentSizes"] =1904          PyAttribute(context, segmentLengthAttr);1905    }1906 1907    // Add operandSegmentSizes attribute.1908    if (!operandSegmentLengths.empty()) {1909      MlirAttribute segmentLengthAttr =1910          mlirDenseI32ArrayGet(context->get(), operandSegmentLengths.size(),1911                               operandSegmentLengths.data());1912      (*attributes)["operandSegmentSizes"] =1913          PyAttribute(context, segmentLengthAttr);1914    }1915  }1916 1917  // Delegate to create.1918  return PyOperation::create(name,1919                             /*results=*/std::move(resultTypes),1920                             /*operands=*/operands,1921                             /*attributes=*/std::move(attributes),1922                             /*successors=*/std::move(successors),1923                             /*regions=*/*regions, location, maybeIp,1924                             !resultTypeList);1925}1926 1927nb::object PyOpView::constructDerived(const nb::object &cls,1928                                      const nb::object &operation) {1929  nb::handle opViewType = nb::type<PyOpView>();1930  nb::object instance = cls.attr("__new__")(cls);1931  opViewType.attr("__init__")(instance, operation);1932  return instance;1933}1934 1935PyOpView::PyOpView(const nb::object &operationObject)1936    // Casting through the PyOperationBase base-class and then back to the1937    // Operation lets us accept any PyOperationBase subclass.1938    : operation(nb::cast<PyOperationBase &>(operationObject).getOperation()),1939      operationObject(operation.getRef().getObject()) {}1940 1941//------------------------------------------------------------------------------1942// PyInsertionPoint.1943//------------------------------------------------------------------------------1944 1945PyInsertionPoint::PyInsertionPoint(const PyBlock &block) : block(block) {}1946 1947PyInsertionPoint::PyInsertionPoint(PyOperationBase &beforeOperationBase)1948    : refOperation(beforeOperationBase.getOperation().getRef()),1949      block((*refOperation)->getBlock()) {}1950 1951PyInsertionPoint::PyInsertionPoint(PyOperationRef beforeOperationRef)1952    : refOperation(beforeOperationRef), block((*refOperation)->getBlock()) {}1953 1954void PyInsertionPoint::insert(PyOperationBase &operationBase) {1955  PyOperation &operation = operationBase.getOperation();1956  if (operation.isAttached())1957    throw nb::value_error(1958        "Attempt to insert operation that is already attached");1959  block.getParentOperation()->checkValid();1960  MlirOperation beforeOp = {nullptr};1961  if (refOperation) {1962    // Insert before operation.1963    (*refOperation)->checkValid();1964    beforeOp = (*refOperation)->get();1965  } else {1966    // Insert at end (before null) is only valid if the block does not1967    // already end in a known terminator (violating this will cause assertion1968    // failures later).1969    if (!mlirOperationIsNull(mlirBlockGetTerminator(block.get()))) {1970      throw nb::index_error("Cannot insert operation at the end of a block "1971                            "that already has a terminator. Did you mean to "1972                            "use 'InsertionPoint.at_block_terminator(block)' "1973                            "versus 'InsertionPoint(block)'?");1974    }1975  }1976  mlirBlockInsertOwnedOperationBefore(block.get(), beforeOp, operation);1977  operation.setAttached();1978}1979 1980PyInsertionPoint PyInsertionPoint::atBlockBegin(PyBlock &block) {1981  MlirOperation firstOp = mlirBlockGetFirstOperation(block.get());1982  if (mlirOperationIsNull(firstOp)) {1983    // Just insert at end.1984    return PyInsertionPoint(block);1985  }1986 1987  // Insert before first op.1988  PyOperationRef firstOpRef = PyOperation::forOperation(1989      block.getParentOperation()->getContext(), firstOp);1990  return PyInsertionPoint{block, std::move(firstOpRef)};1991}1992 1993PyInsertionPoint PyInsertionPoint::atBlockTerminator(PyBlock &block) {1994  MlirOperation terminator = mlirBlockGetTerminator(block.get());1995  if (mlirOperationIsNull(terminator))1996    throw nb::value_error("Block has no terminator");1997  PyOperationRef terminatorOpRef = PyOperation::forOperation(1998      block.getParentOperation()->getContext(), terminator);1999  return PyInsertionPoint{block, std::move(terminatorOpRef)};2000}2001 2002PyInsertionPoint PyInsertionPoint::after(PyOperationBase &op) {2003  PyOperation &operation = op.getOperation();2004  PyBlock block = operation.getBlock();2005  MlirOperation nextOperation = mlirOperationGetNextInBlock(operation);2006  if (mlirOperationIsNull(nextOperation))2007    return PyInsertionPoint(block);2008  PyOperationRef nextOpRef = PyOperation::forOperation(2009      block.getParentOperation()->getContext(), nextOperation);2010  return PyInsertionPoint{block, std::move(nextOpRef)};2011}2012 2013size_t PyMlirContext::getLiveModuleCount() { return liveModules.size(); }2014 2015nb::object PyInsertionPoint::contextEnter(nb::object insertPoint) {2016  return PyThreadContextEntry::pushInsertionPoint(std::move(insertPoint));2017}2018 2019void PyInsertionPoint::contextExit(const nb::object &excType,2020                                   const nb::object &excVal,2021                                   const nb::object &excTb) {2022  PyThreadContextEntry::popInsertionPoint(*this);2023}2024 2025//------------------------------------------------------------------------------2026// PyAttribute.2027//------------------------------------------------------------------------------2028 2029bool PyAttribute::operator==(const PyAttribute &other) const {2030  return mlirAttributeEqual(attr, other.attr);2031}2032 2033nb::object PyAttribute::getCapsule() {2034  return nb::steal<nb::object>(mlirPythonAttributeToCapsule(*this));2035}2036 2037PyAttribute PyAttribute::createFromCapsule(const nb::object &capsule) {2038  MlirAttribute rawAttr = mlirPythonCapsuleToAttribute(capsule.ptr());2039  if (mlirAttributeIsNull(rawAttr))2040    throw nb::python_error();2041  return PyAttribute(2042      PyMlirContext::forContext(mlirAttributeGetContext(rawAttr)), rawAttr);2043}2044 2045nb::object PyAttribute::maybeDownCast() {2046  MlirTypeID mlirTypeID = mlirAttributeGetTypeID(this->get());2047  assert(!mlirTypeIDIsNull(mlirTypeID) &&2048         "mlirTypeID was expected to be non-null.");2049  std::optional<nb::callable> typeCaster = PyGlobals::get().lookupTypeCaster(2050      mlirTypeID, mlirAttributeGetDialect(this->get()));2051  // nb::rv_policy::move means use std::move to move the return value2052  // contents into a new instance that will be owned by Python.2053  nb::object thisObj = nb::cast(this, nb::rv_policy::move);2054  if (!typeCaster)2055    return thisObj;2056  return typeCaster.value()(thisObj);2057}2058 2059//------------------------------------------------------------------------------2060// PyNamedAttribute.2061//------------------------------------------------------------------------------2062 2063PyNamedAttribute::PyNamedAttribute(MlirAttribute attr, std::string ownedName)2064    : ownedName(new std::string(std::move(ownedName))) {2065  namedAttr = mlirNamedAttributeGet(2066      mlirIdentifierGet(mlirAttributeGetContext(attr),2067                        toMlirStringRef(*this->ownedName)),2068      attr);2069}2070 2071//------------------------------------------------------------------------------2072// PyType.2073//------------------------------------------------------------------------------2074 2075bool PyType::operator==(const PyType &other) const {2076  return mlirTypeEqual(type, other.type);2077}2078 2079nb::object PyType::getCapsule() {2080  return nb::steal<nb::object>(mlirPythonTypeToCapsule(*this));2081}2082 2083PyType PyType::createFromCapsule(nb::object capsule) {2084  MlirType rawType = mlirPythonCapsuleToType(capsule.ptr());2085  if (mlirTypeIsNull(rawType))2086    throw nb::python_error();2087  return PyType(PyMlirContext::forContext(mlirTypeGetContext(rawType)),2088                rawType);2089}2090 2091nb::object PyType::maybeDownCast() {2092  MlirTypeID mlirTypeID = mlirTypeGetTypeID(this->get());2093  assert(!mlirTypeIDIsNull(mlirTypeID) &&2094         "mlirTypeID was expected to be non-null.");2095  std::optional<nb::callable> typeCaster = PyGlobals::get().lookupTypeCaster(2096      mlirTypeID, mlirTypeGetDialect(this->get()));2097  // nb::rv_policy::move means use std::move to move the return value2098  // contents into a new instance that will be owned by Python.2099  nb::object thisObj = nb::cast(this, nb::rv_policy::move);2100  if (!typeCaster)2101    return thisObj;2102  return typeCaster.value()(thisObj);2103}2104 2105//------------------------------------------------------------------------------2106// PyTypeID.2107//------------------------------------------------------------------------------2108 2109nb::object PyTypeID::getCapsule() {2110  return nb::steal<nb::object>(mlirPythonTypeIDToCapsule(*this));2111}2112 2113PyTypeID PyTypeID::createFromCapsule(nb::object capsule) {2114  MlirTypeID mlirTypeID = mlirPythonCapsuleToTypeID(capsule.ptr());2115  if (mlirTypeIDIsNull(mlirTypeID))2116    throw nb::python_error();2117  return PyTypeID(mlirTypeID);2118}2119bool PyTypeID::operator==(const PyTypeID &other) const {2120  return mlirTypeIDEqual(typeID, other.typeID);2121}2122 2123//------------------------------------------------------------------------------2124// PyValue and subclasses.2125//------------------------------------------------------------------------------2126 2127nb::object PyValue::getCapsule() {2128  return nb::steal<nb::object>(mlirPythonValueToCapsule(get()));2129}2130 2131nb::object PyValue::maybeDownCast() {2132  MlirType type = mlirValueGetType(get());2133  MlirTypeID mlirTypeID = mlirTypeGetTypeID(type);2134  assert(!mlirTypeIDIsNull(mlirTypeID) &&2135         "mlirTypeID was expected to be non-null.");2136  std::optional<nb::callable> valueCaster =2137      PyGlobals::get().lookupValueCaster(mlirTypeID, mlirTypeGetDialect(type));2138  // nb::rv_policy::move means use std::move to move the return value2139  // contents into a new instance that will be owned by Python.2140  nb::object thisObj = nb::cast(this, nb::rv_policy::move);2141  if (!valueCaster)2142    return thisObj;2143  return valueCaster.value()(thisObj);2144}2145 2146PyValue PyValue::createFromCapsule(nb::object capsule) {2147  MlirValue value = mlirPythonCapsuleToValue(capsule.ptr());2148  if (mlirValueIsNull(value))2149    throw nb::python_error();2150  MlirOperation owner;2151  if (mlirValueIsAOpResult(value))2152    owner = mlirOpResultGetOwner(value);2153  if (mlirValueIsABlockArgument(value))2154    owner = mlirBlockGetParentOperation(mlirBlockArgumentGetOwner(value));2155  if (mlirOperationIsNull(owner))2156    throw nb::python_error();2157  MlirContext ctx = mlirOperationGetContext(owner);2158  PyOperationRef ownerRef =2159      PyOperation::forOperation(PyMlirContext::forContext(ctx), owner);2160  return PyValue(ownerRef, value);2161}2162 2163//------------------------------------------------------------------------------2164// PySymbolTable.2165//------------------------------------------------------------------------------2166 2167PySymbolTable::PySymbolTable(PyOperationBase &operation)2168    : operation(operation.getOperation().getRef()) {2169  symbolTable = mlirSymbolTableCreate(operation.getOperation().get());2170  if (mlirSymbolTableIsNull(symbolTable)) {2171    throw nb::type_error("Operation is not a Symbol Table.");2172  }2173}2174 2175nb::object PySymbolTable::dunderGetItem(const std::string &name) {2176  operation->checkValid();2177  MlirOperation symbol = mlirSymbolTableLookup(2178      symbolTable, mlirStringRefCreate(name.data(), name.length()));2179  if (mlirOperationIsNull(symbol))2180    throw nb::key_error(2181        ("Symbol '" + name + "' not in the symbol table.").c_str());2182 2183  return PyOperation::forOperation(operation->getContext(), symbol,2184                                   operation.getObject())2185      ->createOpView();2186}2187 2188void PySymbolTable::erase(PyOperationBase &symbol) {2189  operation->checkValid();2190  symbol.getOperation().checkValid();2191  mlirSymbolTableErase(symbolTable, symbol.getOperation().get());2192  // The operation is also erased, so we must invalidate it. There may be Python2193  // references to this operation so we don't want to delete it from the list of2194  // live operations here.2195  symbol.getOperation().valid = false;2196}2197 2198void PySymbolTable::dunderDel(const std::string &name) {2199  nb::object operation = dunderGetItem(name);2200  erase(nb::cast<PyOperationBase &>(operation));2201}2202 2203PyStringAttribute PySymbolTable::insert(PyOperationBase &symbol) {2204  operation->checkValid();2205  symbol.getOperation().checkValid();2206  MlirAttribute symbolAttr = mlirOperationGetAttributeByName(2207      symbol.getOperation().get(), mlirSymbolTableGetSymbolAttributeName());2208  if (mlirAttributeIsNull(symbolAttr))2209    throw nb::value_error("Expected operation to have a symbol name.");2210  return PyStringAttribute(2211      symbol.getOperation().getContext(),2212      mlirSymbolTableInsert(symbolTable, symbol.getOperation().get()));2213}2214 2215PyStringAttribute PySymbolTable::getSymbolName(PyOperationBase &symbol) {2216  // Op must already be a symbol.2217  PyOperation &operation = symbol.getOperation();2218  operation.checkValid();2219  MlirStringRef attrName = mlirSymbolTableGetSymbolAttributeName();2220  MlirAttribute existingNameAttr =2221      mlirOperationGetAttributeByName(operation.get(), attrName);2222  if (mlirAttributeIsNull(existingNameAttr))2223    throw nb::value_error("Expected operation to have a symbol name.");2224  return PyStringAttribute(symbol.getOperation().getContext(),2225                           existingNameAttr);2226}2227 2228void PySymbolTable::setSymbolName(PyOperationBase &symbol,2229                                  const std::string &name) {2230  // Op must already be a symbol.2231  PyOperation &operation = symbol.getOperation();2232  operation.checkValid();2233  MlirStringRef attrName = mlirSymbolTableGetSymbolAttributeName();2234  MlirAttribute existingNameAttr =2235      mlirOperationGetAttributeByName(operation.get(), attrName);2236  if (mlirAttributeIsNull(existingNameAttr))2237    throw nb::value_error("Expected operation to have a symbol name.");2238  MlirAttribute newNameAttr =2239      mlirStringAttrGet(operation.getContext()->get(), toMlirStringRef(name));2240  mlirOperationSetAttributeByName(operation.get(), attrName, newNameAttr);2241}2242 2243PyStringAttribute PySymbolTable::getVisibility(PyOperationBase &symbol) {2244  PyOperation &operation = symbol.getOperation();2245  operation.checkValid();2246  MlirStringRef attrName = mlirSymbolTableGetVisibilityAttributeName();2247  MlirAttribute existingVisAttr =2248      mlirOperationGetAttributeByName(operation.get(), attrName);2249  if (mlirAttributeIsNull(existingVisAttr))2250    throw nb::value_error("Expected operation to have a symbol visibility.");2251  return PyStringAttribute(symbol.getOperation().getContext(), existingVisAttr);2252}2253 2254void PySymbolTable::setVisibility(PyOperationBase &symbol,2255                                  const std::string &visibility) {2256  if (visibility != "public" && visibility != "private" &&2257      visibility != "nested")2258    throw nb::value_error(2259        "Expected visibility to be 'public', 'private' or 'nested'");2260  PyOperation &operation = symbol.getOperation();2261  operation.checkValid();2262  MlirStringRef attrName = mlirSymbolTableGetVisibilityAttributeName();2263  MlirAttribute existingVisAttr =2264      mlirOperationGetAttributeByName(operation.get(), attrName);2265  if (mlirAttributeIsNull(existingVisAttr))2266    throw nb::value_error("Expected operation to have a symbol visibility.");2267  MlirAttribute newVisAttr = mlirStringAttrGet(operation.getContext()->get(),2268                                               toMlirStringRef(visibility));2269  mlirOperationSetAttributeByName(operation.get(), attrName, newVisAttr);2270}2271 2272void PySymbolTable::replaceAllSymbolUses(const std::string &oldSymbol,2273                                         const std::string &newSymbol,2274                                         PyOperationBase &from) {2275  PyOperation &fromOperation = from.getOperation();2276  fromOperation.checkValid();2277  if (mlirLogicalResultIsFailure(mlirSymbolTableReplaceAllSymbolUses(2278          toMlirStringRef(oldSymbol), toMlirStringRef(newSymbol),2279          from.getOperation())))2280 2281    throw nb::value_error("Symbol rename failed");2282}2283 2284void PySymbolTable::walkSymbolTables(PyOperationBase &from,2285                                     bool allSymUsesVisible,2286                                     nb::object callback) {2287  PyOperation &fromOperation = from.getOperation();2288  fromOperation.checkValid();2289  struct UserData {2290    PyMlirContextRef context;2291    nb::object callback;2292    bool gotException;2293    std::string exceptionWhat;2294    nb::object exceptionType;2295  };2296  UserData userData{2297      fromOperation.getContext(), std::move(callback), false, {}, {}};2298  mlirSymbolTableWalkSymbolTables(2299      fromOperation.get(), allSymUsesVisible,2300      [](MlirOperation foundOp, bool isVisible, void *calleeUserDataVoid) {2301        UserData *calleeUserData = static_cast<UserData *>(calleeUserDataVoid);2302        auto pyFoundOp =2303            PyOperation::forOperation(calleeUserData->context, foundOp);2304        if (calleeUserData->gotException)2305          return;2306        try {2307          calleeUserData->callback(pyFoundOp.getObject(), isVisible);2308        } catch (nb::python_error &e) {2309          calleeUserData->gotException = true;2310          calleeUserData->exceptionWhat = e.what();2311          calleeUserData->exceptionType = nb::borrow(e.type());2312        }2313      },2314      static_cast<void *>(&userData));2315  if (userData.gotException) {2316    std::string message("Exception raised in callback: ");2317    message.append(userData.exceptionWhat);2318    throw std::runtime_error(message);2319  }2320}2321 2322namespace {2323 2324/// Python wrapper for MlirBlockArgument.2325class PyBlockArgument : public PyConcreteValue<PyBlockArgument> {2326public:2327  static constexpr IsAFunctionTy isaFunction = mlirValueIsABlockArgument;2328  static constexpr const char *pyClassName = "BlockArgument";2329  using PyConcreteValue::PyConcreteValue;2330 2331  static void bindDerived(ClassTy &c) {2332    c.def_prop_ro(2333        "owner",2334        [](PyBlockArgument &self) {2335          return PyBlock(self.getParentOperation(),2336                         mlirBlockArgumentGetOwner(self.get()));2337        },2338        "Returns the block that owns this argument.");2339    c.def_prop_ro(2340        "arg_number",2341        [](PyBlockArgument &self) {2342          return mlirBlockArgumentGetArgNumber(self.get());2343        },2344        "Returns the position of this argument in the block's argument list.");2345    c.def(2346        "set_type",2347        [](PyBlockArgument &self, PyType type) {2348          return mlirBlockArgumentSetType(self.get(), type);2349        },2350        nb::arg("type"), "Sets the type of this block argument.");2351    c.def(2352        "set_location",2353        [](PyBlockArgument &self, PyLocation loc) {2354          return mlirBlockArgumentSetLocation(self.get(), loc);2355        },2356        nb::arg("loc"), "Sets the location of this block argument.");2357  }2358};2359 2360/// A list of block arguments. Internally, these are stored as consecutive2361/// elements, random access is cheap. The argument list is associated with the2362/// operation that contains the block (detached blocks are not allowed in2363/// Python bindings) and extends its lifetime.2364class PyBlockArgumentList2365    : public Sliceable<PyBlockArgumentList, PyBlockArgument> {2366public:2367  static constexpr const char *pyClassName = "BlockArgumentList";2368  using SliceableT = Sliceable<PyBlockArgumentList, PyBlockArgument>;2369 2370  PyBlockArgumentList(PyOperationRef operation, MlirBlock block,2371                      intptr_t startIndex = 0, intptr_t length = -1,2372                      intptr_t step = 1)2373      : Sliceable(startIndex,2374                  length == -1 ? mlirBlockGetNumArguments(block) : length,2375                  step),2376        operation(std::move(operation)), block(block) {}2377 2378  static void bindDerived(ClassTy &c) {2379    c.def_prop_ro(2380        "types",2381        [](PyBlockArgumentList &self) {2382          return getValueTypes(self, self.operation->getContext());2383        },2384        "Returns a list of types for all arguments in this argument list.");2385  }2386 2387private:2388  /// Give the parent CRTP class access to hook implementations below.2389  friend class Sliceable<PyBlockArgumentList, PyBlockArgument>;2390 2391  /// Returns the number of arguments in the list.2392  intptr_t getRawNumElements() {2393    operation->checkValid();2394    return mlirBlockGetNumArguments(block);2395  }2396 2397  /// Returns `pos`-the element in the list.2398  PyBlockArgument getRawElement(intptr_t pos) {2399    MlirValue argument = mlirBlockGetArgument(block, pos);2400    return PyBlockArgument(operation, argument);2401  }2402 2403  /// Returns a sublist of this list.2404  PyBlockArgumentList slice(intptr_t startIndex, intptr_t length,2405                            intptr_t step) {2406    return PyBlockArgumentList(operation, block, startIndex, length, step);2407  }2408 2409  PyOperationRef operation;2410  MlirBlock block;2411};2412 2413/// A list of operation operands. Internally, these are stored as consecutive2414/// elements, random access is cheap. The (returned) operand list is associated2415/// with the operation whose operands these are, and thus extends the lifetime2416/// of this operation.2417class PyOpOperandList : public Sliceable<PyOpOperandList, PyValue> {2418public:2419  static constexpr const char *pyClassName = "OpOperandList";2420  using SliceableT = Sliceable<PyOpOperandList, PyValue>;2421 2422  PyOpOperandList(PyOperationRef operation, intptr_t startIndex = 0,2423                  intptr_t length = -1, intptr_t step = 1)2424      : Sliceable(startIndex,2425                  length == -1 ? mlirOperationGetNumOperands(operation->get())2426                               : length,2427                  step),2428        operation(operation) {}2429 2430  void dunderSetItem(intptr_t index, PyValue value) {2431    index = wrapIndex(index);2432    mlirOperationSetOperand(operation->get(), index, value.get());2433  }2434 2435  static void bindDerived(ClassTy &c) {2436    c.def("__setitem__", &PyOpOperandList::dunderSetItem, nb::arg("index"),2437          nb::arg("value"),2438          "Sets the operand at the specified index to a new value.");2439  }2440 2441private:2442  /// Give the parent CRTP class access to hook implementations below.2443  friend class Sliceable<PyOpOperandList, PyValue>;2444 2445  intptr_t getRawNumElements() {2446    operation->checkValid();2447    return mlirOperationGetNumOperands(operation->get());2448  }2449 2450  PyValue getRawElement(intptr_t pos) {2451    MlirValue operand = mlirOperationGetOperand(operation->get(), pos);2452    MlirOperation owner;2453    if (mlirValueIsAOpResult(operand))2454      owner = mlirOpResultGetOwner(operand);2455    else if (mlirValueIsABlockArgument(operand))2456      owner = mlirBlockGetParentOperation(mlirBlockArgumentGetOwner(operand));2457    else2458      assert(false && "Value must be an block arg or op result.");2459    PyOperationRef pyOwner =2460        PyOperation::forOperation(operation->getContext(), owner);2461    return PyValue(pyOwner, operand);2462  }2463 2464  PyOpOperandList slice(intptr_t startIndex, intptr_t length, intptr_t step) {2465    return PyOpOperandList(operation, startIndex, length, step);2466  }2467 2468  PyOperationRef operation;2469};2470 2471/// A list of operation successors. Internally, these are stored as consecutive2472/// elements, random access is cheap. The (returned) successor list is2473/// associated with the operation whose successors these are, and thus extends2474/// the lifetime of this operation.2475class PyOpSuccessors : public Sliceable<PyOpSuccessors, PyBlock> {2476public:2477  static constexpr const char *pyClassName = "OpSuccessors";2478 2479  PyOpSuccessors(PyOperationRef operation, intptr_t startIndex = 0,2480                 intptr_t length = -1, intptr_t step = 1)2481      : Sliceable(startIndex,2482                  length == -1 ? mlirOperationGetNumSuccessors(operation->get())2483                               : length,2484                  step),2485        operation(operation) {}2486 2487  void dunderSetItem(intptr_t index, PyBlock block) {2488    index = wrapIndex(index);2489    mlirOperationSetSuccessor(operation->get(), index, block.get());2490  }2491 2492  static void bindDerived(ClassTy &c) {2493    c.def("__setitem__", &PyOpSuccessors::dunderSetItem, nb::arg("index"),2494          nb::arg("block"), "Sets the successor block at the specified index.");2495  }2496 2497private:2498  /// Give the parent CRTP class access to hook implementations below.2499  friend class Sliceable<PyOpSuccessors, PyBlock>;2500 2501  intptr_t getRawNumElements() {2502    operation->checkValid();2503    return mlirOperationGetNumSuccessors(operation->get());2504  }2505 2506  PyBlock getRawElement(intptr_t pos) {2507    MlirBlock block = mlirOperationGetSuccessor(operation->get(), pos);2508    return PyBlock(operation, block);2509  }2510 2511  PyOpSuccessors slice(intptr_t startIndex, intptr_t length, intptr_t step) {2512    return PyOpSuccessors(operation, startIndex, length, step);2513  }2514 2515  PyOperationRef operation;2516};2517 2518/// A list of block successors. Internally, these are stored as consecutive2519/// elements, random access is cheap. The (returned) successor list is2520/// associated with the operation and block whose successors these are, and thus2521/// extends the lifetime of this operation and block.2522class PyBlockSuccessors : public Sliceable<PyBlockSuccessors, PyBlock> {2523public:2524  static constexpr const char *pyClassName = "BlockSuccessors";2525 2526  PyBlockSuccessors(PyBlock block, PyOperationRef operation,2527                    intptr_t startIndex = 0, intptr_t length = -1,2528                    intptr_t step = 1)2529      : Sliceable(startIndex,2530                  length == -1 ? mlirBlockGetNumSuccessors(block.get())2531                               : length,2532                  step),2533        operation(operation), block(block) {}2534 2535private:2536  /// Give the parent CRTP class access to hook implementations below.2537  friend class Sliceable<PyBlockSuccessors, PyBlock>;2538 2539  intptr_t getRawNumElements() {2540    block.checkValid();2541    return mlirBlockGetNumSuccessors(block.get());2542  }2543 2544  PyBlock getRawElement(intptr_t pos) {2545    MlirBlock block = mlirBlockGetSuccessor(this->block.get(), pos);2546    return PyBlock(operation, block);2547  }2548 2549  PyBlockSuccessors slice(intptr_t startIndex, intptr_t length, intptr_t step) {2550    return PyBlockSuccessors(block, operation, startIndex, length, step);2551  }2552 2553  PyOperationRef operation;2554  PyBlock block;2555};2556 2557/// A list of block predecessors. The (returned) predecessor list is2558/// associated with the operation and block whose predecessors these are, and2559/// thus extends the lifetime of this operation and block.2560///2561/// WARNING: This Sliceable is more expensive than the others here because2562/// mlirBlockGetPredecessor actually iterates the use-def chain (of block2563/// operands) anew for each indexed access.2564class PyBlockPredecessors : public Sliceable<PyBlockPredecessors, PyBlock> {2565public:2566  static constexpr const char *pyClassName = "BlockPredecessors";2567 2568  PyBlockPredecessors(PyBlock block, PyOperationRef operation,2569                      intptr_t startIndex = 0, intptr_t length = -1,2570                      intptr_t step = 1)2571      : Sliceable(startIndex,2572                  length == -1 ? mlirBlockGetNumPredecessors(block.get())2573                               : length,2574                  step),2575        operation(operation), block(block) {}2576 2577private:2578  /// Give the parent CRTP class access to hook implementations below.2579  friend class Sliceable<PyBlockPredecessors, PyBlock>;2580 2581  intptr_t getRawNumElements() {2582    block.checkValid();2583    return mlirBlockGetNumPredecessors(block.get());2584  }2585 2586  PyBlock getRawElement(intptr_t pos) {2587    MlirBlock block = mlirBlockGetPredecessor(this->block.get(), pos);2588    return PyBlock(operation, block);2589  }2590 2591  PyBlockPredecessors slice(intptr_t startIndex, intptr_t length,2592                            intptr_t step) {2593    return PyBlockPredecessors(block, operation, startIndex, length, step);2594  }2595 2596  PyOperationRef operation;2597  PyBlock block;2598};2599 2600/// A list of operation attributes. Can be indexed by name, producing2601/// attributes, or by index, producing named attributes.2602class PyOpAttributeMap {2603public:2604  PyOpAttributeMap(PyOperationRef operation)2605      : operation(std::move(operation)) {}2606 2607  nb::typed<nb::object, PyAttribute>2608  dunderGetItemNamed(const std::string &name) {2609    MlirAttribute attr = mlirOperationGetAttributeByName(operation->get(),2610                                                         toMlirStringRef(name));2611    if (mlirAttributeIsNull(attr)) {2612      throw nb::key_error("attempt to access a non-existent attribute");2613    }2614    return PyAttribute(operation->getContext(), attr).maybeDownCast();2615  }2616 2617  PyNamedAttribute dunderGetItemIndexed(intptr_t index) {2618    if (index < 0) {2619      index += dunderLen();2620    }2621    if (index < 0 || index >= dunderLen()) {2622      throw nb::index_error("attempt to access out of bounds attribute");2623    }2624    MlirNamedAttribute namedAttr =2625        mlirOperationGetAttribute(operation->get(), index);2626    return PyNamedAttribute(2627        namedAttr.attribute,2628        std::string(mlirIdentifierStr(namedAttr.name).data,2629                    mlirIdentifierStr(namedAttr.name).length));2630  }2631 2632  void dunderSetItem(const std::string &name, const PyAttribute &attr) {2633    mlirOperationSetAttributeByName(operation->get(), toMlirStringRef(name),2634                                    attr);2635  }2636 2637  void dunderDelItem(const std::string &name) {2638    int removed = mlirOperationRemoveAttributeByName(operation->get(),2639                                                     toMlirStringRef(name));2640    if (!removed)2641      throw nb::key_error("attempt to delete a non-existent attribute");2642  }2643 2644  intptr_t dunderLen() {2645    return mlirOperationGetNumAttributes(operation->get());2646  }2647 2648  bool dunderContains(const std::string &name) {2649    return !mlirAttributeIsNull(mlirOperationGetAttributeByName(2650        operation->get(), toMlirStringRef(name)));2651  }2652 2653  static void2654  forEachAttr(MlirOperation op,2655              llvm::function_ref<void(MlirStringRef, MlirAttribute)> fn) {2656    intptr_t n = mlirOperationGetNumAttributes(op);2657    for (intptr_t i = 0; i < n; ++i) {2658      MlirNamedAttribute na = mlirOperationGetAttribute(op, i);2659      MlirStringRef name = mlirIdentifierStr(na.name);2660      fn(name, na.attribute);2661    }2662  }2663 2664  static void bind(nb::module_ &m) {2665    nb::class_<PyOpAttributeMap>(m, "OpAttributeMap")2666        .def("__contains__", &PyOpAttributeMap::dunderContains, nb::arg("name"),2667             "Checks if an attribute with the given name exists in the map.")2668        .def("__len__", &PyOpAttributeMap::dunderLen,2669             "Returns the number of attributes in the map.")2670        .def("__getitem__", &PyOpAttributeMap::dunderGetItemNamed,2671             nb::arg("name"), "Gets an attribute by name.")2672        .def("__getitem__", &PyOpAttributeMap::dunderGetItemIndexed,2673             nb::arg("index"), "Gets a named attribute by index.")2674        .def("__setitem__", &PyOpAttributeMap::dunderSetItem, nb::arg("name"),2675             nb::arg("attr"), "Sets an attribute with the given name.")2676        .def("__delitem__", &PyOpAttributeMap::dunderDelItem, nb::arg("name"),2677             "Deletes an attribute with the given name.")2678        .def(2679            "__iter__",2680            [](PyOpAttributeMap &self) {2681              nb::list keys;2682              PyOpAttributeMap::forEachAttr(2683                  self.operation->get(),2684                  [&](MlirStringRef name, MlirAttribute) {2685                    keys.append(nb::str(name.data, name.length));2686                  });2687              return nb::iter(keys);2688            },2689            "Iterates over attribute names.")2690        .def(2691            "keys",2692            [](PyOpAttributeMap &self) {2693              nb::list out;2694              PyOpAttributeMap::forEachAttr(2695                  self.operation->get(),2696                  [&](MlirStringRef name, MlirAttribute) {2697                    out.append(nb::str(name.data, name.length));2698                  });2699              return out;2700            },2701            "Returns a list of attribute names.")2702        .def(2703            "values",2704            [](PyOpAttributeMap &self) {2705              nb::list out;2706              PyOpAttributeMap::forEachAttr(2707                  self.operation->get(),2708                  [&](MlirStringRef, MlirAttribute attr) {2709                    out.append(PyAttribute(self.operation->getContext(), attr)2710                                   .maybeDownCast());2711                  });2712              return out;2713            },2714            "Returns a list of attribute values.")2715        .def(2716            "items",2717            [](PyOpAttributeMap &self) {2718              nb::list out;2719              PyOpAttributeMap::forEachAttr(2720                  self.operation->get(),2721                  [&](MlirStringRef name, MlirAttribute attr) {2722                    out.append(nb::make_tuple(2723                        nb::str(name.data, name.length),2724                        PyAttribute(self.operation->getContext(), attr)2725                            .maybeDownCast()));2726                  });2727              return out;2728            },2729            "Returns a list of `(name, attribute)` tuples.");2730  }2731 2732private:2733  PyOperationRef operation;2734};2735 2736// see2737// https://raw.githubusercontent.com/python/pythoncapi_compat/master/pythoncapi_compat.h2738 2739#ifndef _Py_CAST2740#define _Py_CAST(type, expr) ((type)(expr))2741#endif2742 2743// Static inline functions should use _Py_NULL rather than using directly NULL2744// to prevent C++ compiler warnings. On C23 and newer and on C++11 and newer,2745// _Py_NULL is defined as nullptr.2746#ifndef _Py_NULL2747#if (defined(__STDC_VERSION__) && __STDC_VERSION__ > 201710L) ||               \2748    (defined(__cplusplus) && __cplusplus >= 201103)2749#define _Py_NULL nullptr2750#else2751#define _Py_NULL NULL2752#endif2753#endif2754 2755// Python 3.10.0a32756#if PY_VERSION_HEX < 0x030A00A32757 2758// bpo-42262 added Py_XNewRef()2759#if !defined(Py_XNewRef)2760[[maybe_unused]] PyObject *_Py_XNewRef(PyObject *obj) {2761  Py_XINCREF(obj);2762  return obj;2763}2764#define Py_XNewRef(obj) _Py_XNewRef(_PyObject_CAST(obj))2765#endif2766 2767// bpo-42262 added Py_NewRef()2768#if !defined(Py_NewRef)2769[[maybe_unused]] PyObject *_Py_NewRef(PyObject *obj) {2770  Py_INCREF(obj);2771  return obj;2772}2773#define Py_NewRef(obj) _Py_NewRef(_PyObject_CAST(obj))2774#endif2775 2776#endif // Python 3.10.0a32777 2778// Python 3.9.0b12779#if PY_VERSION_HEX < 0x030900B1 && !defined(PYPY_VERSION)2780 2781// bpo-40429 added PyThreadState_GetFrame()2782PyFrameObject *PyThreadState_GetFrame(PyThreadState *tstate) {2783  assert(tstate != _Py_NULL && "expected tstate != _Py_NULL");2784  return _Py_CAST(PyFrameObject *, Py_XNewRef(tstate->frame));2785}2786 2787// bpo-40421 added PyFrame_GetBack()2788PyFrameObject *PyFrame_GetBack(PyFrameObject *frame) {2789  assert(frame != _Py_NULL && "expected frame != _Py_NULL");2790  return _Py_CAST(PyFrameObject *, Py_XNewRef(frame->f_back));2791}2792 2793// bpo-40421 added PyFrame_GetCode()2794PyCodeObject *PyFrame_GetCode(PyFrameObject *frame) {2795  assert(frame != _Py_NULL && "expected frame != _Py_NULL");2796  assert(frame->f_code != _Py_NULL && "expected frame->f_code != _Py_NULL");2797  return _Py_CAST(PyCodeObject *, Py_NewRef(frame->f_code));2798}2799 2800#endif // Python 3.9.0b12801 2802MlirLocation tracebackToLocation(MlirContext ctx) {2803  size_t framesLimit =2804      PyGlobals::get().getTracebackLoc().locTracebackFramesLimit();2805  // Use a thread_local here to avoid requiring a large amount of space.2806  thread_local std::array<MlirLocation, PyGlobals::TracebackLoc::kMaxFrames>2807      frames;2808  size_t count = 0;2809 2810  nb::gil_scoped_acquire acquire;2811  PyThreadState *tstate = PyThreadState_GET();2812  PyFrameObject *next;2813  PyFrameObject *pyFrame = PyThreadState_GetFrame(tstate);2814  // In the increment expression:2815  // 1. get the next prev frame;2816  // 2. decrement the ref count on the current frame (in order that it can get2817  //    gc'd, along with any objects in its closure and etc);2818  // 3. set current = next.2819  for (; pyFrame != nullptr && count < framesLimit;2820       next = PyFrame_GetBack(pyFrame), Py_XDECREF(pyFrame), pyFrame = next) {2821    PyCodeObject *code = PyFrame_GetCode(pyFrame);2822    auto fileNameStr =2823        nb::cast<std::string>(nb::borrow<nb::str>(code->co_filename));2824    llvm::StringRef fileName(fileNameStr);2825    if (!PyGlobals::get().getTracebackLoc().isUserTracebackFilename(fileName))2826      continue;2827 2828    // co_qualname and PyCode_Addr2Location added in py3.112829#if PY_VERSION_HEX < 0x030B00F02830    std::string name =2831        nb::cast<std::string>(nb::borrow<nb::str>(code->co_name));2832    llvm::StringRef funcName(name);2833    int startLine = PyFrame_GetLineNumber(pyFrame);2834    MlirLocation loc =2835        mlirLocationFileLineColGet(ctx, wrap(fileName), startLine, 0);2836#else2837    std::string name =2838        nb::cast<std::string>(nb::borrow<nb::str>(code->co_qualname));2839    llvm::StringRef funcName(name);2840    int startLine, startCol, endLine, endCol;2841    int lasti = PyFrame_GetLasti(pyFrame);2842    if (!PyCode_Addr2Location(code, lasti, &startLine, &startCol, &endLine,2843                              &endCol)) {2844      throw nb::python_error();2845    }2846    MlirLocation loc = mlirLocationFileLineColRangeGet(2847        ctx, wrap(fileName), startLine, startCol, endLine, endCol);2848#endif2849 2850    frames[count] = mlirLocationNameGet(ctx, wrap(funcName), loc);2851    ++count;2852  }2853  // When the loop breaks (after the last iter), current frame (if non-null)2854  // is leaked without this.2855  Py_XDECREF(pyFrame);2856 2857  if (count == 0)2858    return mlirLocationUnknownGet(ctx);2859 2860  MlirLocation callee = frames[0];2861  assert(!mlirLocationIsNull(callee) && "expected non-null callee location");2862  if (count == 1)2863    return callee;2864 2865  MlirLocation caller = frames[count - 1];2866  assert(!mlirLocationIsNull(caller) && "expected non-null caller location");2867  for (int i = count - 2; i >= 1; i--)2868    caller = mlirLocationCallSiteGet(frames[i], caller);2869 2870  return mlirLocationCallSiteGet(callee, caller);2871}2872 2873PyLocation2874maybeGetTracebackLocation(const std::optional<PyLocation> &location) {2875  if (location.has_value())2876    return location.value();2877  if (!PyGlobals::get().getTracebackLoc().locTracebacksEnabled())2878    return DefaultingPyLocation::resolve();2879 2880  PyMlirContext &ctx = DefaultingPyMlirContext::resolve();2881  MlirLocation mlirLoc = tracebackToLocation(ctx.get());2882  PyMlirContextRef ref = PyMlirContext::forContext(ctx.get());2883  return {ref, mlirLoc};2884}2885 2886} // namespace2887 2888//------------------------------------------------------------------------------2889// Populates the core exports of the 'ir' submodule.2890//------------------------------------------------------------------------------2891 2892void mlir::python::populateIRCore(nb::module_ &m) {2893  // disable leak warnings which tend to be false positives.2894  nb::set_leak_warnings(false);2895  //----------------------------------------------------------------------------2896  // Enums.2897  //----------------------------------------------------------------------------2898  nb::enum_<MlirDiagnosticSeverity>(m, "DiagnosticSeverity")2899      .value("ERROR", MlirDiagnosticError)2900      .value("WARNING", MlirDiagnosticWarning)2901      .value("NOTE", MlirDiagnosticNote)2902      .value("REMARK", MlirDiagnosticRemark);2903 2904  nb::enum_<MlirWalkOrder>(m, "WalkOrder")2905      .value("PRE_ORDER", MlirWalkPreOrder)2906      .value("POST_ORDER", MlirWalkPostOrder);2907 2908  nb::enum_<MlirWalkResult>(m, "WalkResult")2909      .value("ADVANCE", MlirWalkResultAdvance)2910      .value("INTERRUPT", MlirWalkResultInterrupt)2911      .value("SKIP", MlirWalkResultSkip);2912 2913  //----------------------------------------------------------------------------2914  // Mapping of Diagnostics.2915  //----------------------------------------------------------------------------2916  nb::class_<PyDiagnostic>(m, "Diagnostic")2917      .def_prop_ro("severity", &PyDiagnostic::getSeverity,2918                   "Returns the severity of the diagnostic.")2919      .def_prop_ro("location", &PyDiagnostic::getLocation,2920                   "Returns the location associated with the diagnostic.")2921      .def_prop_ro("message", &PyDiagnostic::getMessage,2922                   "Returns the message text of the diagnostic.")2923      .def_prop_ro("notes", &PyDiagnostic::getNotes,2924                   "Returns a tuple of attached note diagnostics.")2925      .def(2926          "__str__",2927          [](PyDiagnostic &self) -> nb::str {2928            if (!self.isValid())2929              return nb::str("<Invalid Diagnostic>");2930            return self.getMessage();2931          },2932          "Returns the diagnostic message as a string.");2933 2934  nb::class_<PyDiagnostic::DiagnosticInfo>(m, "DiagnosticInfo")2935      .def(2936          "__init__",2937          [](PyDiagnostic::DiagnosticInfo &self, PyDiagnostic diag) {2938            new (&self) PyDiagnostic::DiagnosticInfo(diag.getInfo());2939          },2940          "diag"_a, "Creates a DiagnosticInfo from a Diagnostic.")2941      .def_ro("severity", &PyDiagnostic::DiagnosticInfo::severity,2942              "The severity level of the diagnostic.")2943      .def_ro("location", &PyDiagnostic::DiagnosticInfo::location,2944              "The location associated with the diagnostic.")2945      .def_ro("message", &PyDiagnostic::DiagnosticInfo::message,2946              "The message text of the diagnostic.")2947      .def_ro("notes", &PyDiagnostic::DiagnosticInfo::notes,2948              "List of attached note diagnostics.")2949      .def(2950          "__str__",2951          [](PyDiagnostic::DiagnosticInfo &self) { return self.message; },2952          "Returns the diagnostic message as a string.");2953 2954  nb::class_<PyDiagnosticHandler>(m, "DiagnosticHandler")2955      .def("detach", &PyDiagnosticHandler::detach,2956           "Detaches the diagnostic handler from the context.")2957      .def_prop_ro("attached", &PyDiagnosticHandler::isAttached,2958                   "Returns True if the handler is attached to a context.")2959      .def_prop_ro("had_error", &PyDiagnosticHandler::getHadError,2960                   "Returns True if an error was encountered during diagnostic "2961                   "handling.")2962      .def("__enter__", &PyDiagnosticHandler::contextEnter,2963           "Enters the diagnostic handler as a context manager.")2964      .def("__exit__", &PyDiagnosticHandler::contextExit,2965           nb::arg("exc_type").none(), nb::arg("exc_value").none(),2966           nb::arg("traceback").none(),2967           "Exits the diagnostic handler context manager.");2968 2969  // Expose DefaultThreadPool to python2970  nb::class_<PyThreadPool>(m, "ThreadPool")2971      .def(2972          "__init__", [](PyThreadPool &self) { new (&self) PyThreadPool(); },2973          "Creates a new thread pool with default concurrency.")2974      .def("get_max_concurrency", &PyThreadPool::getMaxConcurrency,2975           "Returns the maximum number of threads in the pool.")2976      .def("_mlir_thread_pool_ptr", &PyThreadPool::_mlir_thread_pool_ptr,2977           "Returns the raw pointer to the LLVM thread pool as a string.");2978 2979  nb::class_<PyMlirContext>(m, "Context")2980      .def(2981          "__init__",2982          [](PyMlirContext &self) {2983            MlirContext context = mlirContextCreateWithThreading(false);2984            new (&self) PyMlirContext(context);2985          },2986          R"(2987            Creates a new MLIR context.2988 2989            The context is the top-level container for all MLIR objects. It owns the storage2990            for types, attributes, locations, and other core IR objects. A context can be2991            configured to allow or disallow unregistered dialects and can have dialects2992            loaded on-demand.)")2993      .def_static("_get_live_count", &PyMlirContext::getLiveCount,2994                  "Gets the number of live Context objects.")2995      .def(2996          "_get_context_again",2997          [](PyMlirContext &self) -> nb::typed<nb::object, PyMlirContext> {2998            PyMlirContextRef ref = PyMlirContext::forContext(self.get());2999            return ref.releaseObject();3000          },3001          "Gets another reference to the same context.")3002      .def("_get_live_module_count", &PyMlirContext::getLiveModuleCount,3003           "Gets the number of live modules owned by this context.")3004      .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyMlirContext::getCapsule,3005                   "Gets a capsule wrapping the MlirContext.")3006      .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR,3007                  &PyMlirContext::createFromCapsule,3008                  "Creates a Context from a capsule wrapping MlirContext.")3009      .def("__enter__", &PyMlirContext::contextEnter,3010           "Enters the context as a context manager.")3011      .def("__exit__", &PyMlirContext::contextExit, nb::arg("exc_type").none(),3012           nb::arg("exc_value").none(), nb::arg("traceback").none(),3013           "Exits the context manager.")3014      .def_prop_ro_static(3015          "current",3016          [](nb::object & /*class*/)3017              -> std::optional<nb::typed<nb::object, PyMlirContext>> {3018            auto *context = PyThreadContextEntry::getDefaultContext();3019            if (!context)3020              return {};3021            return nb::cast(context);3022          },3023          nb::sig("def current(/) -> Context | None"),3024          "Gets the Context bound to the current thread or returns None if no "3025          "context is set.")3026      .def_prop_ro(3027          "dialects",3028          [](PyMlirContext &self) { return PyDialects(self.getRef()); },3029          "Gets a container for accessing dialects by name.")3030      .def_prop_ro(3031          "d", [](PyMlirContext &self) { return PyDialects(self.getRef()); },3032          "Alias for `dialects`.")3033      .def(3034          "get_dialect_descriptor",3035          [=](PyMlirContext &self, std::string &name) {3036            MlirDialect dialect = mlirContextGetOrLoadDialect(3037                self.get(), {name.data(), name.size()});3038            if (mlirDialectIsNull(dialect)) {3039              throw nb::value_error(3040                  (Twine("Dialect '") + name + "' not found").str().c_str());3041            }3042            return PyDialectDescriptor(self.getRef(), dialect);3043          },3044          nb::arg("dialect_name"),3045          "Gets or loads a dialect by name, returning its descriptor object.")3046      .def_prop_rw(3047          "allow_unregistered_dialects",3048          [](PyMlirContext &self) -> bool {3049            return mlirContextGetAllowUnregisteredDialects(self.get());3050          },3051          [](PyMlirContext &self, bool value) {3052            mlirContextSetAllowUnregisteredDialects(self.get(), value);3053          },3054          "Controls whether unregistered dialects are allowed in this context.")3055      .def("attach_diagnostic_handler", &PyMlirContext::attachDiagnosticHandler,3056           nb::arg("callback"),3057           "Attaches a diagnostic handler that will receive callbacks.")3058      .def(3059          "enable_multithreading",3060          [](PyMlirContext &self, bool enable) {3061            mlirContextEnableMultithreading(self.get(), enable);3062          },3063          nb::arg("enable"),3064          R"(3065            Enables or disables multi-threading support in the context.3066 3067            Args:3068              enable: Whether to enable (True) or disable (False) multi-threading.3069          )")3070      .def(3071          "set_thread_pool",3072          [](PyMlirContext &self, PyThreadPool &pool) {3073            // we should disable multi-threading first before setting3074            // new thread pool otherwise the assert in3075            // MLIRContext::setThreadPool will be raised.3076            mlirContextEnableMultithreading(self.get(), false);3077            mlirContextSetThreadPool(self.get(), pool.get());3078          },3079          R"(3080            Sets a custom thread pool for the context to use.3081 3082            Args:3083              pool: A ThreadPool object to use for parallel operations.3084 3085            Note:3086              Multi-threading is automatically disabled before setting the thread pool.)")3087      .def(3088          "get_num_threads",3089          [](PyMlirContext &self) {3090            return mlirContextGetNumThreads(self.get());3091          },3092          "Gets the number of threads in the context's thread pool.")3093      .def(3094          "_mlir_thread_pool_ptr",3095          [](PyMlirContext &self) {3096            MlirLlvmThreadPool pool = mlirContextGetThreadPool(self.get());3097            std::stringstream ss;3098            ss << pool.ptr;3099            return ss.str();3100          },3101          "Gets the raw pointer to the LLVM thread pool as a string.")3102      .def(3103          "is_registered_operation",3104          [](PyMlirContext &self, std::string &name) {3105            return mlirContextIsRegisteredOperation(3106                self.get(), MlirStringRef{name.data(), name.size()});3107          },3108          nb::arg("operation_name"),3109          R"(3110            Checks whether an operation with the given name is registered.3111 3112            Args:3113              operation_name: The fully qualified name of the operation (e.g., `arith.addf`).3114 3115            Returns:3116              True if the operation is registered, False otherwise.)")3117      .def(3118          "append_dialect_registry",3119          [](PyMlirContext &self, PyDialectRegistry &registry) {3120            mlirContextAppendDialectRegistry(self.get(), registry);3121          },3122          nb::arg("registry"),3123          R"(3124            Appends the contents of a dialect registry to the context.3125 3126            Args:3127              registry: A DialectRegistry containing dialects to append.)")3128      .def_prop_rw("emit_error_diagnostics",3129                   &PyMlirContext::getEmitErrorDiagnostics,3130                   &PyMlirContext::setEmitErrorDiagnostics,3131                   R"(3132            Controls whether error diagnostics are emitted to diagnostic handlers.3133 3134            By default, error diagnostics are captured and reported through MLIRError exceptions.)")3135      .def(3136          "load_all_available_dialects",3137          [](PyMlirContext &self) {3138            mlirContextLoadAllAvailableDialects(self.get());3139          },3140          R"(3141            Loads all dialects available in the registry into the context.3142 3143            This eagerly loads all dialects that have been registered, making them3144            immediately available for use.)");3145 3146  //----------------------------------------------------------------------------3147  // Mapping of PyDialectDescriptor3148  //----------------------------------------------------------------------------3149  nb::class_<PyDialectDescriptor>(m, "DialectDescriptor")3150      .def_prop_ro(3151          "namespace",3152          [](PyDialectDescriptor &self) {3153            MlirStringRef ns = mlirDialectGetNamespace(self.get());3154            return nb::str(ns.data, ns.length);3155          },3156          "Returns the namespace of the dialect.")3157      .def(3158          "__repr__",3159          [](PyDialectDescriptor &self) {3160            MlirStringRef ns = mlirDialectGetNamespace(self.get());3161            std::string repr("<DialectDescriptor ");3162            repr.append(ns.data, ns.length);3163            repr.append(">");3164            return repr;3165          },3166          nb::sig("def __repr__(self) -> str"),3167          "Returns a string representation of the dialect descriptor.");3168 3169  //----------------------------------------------------------------------------3170  // Mapping of PyDialects3171  //----------------------------------------------------------------------------3172  nb::class_<PyDialects>(m, "Dialects")3173      .def(3174          "__getitem__",3175          [=](PyDialects &self, std::string keyName) {3176            MlirDialect dialect =3177                self.getDialectForKey(keyName, /*attrError=*/false);3178            nb::object descriptor =3179                nb::cast(PyDialectDescriptor{self.getContext(), dialect});3180            return createCustomDialectWrapper(keyName, std::move(descriptor));3181          },3182          "Gets a dialect by name using subscript notation.")3183      .def(3184          "__getattr__",3185          [=](PyDialects &self, std::string attrName) {3186            MlirDialect dialect =3187                self.getDialectForKey(attrName, /*attrError=*/true);3188            nb::object descriptor =3189                nb::cast(PyDialectDescriptor{self.getContext(), dialect});3190            return createCustomDialectWrapper(attrName, std::move(descriptor));3191          },3192          "Gets a dialect by name using attribute notation.");3193 3194  //----------------------------------------------------------------------------3195  // Mapping of PyDialect3196  //----------------------------------------------------------------------------3197  nb::class_<PyDialect>(m, "Dialect")3198      .def(nb::init<nb::object>(), nb::arg("descriptor"),3199           "Creates a Dialect from a DialectDescriptor.")3200      .def_prop_ro(3201          "descriptor", [](PyDialect &self) { return self.getDescriptor(); },3202          "Returns the DialectDescriptor for this dialect.")3203      .def(3204          "__repr__",3205          [](const nb::object &self) {3206            auto clazz = self.attr("__class__");3207            return nb::str("<Dialect ") +3208                   self.attr("descriptor").attr("namespace") +3209                   nb::str(" (class ") + clazz.attr("__module__") +3210                   nb::str(".") + clazz.attr("__name__") + nb::str(")>");3211          },3212          nb::sig("def __repr__(self) -> str"),3213          "Returns a string representation of the dialect.");3214 3215  //----------------------------------------------------------------------------3216  // Mapping of PyDialectRegistry3217  //----------------------------------------------------------------------------3218  nb::class_<PyDialectRegistry>(m, "DialectRegistry")3219      .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyDialectRegistry::getCapsule,3220                   "Gets a capsule wrapping the MlirDialectRegistry.")3221      .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR,3222                  &PyDialectRegistry::createFromCapsule,3223                  "Creates a DialectRegistry from a capsule wrapping "3224                  "`MlirDialectRegistry`.")3225      .def(nb::init<>(), "Creates a new empty dialect registry.");3226 3227  //----------------------------------------------------------------------------3228  // Mapping of Location3229  //----------------------------------------------------------------------------3230  nb::class_<PyLocation>(m, "Location")3231      .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyLocation::getCapsule,3232                   "Gets a capsule wrapping the MlirLocation.")3233      .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyLocation::createFromCapsule,3234                  "Creates a Location from a capsule wrapping MlirLocation.")3235      .def("__enter__", &PyLocation::contextEnter,3236           "Enters the location as a context manager.")3237      .def("__exit__", &PyLocation::contextExit, nb::arg("exc_type").none(),3238           nb::arg("exc_value").none(), nb::arg("traceback").none(),3239           "Exits the location context manager.")3240      .def(3241          "__eq__",3242          [](PyLocation &self, PyLocation &other) -> bool {3243            return mlirLocationEqual(self, other);3244          },3245          "Compares two locations for equality.")3246      .def(3247          "__eq__", [](PyLocation &self, nb::object other) { return false; },3248          "Compares location with non-location object (always returns False).")3249      .def_prop_ro_static(3250          "current",3251          [](nb::object & /*class*/) -> std::optional<PyLocation *> {3252            auto *loc = PyThreadContextEntry::getDefaultLocation();3253            if (!loc)3254              return std::nullopt;3255            return loc;3256          },3257          // clang-format off3258          nb::sig("def current(/) -> Location | None"),3259          // clang-format on3260          "Gets the Location bound to the current thread or raises ValueError.")3261      .def_static(3262          "unknown",3263          [](DefaultingPyMlirContext context) {3264            return PyLocation(context->getRef(),3265                              mlirLocationUnknownGet(context->get()));3266          },3267          nb::arg("context") = nb::none(),3268          "Gets a Location representing an unknown location.")3269      .def_static(3270          "callsite",3271          [](PyLocation callee, const std::vector<PyLocation> &frames,3272             DefaultingPyMlirContext context) {3273            if (frames.empty())3274              throw nb::value_error("No caller frames provided.");3275            MlirLocation caller = frames.back().get();3276            for (const PyLocation &frame :3277                 llvm::reverse(llvm::ArrayRef(frames).drop_back()))3278              caller = mlirLocationCallSiteGet(frame.get(), caller);3279            return PyLocation(context->getRef(),3280                              mlirLocationCallSiteGet(callee.get(), caller));3281          },3282          nb::arg("callee"), nb::arg("frames"), nb::arg("context") = nb::none(),3283          "Gets a Location representing a caller and callsite.")3284      .def("is_a_callsite", mlirLocationIsACallSite,3285           "Returns True if this location is a CallSiteLoc.")3286      .def_prop_ro(3287          "callee",3288          [](PyLocation &self) {3289            return PyLocation(self.getContext(),3290                              mlirLocationCallSiteGetCallee(self));3291          },3292          "Gets the callee location from a CallSiteLoc.")3293      .def_prop_ro(3294          "caller",3295          [](PyLocation &self) {3296            return PyLocation(self.getContext(),3297                              mlirLocationCallSiteGetCaller(self));3298          },3299          "Gets the caller location from a CallSiteLoc.")3300      .def_static(3301          "file",3302          [](std::string filename, int line, int col,3303             DefaultingPyMlirContext context) {3304            return PyLocation(3305                context->getRef(),3306                mlirLocationFileLineColGet(3307                    context->get(), toMlirStringRef(filename), line, col));3308          },3309          nb::arg("filename"), nb::arg("line"), nb::arg("col"),3310          nb::arg("context") = nb::none(),3311          "Gets a Location representing a file, line and column.")3312      .def_static(3313          "file",3314          [](std::string filename, int startLine, int startCol, int endLine,3315             int endCol, DefaultingPyMlirContext context) {3316            return PyLocation(context->getRef(),3317                              mlirLocationFileLineColRangeGet(3318                                  context->get(), toMlirStringRef(filename),3319                                  startLine, startCol, endLine, endCol));3320          },3321          nb::arg("filename"), nb::arg("start_line"), nb::arg("start_col"),3322          nb::arg("end_line"), nb::arg("end_col"),3323          nb::arg("context") = nb::none(),3324          "Gets a Location representing a file, line and column range.")3325      .def("is_a_file", mlirLocationIsAFileLineColRange,3326           "Returns True if this location is a FileLineColLoc.")3327      .def_prop_ro(3328          "filename",3329          [](MlirLocation loc) {3330            return mlirIdentifierStr(3331                mlirLocationFileLineColRangeGetFilename(loc));3332          },3333          "Gets the filename from a FileLineColLoc.")3334      .def_prop_ro("start_line", mlirLocationFileLineColRangeGetStartLine,3335                   "Gets the start line number from a `FileLineColLoc`.")3336      .def_prop_ro("start_col", mlirLocationFileLineColRangeGetStartColumn,3337                   "Gets the start column number from a `FileLineColLoc`.")3338      .def_prop_ro("end_line", mlirLocationFileLineColRangeGetEndLine,3339                   "Gets the end line number from a `FileLineColLoc`.")3340      .def_prop_ro("end_col", mlirLocationFileLineColRangeGetEndColumn,3341                   "Gets the end column number from a `FileLineColLoc`.")3342      .def_static(3343          "fused",3344          [](const std::vector<PyLocation> &pyLocations,3345             std::optional<PyAttribute> metadata,3346             DefaultingPyMlirContext context) {3347            llvm::SmallVector<MlirLocation, 4> locations;3348            locations.reserve(pyLocations.size());3349            for (auto &pyLocation : pyLocations)3350              locations.push_back(pyLocation.get());3351            MlirLocation location = mlirLocationFusedGet(3352                context->get(), locations.size(), locations.data(),3353                metadata ? metadata->get() : MlirAttribute{0});3354            return PyLocation(context->getRef(), location);3355          },3356          nb::arg("locations"), nb::arg("metadata") = nb::none(),3357          nb::arg("context") = nb::none(),3358          "Gets a Location representing a fused location with optional "3359          "metadata.")3360      .def("is_a_fused", mlirLocationIsAFused,3361           "Returns True if this location is a `FusedLoc`.")3362      .def_prop_ro(3363          "locations",3364          [](PyLocation &self) {3365            unsigned numLocations = mlirLocationFusedGetNumLocations(self);3366            std::vector<MlirLocation> locations(numLocations);3367            if (numLocations)3368              mlirLocationFusedGetLocations(self, locations.data());3369            std::vector<PyLocation> pyLocations{};3370            pyLocations.reserve(numLocations);3371            for (unsigned i = 0; i < numLocations; ++i)3372              pyLocations.emplace_back(self.getContext(), locations[i]);3373            return pyLocations;3374          },3375          "Gets the list of locations from a `FusedLoc`.")3376      .def_static(3377          "name",3378          [](std::string name, std::optional<PyLocation> childLoc,3379             DefaultingPyMlirContext context) {3380            return PyLocation(3381                context->getRef(),3382                mlirLocationNameGet(3383                    context->get(), toMlirStringRef(name),3384                    childLoc ? childLoc->get()3385                             : mlirLocationUnknownGet(context->get())));3386          },3387          nb::arg("name"), nb::arg("childLoc") = nb::none(),3388          nb::arg("context") = nb::none(),3389          "Gets a Location representing a named location with optional child "3390          "location.")3391      .def("is_a_name", mlirLocationIsAName,3392           "Returns True if this location is a `NameLoc`.")3393      .def_prop_ro(3394          "name_str",3395          [](MlirLocation loc) {3396            return mlirIdentifierStr(mlirLocationNameGetName(loc));3397          },3398          "Gets the name string from a `NameLoc`.")3399      .def_prop_ro(3400          "child_loc",3401          [](PyLocation &self) {3402            return PyLocation(self.getContext(),3403                              mlirLocationNameGetChildLoc(self));3404          },3405          "Gets the child location from a `NameLoc`.")3406      .def_static(3407          "from_attr",3408          [](PyAttribute &attribute, DefaultingPyMlirContext context) {3409            return PyLocation(context->getRef(),3410                              mlirLocationFromAttribute(attribute));3411          },3412          nb::arg("attribute"), nb::arg("context") = nb::none(),3413          "Gets a Location from a `LocationAttr`.")3414      .def_prop_ro(3415          "context",3416          [](PyLocation &self) -> nb::typed<nb::object, PyMlirContext> {3417            return self.getContext().getObject();3418          },3419          "Context that owns the `Location`.")3420      .def_prop_ro(3421          "attr",3422          [](PyLocation &self) {3423            return PyAttribute(self.getContext(),3424                               mlirLocationGetAttribute(self));3425          },3426          "Get the underlying `LocationAttr`.")3427      .def(3428          "emit_error",3429          [](PyLocation &self, std::string message) {3430            mlirEmitError(self, message.c_str());3431          },3432          nb::arg("message"),3433          R"(3434            Emits an error diagnostic at this location.3435 3436            Args:3437              message: The error message to emit.)")3438      .def(3439          "__repr__",3440          [](PyLocation &self) {3441            PyPrintAccumulator printAccum;3442            mlirLocationPrint(self, printAccum.getCallback(),3443                              printAccum.getUserData());3444            return printAccum.join();3445          },3446          "Returns the assembly representation of the location.");3447 3448  //----------------------------------------------------------------------------3449  // Mapping of Module3450  //----------------------------------------------------------------------------3451  nb::class_<PyModule>(m, "Module", nb::is_weak_referenceable())3452      .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyModule::getCapsule,3453                   "Gets a capsule wrapping the MlirModule.")3454      .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyModule::createFromCapsule,3455                  R"(3456                    Creates a Module from a `MlirModule` wrapped by a capsule (i.e. `module._CAPIPtr`).3457 3458                    This returns a new object **BUT** `_clear_mlir_module(module)` must be called to3459                    prevent double-frees (of the underlying `mlir::Module`).)")3460      .def("_clear_mlir_module", &PyModule::clearMlirModule,3461           R"(3462             Clears the internal MLIR module reference.3463 3464             This is used internally to prevent double-free when ownership is transferred3465             via the C API capsule mechanism. Not intended for normal use.)")3466      .def_static(3467          "parse",3468          [](const std::string &moduleAsm, DefaultingPyMlirContext context)3469              -> nb::typed<nb::object, PyModule> {3470            PyMlirContext::ErrorCapture errors(context->getRef());3471            MlirModule module = mlirModuleCreateParse(3472                context->get(), toMlirStringRef(moduleAsm));3473            if (mlirModuleIsNull(module))3474              throw MLIRError("Unable to parse module assembly", errors.take());3475            return PyModule::forModule(module).releaseObject();3476          },3477          nb::arg("asm"), nb::arg("context") = nb::none(),3478          kModuleParseDocstring)3479      .def_static(3480          "parse",3481          [](nb::bytes moduleAsm, DefaultingPyMlirContext context)3482              -> nb::typed<nb::object, PyModule> {3483            PyMlirContext::ErrorCapture errors(context->getRef());3484            MlirModule module = mlirModuleCreateParse(3485                context->get(), toMlirStringRef(moduleAsm));3486            if (mlirModuleIsNull(module))3487              throw MLIRError("Unable to parse module assembly", errors.take());3488            return PyModule::forModule(module).releaseObject();3489          },3490          nb::arg("asm"), nb::arg("context") = nb::none(),3491          kModuleParseDocstring)3492      .def_static(3493          "parseFile",3494          [](const std::string &path, DefaultingPyMlirContext context)3495              -> nb::typed<nb::object, PyModule> {3496            PyMlirContext::ErrorCapture errors(context->getRef());3497            MlirModule module = mlirModuleCreateParseFromFile(3498                context->get(), toMlirStringRef(path));3499            if (mlirModuleIsNull(module))3500              throw MLIRError("Unable to parse module assembly", errors.take());3501            return PyModule::forModule(module).releaseObject();3502          },3503          nb::arg("path"), nb::arg("context") = nb::none(),3504          kModuleParseDocstring)3505      .def_static(3506          "create",3507          [](const std::optional<PyLocation> &loc)3508              -> nb::typed<nb::object, PyModule> {3509            PyLocation pyLoc = maybeGetTracebackLocation(loc);3510            MlirModule module = mlirModuleCreateEmpty(pyLoc.get());3511            return PyModule::forModule(module).releaseObject();3512          },3513          nb::arg("loc") = nb::none(), "Creates an empty module.")3514      .def_prop_ro(3515          "context",3516          [](PyModule &self) -> nb::typed<nb::object, PyMlirContext> {3517            return self.getContext().getObject();3518          },3519          "Context that created the `Module`.")3520      .def_prop_ro(3521          "operation",3522          [](PyModule &self) -> nb::typed<nb::object, PyOperation> {3523            return PyOperation::forOperation(self.getContext(),3524                                             mlirModuleGetOperation(self.get()),3525                                             self.getRef().releaseObject())3526                .releaseObject();3527          },3528          "Accesses the module as an operation.")3529      .def_prop_ro(3530          "body",3531          [](PyModule &self) {3532            PyOperationRef moduleOp = PyOperation::forOperation(3533                self.getContext(), mlirModuleGetOperation(self.get()),3534                self.getRef().releaseObject());3535            PyBlock returnBlock(moduleOp, mlirModuleGetBody(self.get()));3536            return returnBlock;3537          },3538          "Return the block for this module.")3539      .def(3540          "dump",3541          [](PyModule &self) {3542            mlirOperationDump(mlirModuleGetOperation(self.get()));3543          },3544          kDumpDocstring)3545      .def(3546          "__str__",3547          [](const nb::object &self) {3548            // Defer to the operation's __str__.3549            return self.attr("operation").attr("__str__")();3550          },3551          nb::sig("def __str__(self) -> str"),3552          R"(3553            Gets the assembly form of the operation with default options.3554 3555            If more advanced control over the assembly formatting or I/O options is needed,3556            use the dedicated print or get_asm method, which supports keyword arguments to3557            customize behavior.3558          )")3559      .def(3560          "__eq__",3561          [](PyModule &self, PyModule &other) {3562            return mlirModuleEqual(self.get(), other.get());3563          },3564          "other"_a, "Compares two modules for equality.")3565      .def(3566          "__hash__",3567          [](PyModule &self) { return mlirModuleHashValue(self.get()); },3568          "Returns the hash value of the module.");3569 3570  //----------------------------------------------------------------------------3571  // Mapping of Operation.3572  //----------------------------------------------------------------------------3573  nb::class_<PyOperationBase>(m, "_OperationBase")3574      .def_prop_ro(3575          MLIR_PYTHON_CAPI_PTR_ATTR,3576          [](PyOperationBase &self) {3577            return self.getOperation().getCapsule();3578          },3579          "Gets a capsule wrapping the `MlirOperation`.")3580      .def(3581          "__eq__",3582          [](PyOperationBase &self, PyOperationBase &other) {3583            return mlirOperationEqual(self.getOperation().get(),3584                                      other.getOperation().get());3585          },3586          "Compares two operations for equality.")3587      .def(3588          "__eq__",3589          [](PyOperationBase &self, nb::object other) { return false; },3590          "Compares operation with non-operation object (always returns "3591          "False).")3592      .def(3593          "__hash__",3594          [](PyOperationBase &self) {3595            return mlirOperationHashValue(self.getOperation().get());3596          },3597          "Returns the hash value of the operation.")3598      .def_prop_ro(3599          "attributes",3600          [](PyOperationBase &self) {3601            return PyOpAttributeMap(self.getOperation().getRef());3602          },3603          "Returns a dictionary-like map of operation attributes.")3604      .def_prop_ro(3605          "context",3606          [](PyOperationBase &self) -> nb::typed<nb::object, PyMlirContext> {3607            PyOperation &concreteOperation = self.getOperation();3608            concreteOperation.checkValid();3609            return concreteOperation.getContext().getObject();3610          },3611          "Context that owns the operation.")3612      .def_prop_ro(3613          "name",3614          [](PyOperationBase &self) {3615            auto &concreteOperation = self.getOperation();3616            concreteOperation.checkValid();3617            MlirOperation operation = concreteOperation.get();3618            return mlirIdentifierStr(mlirOperationGetName(operation));3619          },3620          "Returns the fully qualified name of the operation.")3621      .def_prop_ro(3622          "operands",3623          [](PyOperationBase &self) {3624            return PyOpOperandList(self.getOperation().getRef());3625          },3626          "Returns the list of operation operands.")3627      .def_prop_ro(3628          "regions",3629          [](PyOperationBase &self) {3630            return PyRegionList(self.getOperation().getRef());3631          },3632          "Returns the list of operation regions.")3633      .def_prop_ro(3634          "results",3635          [](PyOperationBase &self) {3636            return PyOpResultList(self.getOperation().getRef());3637          },3638          "Returns the list of Operation results.")3639      .def_prop_ro(3640          "result",3641          [](PyOperationBase &self) -> nb::typed<nb::object, PyOpResult> {3642            auto &operation = self.getOperation();3643            return PyOpResult(operation.getRef(), getUniqueResult(operation))3644                .maybeDownCast();3645          },3646          "Shortcut to get an op result if it has only one (throws an error "3647          "otherwise).")3648      .def_prop_rw(3649          "location",3650          [](PyOperationBase &self) {3651            PyOperation &operation = self.getOperation();3652            return PyLocation(operation.getContext(),3653                              mlirOperationGetLocation(operation.get()));3654          },3655          [](PyOperationBase &self, const PyLocation &location) {3656            PyOperation &operation = self.getOperation();3657            mlirOperationSetLocation(operation.get(), location.get());3658          },3659          nb::for_getter("Returns the source location the operation was "3660                         "defined or derived from."),3661          nb::for_setter("Sets the source location the operation was defined "3662                         "or derived from."))3663      .def_prop_ro(3664          "parent",3665          [](PyOperationBase &self)3666              -> std::optional<nb::typed<nb::object, PyOperation>> {3667            auto parent = self.getOperation().getParentOperation();3668            if (parent)3669              return parent->getObject();3670            return {};3671          },3672          "Returns the parent operation, or `None` if at top level.")3673      .def(3674          "__str__",3675          [](PyOperationBase &self) {3676            return self.getAsm(/*binary=*/false,3677                               /*largeElementsLimit=*/std::nullopt,3678                               /*largeResourceLimit=*/std::nullopt,3679                               /*enableDebugInfo=*/false,3680                               /*prettyDebugInfo=*/false,3681                               /*printGenericOpForm=*/false,3682                               /*useLocalScope=*/false,3683                               /*useNameLocAsPrefix=*/false,3684                               /*assumeVerified=*/false,3685                               /*skipRegions=*/false);3686          },3687          nb::sig("def __str__(self) -> str"),3688          "Returns the assembly form of the operation.")3689      .def("print",3690           nb::overload_cast<PyAsmState &, nb::object, bool>(3691               &PyOperationBase::print),3692           nb::arg("state"), nb::arg("file") = nb::none(),3693           nb::arg("binary") = false,3694           R"(3695             Prints the assembly form of the operation to a file like object.3696 3697             Args:3698               state: `AsmState` capturing the operation numbering and flags.3699               file: Optional file like object to write to. Defaults to sys.stdout.3700               binary: Whether to write `bytes` (True) or `str` (False). Defaults to False.)")3701      .def("print",3702           nb::overload_cast<std::optional<int64_t>, std::optional<int64_t>,3703                             bool, bool, bool, bool, bool, bool, nb::object,3704                             bool, bool>(&PyOperationBase::print),3705           // Careful: Lots of arguments must match up with print method.3706           nb::arg("large_elements_limit") = nb::none(),3707           nb::arg("large_resource_limit") = nb::none(),3708           nb::arg("enable_debug_info") = false,3709           nb::arg("pretty_debug_info") = false,3710           nb::arg("print_generic_op_form") = false,3711           nb::arg("use_local_scope") = false,3712           nb::arg("use_name_loc_as_prefix") = false,3713           nb::arg("assume_verified") = false, nb::arg("file") = nb::none(),3714           nb::arg("binary") = false, nb::arg("skip_regions") = false,3715           R"(3716             Prints the assembly form of the operation to a file like object.3717 3718             Args:3719               large_elements_limit: Whether to elide elements attributes above this3720                 number of elements. Defaults to None (no limit).3721               large_resource_limit: Whether to elide resource attributes above this3722                 number of characters. Defaults to None (no limit). If large_elements_limit3723                 is set and this is None, the behavior will be to use large_elements_limit3724                 as large_resource_limit.3725               enable_debug_info: Whether to print debug/location information. Defaults3726                 to False.3727               pretty_debug_info: Whether to format debug information for easier reading3728                 by a human (warning: the result is unparseable). Defaults to False.3729               print_generic_op_form: Whether to print the generic assembly forms of all3730                 ops. Defaults to False.3731               use_local_scope: Whether to print in a way that is more optimized for3732                 multi-threaded access but may not be consistent with how the overall3733                 module prints.3734               use_name_loc_as_prefix: Whether to use location attributes (NameLoc) as3735                 prefixes for the SSA identifiers. Defaults to False.3736               assume_verified: By default, if not printing generic form, the verifier3737                 will be run and if it fails, generic form will be printed with a comment3738                 about failed verification. While a reasonable default for interactive use,3739                 for systematic use, it is often better for the caller to verify explicitly3740                 and report failures in a more robust fashion. Set this to True if doing this3741                 in order to avoid running a redundant verification. If the IR is actually3742                 invalid, behavior is undefined.3743               file: The file like object to write to. Defaults to sys.stdout.3744               binary: Whether to write bytes (True) or str (False). Defaults to False.3745               skip_regions: Whether to skip printing regions. Defaults to False.)")3746      .def("write_bytecode", &PyOperationBase::writeBytecode, nb::arg("file"),3747           nb::arg("desired_version") = nb::none(),3748           R"(3749             Write the bytecode form of the operation to a file like object.3750 3751             Args:3752               file: The file like object to write to.3753               desired_version: Optional version of bytecode to emit.3754             Returns:3755               The bytecode writer status.)")3756      .def("get_asm", &PyOperationBase::getAsm,3757           // Careful: Lots of arguments must match up with get_asm method.3758           nb::arg("binary") = false,3759           nb::arg("large_elements_limit") = nb::none(),3760           nb::arg("large_resource_limit") = nb::none(),3761           nb::arg("enable_debug_info") = false,3762           nb::arg("pretty_debug_info") = false,3763           nb::arg("print_generic_op_form") = false,3764           nb::arg("use_local_scope") = false,3765           nb::arg("use_name_loc_as_prefix") = false,3766           nb::arg("assume_verified") = false, nb::arg("skip_regions") = false,3767           R"(3768            Gets the assembly form of the operation with all options available.3769 3770            Args:3771              binary: Whether to return a bytes (True) or str (False) object. Defaults to3772                False.3773              ... others ...: See the print() method for common keyword arguments for3774                configuring the printout.3775            Returns:3776              Either a bytes or str object, depending on the setting of the `binary`3777              argument.)")3778      .def("verify", &PyOperationBase::verify,3779           "Verify the operation. Raises MLIRError if verification fails, and "3780           "returns true otherwise.")3781      .def("move_after", &PyOperationBase::moveAfter, nb::arg("other"),3782           "Puts self immediately after the other operation in its parent "3783           "block.")3784      .def("move_before", &PyOperationBase::moveBefore, nb::arg("other"),3785           "Puts self immediately before the other operation in its parent "3786           "block.")3787      .def("is_before_in_block", &PyOperationBase::isBeforeInBlock,3788           nb::arg("other"),3789           R"(3790             Checks if this operation is before another in the same block.3791 3792             Args:3793               other: Another operation in the same parent block.3794 3795             Returns:3796               True if this operation is before `other` in the operation list of the parent block.)")3797      .def(3798          "clone",3799          [](PyOperationBase &self,3800             const nb::object &ip) -> nb::typed<nb::object, PyOperation> {3801            return self.getOperation().clone(ip);3802          },3803          nb::arg("ip") = nb::none(),3804          R"(3805            Creates a deep copy of the operation.3806 3807            Args:3808              ip: Optional insertion point where the cloned operation should be inserted.3809                If None, the current insertion point is used. If False, the operation3810                remains detached.3811 3812            Returns:3813              A new Operation that is a clone of this operation.)")3814      .def(3815          "detach_from_parent",3816          [](PyOperationBase &self) -> nb::typed<nb::object, PyOpView> {3817            PyOperation &operation = self.getOperation();3818            operation.checkValid();3819            if (!operation.isAttached())3820              throw nb::value_error("Detached operation has no parent.");3821 3822            operation.detachFromParent();3823            return operation.createOpView();3824          },3825          "Detaches the operation from its parent block.")3826      .def_prop_ro(3827          "attached",3828          [](PyOperationBase &self) {3829            PyOperation &operation = self.getOperation();3830            operation.checkValid();3831            return operation.isAttached();3832          },3833          "Reports if the operation is attached to its parent block.")3834      .def(3835          "erase", [](PyOperationBase &self) { self.getOperation().erase(); },3836          R"(3837            Erases the operation and frees its memory.3838 3839            Note:3840              After erasing, any Python references to the operation become invalid.)")3841      .def("walk", &PyOperationBase::walk, nb::arg("callback"),3842           nb::arg("walk_order") = MlirWalkPostOrder,3843           // clang-format off3844          nb::sig("def walk(self, callback: Callable[[Operation], WalkResult], walk_order: WalkOrder) -> None"),3845           // clang-format on3846           R"(3847             Walks the operation tree with a callback function.3848 3849             Args:3850               callback: A callable that takes an Operation and returns a WalkResult.3851               walk_order: The order of traversal (PRE_ORDER or POST_ORDER).)");3852 3853  nb::class_<PyOperation, PyOperationBase>(m, "Operation")3854      .def_static(3855          "create",3856          [](std::string_view name,3857             std::optional<std::vector<PyType *>> results,3858             std::optional<std::vector<PyValue *>> operands,3859             std::optional<nb::dict> attributes,3860             std::optional<std::vector<PyBlock *>> successors, int regions,3861             const std::optional<PyLocation> &location,3862             const nb::object &maybeIp,3863             bool inferType) -> nb::typed<nb::object, PyOperation> {3864            // Unpack/validate operands.3865            llvm::SmallVector<MlirValue, 4> mlirOperands;3866            if (operands) {3867              mlirOperands.reserve(operands->size());3868              for (PyValue *operand : *operands) {3869                if (!operand)3870                  throw nb::value_error("operand value cannot be None");3871                mlirOperands.push_back(operand->get());3872              }3873            }3874 3875            PyLocation pyLoc = maybeGetTracebackLocation(location);3876            return PyOperation::create(name, results, mlirOperands, attributes,3877                                       successors, regions, pyLoc, maybeIp,3878                                       inferType);3879          },3880          nb::arg("name"), nb::arg("results") = nb::none(),3881          nb::arg("operands") = nb::none(), nb::arg("attributes") = nb::none(),3882          nb::arg("successors") = nb::none(), nb::arg("regions") = 0,3883          nb::arg("loc") = nb::none(), nb::arg("ip") = nb::none(),3884          nb::arg("infer_type") = false,3885          R"(3886            Creates a new operation.3887 3888            Args:3889              name: Operation name (e.g. `dialect.operation`).3890              results: Optional sequence of Type representing op result types.3891              operands: Optional operands of the operation.3892              attributes: Optional Dict of {str: Attribute}.3893              successors: Optional List of Block for the operation's successors.3894              regions: Number of regions to create (default = 0).3895              location: Optional Location object (defaults to resolve from context manager).3896              ip: Optional InsertionPoint (defaults to resolve from context manager or set to False to disable insertion, even with an insertion point set in the context manager).3897              infer_type: Whether to infer result types (default = False).3898            Returns:3899              A new detached Operation object. Detached operations can be added to blocks, which causes them to become attached.)")3900      .def_static(3901          "parse",3902          [](const std::string &sourceStr, const std::string &sourceName,3903             DefaultingPyMlirContext context)3904              -> nb::typed<nb::object, PyOpView> {3905            return PyOperation::parse(context->getRef(), sourceStr, sourceName)3906                ->createOpView();3907          },3908          nb::arg("source"), nb::kw_only(), nb::arg("source_name") = "",3909          nb::arg("context") = nb::none(),3910          "Parses an operation. Supports both text assembly format and binary "3911          "bytecode format.")3912      .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyOperation::getCapsule,3913                   "Gets a capsule wrapping the MlirOperation.")3914      .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR,3915                  &PyOperation::createFromCapsule,3916                  "Creates an Operation from a capsule wrapping MlirOperation.")3917      .def_prop_ro(3918          "operation",3919          [](nb::object self) -> nb::typed<nb::object, PyOperation> {3920            return self;3921          },3922          "Returns self (the operation).")3923      .def_prop_ro(3924          "opview",3925          [](PyOperation &self) -> nb::typed<nb::object, PyOpView> {3926            return self.createOpView();3927          },3928          R"(3929            Returns an OpView of this operation.3930 3931            Note:3932              If the operation has a registered and loaded dialect then this OpView will3933              be concrete wrapper class.)")3934      .def_prop_ro("block", &PyOperation::getBlock,3935                   "Returns the block containing this operation.")3936      .def_prop_ro(3937          "successors",3938          [](PyOperationBase &self) {3939            return PyOpSuccessors(self.getOperation().getRef());3940          },3941          "Returns the list of Operation successors.")3942      .def("_set_invalid", &PyOperation::setInvalid,3943           "Invalidate the operation.");3944 3945  auto opViewClass =3946      nb::class_<PyOpView, PyOperationBase>(m, "OpView")3947          .def(nb::init<nb::typed<nb::object, PyOperation>>(),3948               nb::arg("operation"))3949          .def(3950              "__init__",3951              [](PyOpView *self, std::string_view name,3952                 std::tuple<int, bool> opRegionSpec,3953                 nb::object operandSegmentSpecObj,3954                 nb::object resultSegmentSpecObj,3955                 std::optional<nb::list> resultTypeList, nb::list operandList,3956                 std::optional<nb::dict> attributes,3957                 std::optional<std::vector<PyBlock *>> successors,3958                 std::optional<int> regions,3959                 const std::optional<PyLocation> &location,3960                 const nb::object &maybeIp) {3961                PyLocation pyLoc = maybeGetTracebackLocation(location);3962                new (self) PyOpView(PyOpView::buildGeneric(3963                    name, opRegionSpec, operandSegmentSpecObj,3964                    resultSegmentSpecObj, resultTypeList, operandList,3965                    attributes, successors, regions, pyLoc, maybeIp));3966              },3967              nb::arg("name"), nb::arg("opRegionSpec"),3968              nb::arg("operandSegmentSpecObj") = nb::none(),3969              nb::arg("resultSegmentSpecObj") = nb::none(),3970              nb::arg("results") = nb::none(), nb::arg("operands") = nb::none(),3971              nb::arg("attributes") = nb::none(),3972              nb::arg("successors") = nb::none(),3973              nb::arg("regions") = nb::none(), nb::arg("loc") = nb::none(),3974              nb::arg("ip") = nb::none())3975          .def_prop_ro(3976              "operation",3977              [](PyOpView &self) -> nb::typed<nb::object, PyOperation> {3978                return self.getOperationObject();3979              })3980          .def_prop_ro("opview",3981                       [](nb::object self) -> nb::typed<nb::object, PyOpView> {3982                         return self;3983                       })3984          .def(3985              "__str__",3986              [](PyOpView &self) { return nb::str(self.getOperationObject()); })3987          .def_prop_ro(3988              "successors",3989              [](PyOperationBase &self) {3990                return PyOpSuccessors(self.getOperation().getRef());3991              },3992              "Returns the list of Operation successors.")3993          .def(3994              "_set_invalid",3995              [](PyOpView &self) { self.getOperation().setInvalid(); },3996              "Invalidate the operation.");3997  opViewClass.attr("_ODS_REGIONS") = nb::make_tuple(0, true);3998  opViewClass.attr("_ODS_OPERAND_SEGMENTS") = nb::none();3999  opViewClass.attr("_ODS_RESULT_SEGMENTS") = nb::none();4000  // It is faster to pass the operation_name, ods_regions, and4001  // ods_operand_segments/ods_result_segments as arguments to the constructor,4002  // rather than to access them as attributes.4003  opViewClass.attr("build_generic") = classmethod(4004      [](nb::handle cls, std::optional<nb::list> resultTypeList,4005         nb::list operandList, std::optional<nb::dict> attributes,4006         std::optional<std::vector<PyBlock *>> successors,4007         std::optional<int> regions, std::optional<PyLocation> location,4008         const nb::object &maybeIp) {4009        std::string name = nb::cast<std::string>(cls.attr("OPERATION_NAME"));4010        std::tuple<int, bool> opRegionSpec =4011            nb::cast<std::tuple<int, bool>>(cls.attr("_ODS_REGIONS"));4012        nb::object operandSegmentSpec = cls.attr("_ODS_OPERAND_SEGMENTS");4013        nb::object resultSegmentSpec = cls.attr("_ODS_RESULT_SEGMENTS");4014        PyLocation pyLoc = maybeGetTracebackLocation(location);4015        return PyOpView::buildGeneric(name, opRegionSpec, operandSegmentSpec,4016                                      resultSegmentSpec, resultTypeList,4017                                      operandList, attributes, successors,4018                                      regions, pyLoc, maybeIp);4019      },4020      nb::arg("cls"), nb::arg("results") = nb::none(),4021      nb::arg("operands") = nb::none(), nb::arg("attributes") = nb::none(),4022      nb::arg("successors") = nb::none(), nb::arg("regions") = nb::none(),4023      nb::arg("loc") = nb::none(), nb::arg("ip") = nb::none(),4024      "Builds a specific, generated OpView based on class level attributes.");4025  opViewClass.attr("parse") = classmethod(4026      [](const nb::object &cls, const std::string &sourceStr,4027         const std::string &sourceName,4028         DefaultingPyMlirContext context) -> nb::typed<nb::object, PyOpView> {4029        PyOperationRef parsed =4030            PyOperation::parse(context->getRef(), sourceStr, sourceName);4031 4032        // Check if the expected operation was parsed, and cast to to the4033        // appropriate `OpView` subclass if successful.4034        // NOTE: This accesses attributes that have been automatically added to4035        // `OpView` subclasses, and is not intended to be used on `OpView`4036        // directly.4037        std::string clsOpName =4038            nb::cast<std::string>(cls.attr("OPERATION_NAME"));4039        MlirStringRef identifier =4040            mlirIdentifierStr(mlirOperationGetName(*parsed.get()));4041        std::string_view parsedOpName(identifier.data, identifier.length);4042        if (clsOpName != parsedOpName)4043          throw MLIRError(Twine("Expected a '") + clsOpName + "' op, got: '" +4044                          parsedOpName + "'");4045        return PyOpView::constructDerived(cls, parsed.getObject());4046      },4047      nb::arg("cls"), nb::arg("source"), nb::kw_only(),4048      nb::arg("source_name") = "", nb::arg("context") = nb::none(),4049      "Parses a specific, generated OpView based on class level attributes.");4050 4051  //----------------------------------------------------------------------------4052  // Mapping of PyRegion.4053  //----------------------------------------------------------------------------4054  nb::class_<PyRegion>(m, "Region")4055      .def_prop_ro(4056          "blocks",4057          [](PyRegion &self) {4058            return PyBlockList(self.getParentOperation(), self.get());4059          },4060          "Returns a forward-optimized sequence of blocks.")4061      .def_prop_ro(4062          "owner",4063          [](PyRegion &self) -> nb::typed<nb::object, PyOpView> {4064            return self.getParentOperation()->createOpView();4065          },4066          "Returns the operation owning this region.")4067      .def(4068          "__iter__",4069          [](PyRegion &self) {4070            self.checkValid();4071            MlirBlock firstBlock = mlirRegionGetFirstBlock(self.get());4072            return PyBlockIterator(self.getParentOperation(), firstBlock);4073          },4074          "Iterates over blocks in the region.")4075      .def(4076          "__eq__",4077          [](PyRegion &self, PyRegion &other) {4078            return self.get().ptr == other.get().ptr;4079          },4080          "Compares two regions for pointer equality.")4081      .def(4082          "__eq__", [](PyRegion &self, nb::object &other) { return false; },4083          "Compares region with non-region object (always returns False).");4084 4085  //----------------------------------------------------------------------------4086  // Mapping of PyBlock.4087  //----------------------------------------------------------------------------4088  nb::class_<PyBlock>(m, "Block")4089      .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyBlock::getCapsule,4090                   "Gets a capsule wrapping the MlirBlock.")4091      .def_prop_ro(4092          "owner",4093          [](PyBlock &self) -> nb::typed<nb::object, PyOpView> {4094            return self.getParentOperation()->createOpView();4095          },4096          "Returns the owning operation of this block.")4097      .def_prop_ro(4098          "region",4099          [](PyBlock &self) {4100            MlirRegion region = mlirBlockGetParentRegion(self.get());4101            return PyRegion(self.getParentOperation(), region);4102          },4103          "Returns the owning region of this block.")4104      .def_prop_ro(4105          "arguments",4106          [](PyBlock &self) {4107            return PyBlockArgumentList(self.getParentOperation(), self.get());4108          },4109          "Returns a list of block arguments.")4110      .def(4111          "add_argument",4112          [](PyBlock &self, const PyType &type, const PyLocation &loc) {4113            return PyBlockArgument(self.getParentOperation(),4114                                   mlirBlockAddArgument(self.get(), type, loc));4115          },4116          "type"_a, "loc"_a,4117          R"(4118            Appends an argument of the specified type to the block.4119 4120            Args:4121              type: The type of the argument to add.4122              loc: The source location for the argument.4123 4124            Returns:4125              The newly added block argument.)")4126      .def(4127          "erase_argument",4128          [](PyBlock &self, unsigned index) {4129            return mlirBlockEraseArgument(self.get(), index);4130          },4131          nb::arg("index"),4132          R"(4133            Erases the argument at the specified index.4134 4135            Args:4136              index: The index of the argument to erase.)")4137      .def_prop_ro(4138          "operations",4139          [](PyBlock &self) {4140            return PyOperationList(self.getParentOperation(), self.get());4141          },4142          "Returns a forward-optimized sequence of operations.")4143      .def_static(4144          "create_at_start",4145          [](PyRegion &parent, const nb::sequence &pyArgTypes,4146             const std::optional<nb::sequence> &pyArgLocs) {4147            parent.checkValid();4148            MlirBlock block = createBlock(pyArgTypes, pyArgLocs);4149            mlirRegionInsertOwnedBlock(parent, 0, block);4150            return PyBlock(parent.getParentOperation(), block);4151          },4152          nb::arg("parent"), nb::arg("arg_types") = nb::list(),4153          nb::arg("arg_locs") = std::nullopt,4154          "Creates and returns a new Block at the beginning of the given "4155          "region (with given argument types and locations).")4156      .def(4157          "append_to",4158          [](PyBlock &self, PyRegion &region) {4159            MlirBlock b = self.get();4160            if (!mlirRegionIsNull(mlirBlockGetParentRegion(b)))4161              mlirBlockDetach(b);4162            mlirRegionAppendOwnedBlock(region.get(), b);4163          },4164          nb::arg("region"),4165          R"(4166            Appends this block to a region.4167 4168            Transfers ownership if the block is currently owned by another region.4169 4170            Args:4171              region: The region to append the block to.)")4172      .def(4173          "create_before",4174          [](PyBlock &self, const nb::args &pyArgTypes,4175             const std::optional<nb::sequence> &pyArgLocs) {4176            self.checkValid();4177            MlirBlock block =4178                createBlock(nb::cast<nb::sequence>(pyArgTypes), pyArgLocs);4179            MlirRegion region = mlirBlockGetParentRegion(self.get());4180            mlirRegionInsertOwnedBlockBefore(region, self.get(), block);4181            return PyBlock(self.getParentOperation(), block);4182          },4183          nb::arg("arg_types"), nb::kw_only(),4184          nb::arg("arg_locs") = std::nullopt,4185          "Creates and returns a new Block before this block "4186          "(with given argument types and locations).")4187      .def(4188          "create_after",4189          [](PyBlock &self, const nb::args &pyArgTypes,4190             const std::optional<nb::sequence> &pyArgLocs) {4191            self.checkValid();4192            MlirBlock block =4193                createBlock(nb::cast<nb::sequence>(pyArgTypes), pyArgLocs);4194            MlirRegion region = mlirBlockGetParentRegion(self.get());4195            mlirRegionInsertOwnedBlockAfter(region, self.get(), block);4196            return PyBlock(self.getParentOperation(), block);4197          },4198          nb::arg("arg_types"), nb::kw_only(),4199          nb::arg("arg_locs") = std::nullopt,4200          "Creates and returns a new Block after this block "4201          "(with given argument types and locations).")4202      .def(4203          "__iter__",4204          [](PyBlock &self) {4205            self.checkValid();4206            MlirOperation firstOperation =4207                mlirBlockGetFirstOperation(self.get());4208            return PyOperationIterator(self.getParentOperation(),4209                                       firstOperation);4210          },4211          "Iterates over operations in the block.")4212      .def(4213          "__eq__",4214          [](PyBlock &self, PyBlock &other) {4215            return self.get().ptr == other.get().ptr;4216          },4217          "Compares two blocks for pointer equality.")4218      .def(4219          "__eq__", [](PyBlock &self, nb::object &other) { return false; },4220          "Compares block with non-block object (always returns False).")4221      .def(4222          "__hash__",4223          [](PyBlock &self) {4224            return static_cast<size_t>(llvm::hash_value(self.get().ptr));4225          },4226          "Returns the hash value of the block.")4227      .def(4228          "__str__",4229          [](PyBlock &self) {4230            self.checkValid();4231            PyPrintAccumulator printAccum;4232            mlirBlockPrint(self.get(), printAccum.getCallback(),4233                           printAccum.getUserData());4234            return printAccum.join();4235          },4236          "Returns the assembly form of the block.")4237      .def(4238          "append",4239          [](PyBlock &self, PyOperationBase &operation) {4240            if (operation.getOperation().isAttached())4241              operation.getOperation().detachFromParent();4242 4243            MlirOperation mlirOperation = operation.getOperation().get();4244            mlirBlockAppendOwnedOperation(self.get(), mlirOperation);4245            operation.getOperation().setAttached(4246                self.getParentOperation().getObject());4247          },4248          nb::arg("operation"),4249          R"(4250            Appends an operation to this block.4251 4252            If the operation is currently in another block, it will be moved.4253 4254            Args:4255              operation: The operation to append to the block.)")4256      .def_prop_ro(4257          "successors",4258          [](PyBlock &self) {4259            return PyBlockSuccessors(self, self.getParentOperation());4260          },4261          "Returns the list of Block successors.")4262      .def_prop_ro(4263          "predecessors",4264          [](PyBlock &self) {4265            return PyBlockPredecessors(self, self.getParentOperation());4266          },4267          "Returns the list of Block predecessors.");4268 4269  //----------------------------------------------------------------------------4270  // Mapping of PyInsertionPoint.4271  //----------------------------------------------------------------------------4272 4273  nb::class_<PyInsertionPoint>(m, "InsertionPoint")4274      .def(nb::init<PyBlock &>(), nb::arg("block"),4275           "Inserts after the last operation but still inside the block.")4276      .def("__enter__", &PyInsertionPoint::contextEnter,4277           "Enters the insertion point as a context manager.")4278      .def("__exit__", &PyInsertionPoint::contextExit,4279           nb::arg("exc_type").none(), nb::arg("exc_value").none(),4280           nb::arg("traceback").none(),4281           "Exits the insertion point context manager.")4282      .def_prop_ro_static(4283          "current",4284          [](nb::object & /*class*/) {4285            auto *ip = PyThreadContextEntry::getDefaultInsertionPoint();4286            if (!ip)4287              throw nb::value_error("No current InsertionPoint");4288            return ip;4289          },4290          nb::sig("def current(/) -> InsertionPoint"),4291          "Gets the InsertionPoint bound to the current thread or raises "4292          "ValueError if none has been set.")4293      .def(nb::init<PyOperationBase &>(), nb::arg("beforeOperation"),4294           "Inserts before a referenced operation.")4295      .def_static("at_block_begin", &PyInsertionPoint::atBlockBegin,4296                  nb::arg("block"),4297                  R"(4298                    Creates an insertion point at the beginning of a block.4299 4300                    Args:4301                      block: The block at whose beginning operations should be inserted.4302 4303                    Returns:4304                      An InsertionPoint at the block's beginning.)")4305      .def_static("at_block_terminator", &PyInsertionPoint::atBlockTerminator,4306                  nb::arg("block"),4307                  R"(4308                    Creates an insertion point before a block's terminator.4309 4310                    Args:4311                      block: The block whose terminator to insert before.4312 4313                    Returns:4314                      An InsertionPoint before the terminator.4315 4316                    Raises:4317                      ValueError: If the block has no terminator.)")4318      .def_static("after", &PyInsertionPoint::after, nb::arg("operation"),4319                  R"(4320                    Creates an insertion point immediately after an operation.4321 4322                    Args:4323                      operation: The operation after which to insert.4324 4325                    Returns:4326                      An InsertionPoint after the operation.)")4327      .def("insert", &PyInsertionPoint::insert, nb::arg("operation"),4328           R"(4329             Inserts an operation at this insertion point.4330 4331             Args:4332               operation: The operation to insert.)")4333      .def_prop_ro(4334          "block", [](PyInsertionPoint &self) { return self.getBlock(); },4335          "Returns the block that this `InsertionPoint` points to.")4336      .def_prop_ro(4337          "ref_operation",4338          [](PyInsertionPoint &self)4339              -> std::optional<nb::typed<nb::object, PyOperation>> {4340            auto refOperation = self.getRefOperation();4341            if (refOperation)4342              return refOperation->getObject();4343            return {};4344          },4345          "The reference operation before which new operations are "4346          "inserted, or None if the insertion point is at the end of "4347          "the block.");4348 4349  //----------------------------------------------------------------------------4350  // Mapping of PyAttribute.4351  //----------------------------------------------------------------------------4352  nb::class_<PyAttribute>(m, "Attribute")4353      // Delegate to the PyAttribute copy constructor, which will also lifetime4354      // extend the backing context which owns the MlirAttribute.4355      .def(nb::init<PyAttribute &>(), nb::arg("cast_from_type"),4356           "Casts the passed attribute to the generic `Attribute`.")4357      .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyAttribute::getCapsule,4358                   "Gets a capsule wrapping the MlirAttribute.")4359      .def_static(4360          MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyAttribute::createFromCapsule,4361          "Creates an Attribute from a capsule wrapping `MlirAttribute`.")4362      .def_static(4363          "parse",4364          [](const std::string &attrSpec, DefaultingPyMlirContext context)4365              -> nb::typed<nb::object, PyAttribute> {4366            PyMlirContext::ErrorCapture errors(context->getRef());4367            MlirAttribute attr = mlirAttributeParseGet(4368                context->get(), toMlirStringRef(attrSpec));4369            if (mlirAttributeIsNull(attr))4370              throw MLIRError("Unable to parse attribute", errors.take());4371            return PyAttribute(context.get()->getRef(), attr).maybeDownCast();4372          },4373          nb::arg("asm"), nb::arg("context") = nb::none(),4374          "Parses an attribute from an assembly form. Raises an `MLIRError` on "4375          "failure.")4376      .def_prop_ro(4377          "context",4378          [](PyAttribute &self) -> nb::typed<nb::object, PyMlirContext> {4379            return self.getContext().getObject();4380          },4381          "Context that owns the `Attribute`.")4382      .def_prop_ro(4383          "type",4384          [](PyAttribute &self) -> nb::typed<nb::object, PyType> {4385            return PyType(self.getContext(), mlirAttributeGetType(self))4386                .maybeDownCast();4387          },4388          "Returns the type of the `Attribute`.")4389      .def(4390          "get_named",4391          [](PyAttribute &self, std::string name) {4392            return PyNamedAttribute(self, std::move(name));4393          },4394          nb::keep_alive<0, 1>(),4395          R"(4396            Binds a name to the attribute, creating a `NamedAttribute`.4397 4398            Args:4399              name: The name to bind to the `Attribute`.4400 4401            Returns:4402              A `NamedAttribute` with the given name and this attribute.)")4403      .def(4404          "__eq__",4405          [](PyAttribute &self, PyAttribute &other) { return self == other; },4406          "Compares two attributes for equality.")4407      .def(4408          "__eq__", [](PyAttribute &self, nb::object &other) { return false; },4409          "Compares attribute with non-attribute object (always returns "4410          "False).")4411      .def(4412          "__hash__",4413          [](PyAttribute &self) {4414            return static_cast<size_t>(llvm::hash_value(self.get().ptr));4415          },4416          "Returns the hash value of the attribute.")4417      .def(4418          "dump", [](PyAttribute &self) { mlirAttributeDump(self); },4419          kDumpDocstring)4420      .def(4421          "__str__",4422          [](PyAttribute &self) {4423            PyPrintAccumulator printAccum;4424            mlirAttributePrint(self, printAccum.getCallback(),4425                               printAccum.getUserData());4426            return printAccum.join();4427          },4428          "Returns the assembly form of the Attribute.")4429      .def(4430          "__repr__",4431          [](PyAttribute &self) {4432            // Generally, assembly formats are not printed for __repr__ because4433            // this can cause exceptionally long debug output and exceptions.4434            // However, attribute values are generally considered useful and4435            // are printed. This may need to be re-evaluated if debug dumps end4436            // up being excessive.4437            PyPrintAccumulator printAccum;4438            printAccum.parts.append("Attribute(");4439            mlirAttributePrint(self, printAccum.getCallback(),4440                               printAccum.getUserData());4441            printAccum.parts.append(")");4442            return printAccum.join();4443          },4444          "Returns a string representation of the attribute.")4445      .def_prop_ro(4446          "typeid",4447          [](PyAttribute &self) {4448            MlirTypeID mlirTypeID = mlirAttributeGetTypeID(self);4449            assert(!mlirTypeIDIsNull(mlirTypeID) &&4450                   "mlirTypeID was expected to be non-null.");4451            return PyTypeID(mlirTypeID);4452          },4453          "Returns the `TypeID` of the attribute.")4454      .def(4455          MLIR_PYTHON_MAYBE_DOWNCAST_ATTR,4456          [](PyAttribute &self) -> nb::typed<nb::object, PyAttribute> {4457            return self.maybeDownCast();4458          },4459          "Downcasts the attribute to a more specific attribute if possible.");4460 4461  //----------------------------------------------------------------------------4462  // Mapping of PyNamedAttribute4463  //----------------------------------------------------------------------------4464  nb::class_<PyNamedAttribute>(m, "NamedAttribute")4465      .def(4466          "__repr__",4467          [](PyNamedAttribute &self) {4468            PyPrintAccumulator printAccum;4469            printAccum.parts.append("NamedAttribute(");4470            printAccum.parts.append(4471                nb::str(mlirIdentifierStr(self.namedAttr.name).data,4472                        mlirIdentifierStr(self.namedAttr.name).length));4473            printAccum.parts.append("=");4474            mlirAttributePrint(self.namedAttr.attribute,4475                               printAccum.getCallback(),4476                               printAccum.getUserData());4477            printAccum.parts.append(")");4478            return printAccum.join();4479          },4480          "Returns a string representation of the named attribute.")4481      .def_prop_ro(4482          "name",4483          [](PyNamedAttribute &self) {4484            return mlirIdentifierStr(self.namedAttr.name);4485          },4486          "The name of the `NamedAttribute` binding.")4487      .def_prop_ro(4488          "attr",4489          [](PyNamedAttribute &self) { return self.namedAttr.attribute; },4490          nb::keep_alive<0, 1>(), nb::sig("def attr(self) -> Attribute"),4491          "The underlying generic attribute of the `NamedAttribute` binding.");4492 4493  //----------------------------------------------------------------------------4494  // Mapping of PyType.4495  //----------------------------------------------------------------------------4496  nb::class_<PyType>(m, "Type")4497      // Delegate to the PyType copy constructor, which will also lifetime4498      // extend the backing context which owns the MlirType.4499      .def(nb::init<PyType &>(), nb::arg("cast_from_type"),4500           "Casts the passed type to the generic `Type`.")4501      .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyType::getCapsule,4502                   "Gets a capsule wrapping the `MlirType`.")4503      .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyType::createFromCapsule,4504                  "Creates a Type from a capsule wrapping `MlirType`.")4505      .def_static(4506          "parse",4507          [](std::string typeSpec,4508             DefaultingPyMlirContext context) -> nb::typed<nb::object, PyType> {4509            PyMlirContext::ErrorCapture errors(context->getRef());4510            MlirType type =4511                mlirTypeParseGet(context->get(), toMlirStringRef(typeSpec));4512            if (mlirTypeIsNull(type))4513              throw MLIRError("Unable to parse type", errors.take());4514            return PyType(context.get()->getRef(), type).maybeDownCast();4515          },4516          nb::arg("asm"), nb::arg("context") = nb::none(),4517          R"(4518            Parses the assembly form of a type.4519 4520            Returns a Type object or raises an `MLIRError` if the type cannot be parsed.4521 4522            See also: https://mlir.llvm.org/docs/LangRef/#type-system)")4523      .def_prop_ro(4524          "context",4525          [](PyType &self) -> nb::typed<nb::object, PyMlirContext> {4526            return self.getContext().getObject();4527          },4528          "Context that owns the `Type`.")4529      .def(4530          "__eq__", [](PyType &self, PyType &other) { return self == other; },4531          "Compares two types for equality.")4532      .def(4533          "__eq__", [](PyType &self, nb::object &other) { return false; },4534          nb::arg("other").none(),4535          "Compares type with non-type object (always returns False).")4536      .def(4537          "__hash__",4538          [](PyType &self) {4539            return static_cast<size_t>(llvm::hash_value(self.get().ptr));4540          },4541          "Returns the hash value of the `Type`.")4542      .def(4543          "dump", [](PyType &self) { mlirTypeDump(self); }, kDumpDocstring)4544      .def(4545          "__str__",4546          [](PyType &self) {4547            PyPrintAccumulator printAccum;4548            mlirTypePrint(self, printAccum.getCallback(),4549                          printAccum.getUserData());4550            return printAccum.join();4551          },4552          "Returns the assembly form of the `Type`.")4553      .def(4554          "__repr__",4555          [](PyType &self) {4556            // Generally, assembly formats are not printed for __repr__ because4557            // this can cause exceptionally long debug output and exceptions.4558            // However, types are an exception as they typically have compact4559            // assembly forms and printing them is useful.4560            PyPrintAccumulator printAccum;4561            printAccum.parts.append("Type(");4562            mlirTypePrint(self, printAccum.getCallback(),4563                          printAccum.getUserData());4564            printAccum.parts.append(")");4565            return printAccum.join();4566          },4567          "Returns a string representation of the `Type`.")4568      .def(4569          MLIR_PYTHON_MAYBE_DOWNCAST_ATTR,4570          [](PyType &self) -> nb::typed<nb::object, PyType> {4571            return self.maybeDownCast();4572          },4573          "Downcasts the Type to a more specific `Type` if possible.")4574      .def_prop_ro(4575          "typeid",4576          [](PyType &self) {4577            MlirTypeID mlirTypeID = mlirTypeGetTypeID(self);4578            if (!mlirTypeIDIsNull(mlirTypeID))4579              return PyTypeID(mlirTypeID);4580            auto origRepr = nb::cast<std::string>(nb::repr(nb::cast(self)));4581            throw nb::value_error(4582                (origRepr + llvm::Twine(" has no typeid.")).str().c_str());4583          },4584          "Returns the `TypeID` of the `Type`, or raises `ValueError` if "4585          "`Type` has no "4586          "`TypeID`.");4587 4588  //----------------------------------------------------------------------------4589  // Mapping of PyTypeID.4590  //----------------------------------------------------------------------------4591  nb::class_<PyTypeID>(m, "TypeID")4592      .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyTypeID::getCapsule,4593                   "Gets a capsule wrapping the `MlirTypeID`.")4594      .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyTypeID::createFromCapsule,4595                  "Creates a `TypeID` from a capsule wrapping `MlirTypeID`.")4596      // Note, this tests whether the underlying TypeIDs are the same,4597      // not whether the wrapper MlirTypeIDs are the same, nor whether4598      // the Python objects are the same (i.e., PyTypeID is a value type).4599      .def(4600          "__eq__",4601          [](PyTypeID &self, PyTypeID &other) { return self == other; },4602          "Compares two `TypeID`s for equality.")4603      .def(4604          "__eq__",4605          [](PyTypeID &self, const nb::object &other) { return false; },4606          "Compares TypeID with non-TypeID object (always returns False).")4607      // Note, this gives the hash value of the underlying TypeID, not the4608      // hash value of the Python object, nor the hash value of the4609      // MlirTypeID wrapper.4610      .def(4611          "__hash__",4612          [](PyTypeID &self) {4613            return static_cast<size_t>(mlirTypeIDHashValue(self));4614          },4615          "Returns the hash value of the `TypeID`.");4616 4617  //----------------------------------------------------------------------------4618  // Mapping of Value.4619  //----------------------------------------------------------------------------4620  m.attr("_T") = nb::type_var("_T", nb::arg("bound") = m.attr("Type"));4621 4622  nb::class_<PyValue>(m, "Value", nb::is_generic(),4623                      nb::sig("class Value(Generic[_T])"))4624      .def(nb::init<PyValue &>(), nb::keep_alive<0, 1>(), nb::arg("value"),4625           "Creates a Value reference from another `Value`.")4626      .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyValue::getCapsule,4627                   "Gets a capsule wrapping the `MlirValue`.")4628      .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyValue::createFromCapsule,4629                  "Creates a `Value` from a capsule wrapping `MlirValue`.")4630      .def_prop_ro(4631          "context",4632          [](PyValue &self) -> nb::typed<nb::object, PyMlirContext> {4633            return self.getParentOperation()->getContext().getObject();4634          },4635          "Context in which the value lives.")4636      .def(4637          "dump", [](PyValue &self) { mlirValueDump(self.get()); },4638          kDumpDocstring)4639      .def_prop_ro(4640          "owner",4641          [](PyValue &self) -> nb::object {4642            MlirValue v = self.get();4643            if (mlirValueIsAOpResult(v)) {4644              assert(mlirOperationEqual(self.getParentOperation()->get(),4645                                        mlirOpResultGetOwner(self.get())) &&4646                     "expected the owner of the value in Python to match "4647                     "that in "4648                     "the IR");4649              return self.getParentOperation().getObject();4650            }4651 4652            if (mlirValueIsABlockArgument(v)) {4653              MlirBlock block = mlirBlockArgumentGetOwner(self.get());4654              return nb::cast(PyBlock(self.getParentOperation(), block));4655            }4656 4657            assert(false && "Value must be a block argument or an op result");4658            return nb::none();4659          },4660          "Returns the owner of the value (`Operation` for results, `Block` "4661          "for "4662          "arguments).")4663      .def_prop_ro(4664          "uses",4665          [](PyValue &self) {4666            return PyOpOperandIterator(mlirValueGetFirstUse(self.get()));4667          },4668          "Returns an iterator over uses of this value.")4669      .def(4670          "__eq__",4671          [](PyValue &self, PyValue &other) {4672            return self.get().ptr == other.get().ptr;4673          },4674          "Compares two values for pointer equality.")4675      .def(4676          "__eq__", [](PyValue &self, nb::object other) { return false; },4677          "Compares value with non-value object (always returns False).")4678      .def(4679          "__hash__",4680          [](PyValue &self) {4681            return static_cast<size_t>(llvm::hash_value(self.get().ptr));4682          },4683          "Returns the hash value of the value.")4684      .def(4685          "__str__",4686          [](PyValue &self) {4687            PyPrintAccumulator printAccum;4688            printAccum.parts.append("Value(");4689            mlirValuePrint(self.get(), printAccum.getCallback(),4690                           printAccum.getUserData());4691            printAccum.parts.append(")");4692            return printAccum.join();4693          },4694          R"(4695            Returns the string form of the value.4696 4697            If the value is a block argument, this is the assembly form of its type and the4698            position in the argument list. If the value is an operation result, this is4699            equivalent to printing the operation that produced it.4700          )")4701      .def(4702          "get_name",4703          [](PyValue &self, bool useLocalScope, bool useNameLocAsPrefix) {4704            PyPrintAccumulator printAccum;4705            MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();4706            if (useLocalScope)4707              mlirOpPrintingFlagsUseLocalScope(flags);4708            if (useNameLocAsPrefix)4709              mlirOpPrintingFlagsPrintNameLocAsPrefix(flags);4710            MlirAsmState valueState =4711                mlirAsmStateCreateForValue(self.get(), flags);4712            mlirValuePrintAsOperand(self.get(), valueState,4713                                    printAccum.getCallback(),4714                                    printAccum.getUserData());4715            mlirOpPrintingFlagsDestroy(flags);4716            mlirAsmStateDestroy(valueState);4717            return printAccum.join();4718          },4719          nb::arg("use_local_scope") = false,4720          nb::arg("use_name_loc_as_prefix") = false,4721          R"(4722            Returns the string form of value as an operand.4723 4724            Args:4725              use_local_scope: Whether to use local scope for naming.4726              use_name_loc_as_prefix: Whether to use the location attribute (NameLoc) as prefix.4727 4728            Returns:4729              The value's name as it appears in IR (e.g., `%0`, `%arg0`).)")4730      .def(4731          "get_name",4732          [](PyValue &self, PyAsmState &state) {4733            PyPrintAccumulator printAccum;4734            MlirAsmState valueState = state.get();4735            mlirValuePrintAsOperand(self.get(), valueState,4736                                    printAccum.getCallback(),4737                                    printAccum.getUserData());4738            return printAccum.join();4739          },4740          nb::arg("state"),4741          "Returns the string form of value as an operand (i.e., the ValueID).")4742      .def_prop_ro(4743          "type",4744          [](PyValue &self) -> nb::typed<nb::object, PyType> {4745            return PyType(self.getParentOperation()->getContext(),4746                          mlirValueGetType(self.get()))4747                .maybeDownCast();4748          },4749          "Returns the type of the value.")4750      .def(4751          "set_type",4752          [](PyValue &self, const PyType &type) {4753            mlirValueSetType(self.get(), type);4754          },4755          nb::arg("type"), "Sets the type of the value.",4756          nb::sig("def set_type(self, type: _T)"))4757      .def(4758          "replace_all_uses_with",4759          [](PyValue &self, PyValue &with) {4760            mlirValueReplaceAllUsesOfWith(self.get(), with.get());4761          },4762          "Replace all uses of value with the new value, updating anything in "4763          "the IR that uses `self` to use the other value instead.")4764      .def(4765          "replace_all_uses_except",4766          [](PyValue &self, PyValue &with, PyOperation &exception) {4767            MlirOperation exceptedUser = exception.get();4768            mlirValueReplaceAllUsesExcept(self, with, 1, &exceptedUser);4769          },4770          nb::arg("with_"), nb::arg("exceptions"),4771          kValueReplaceAllUsesExceptDocstring)4772      .def(4773          "replace_all_uses_except",4774          [](PyValue &self, PyValue &with, const nb::list &exceptions) {4775            // Convert Python list to a SmallVector of MlirOperations4776            llvm::SmallVector<MlirOperation> exceptionOps;4777            for (nb::handle exception : exceptions) {4778              exceptionOps.push_back(nb::cast<PyOperation &>(exception).get());4779            }4780 4781            mlirValueReplaceAllUsesExcept(4782                self, with, static_cast<intptr_t>(exceptionOps.size()),4783                exceptionOps.data());4784          },4785          nb::arg("with_"), nb::arg("exceptions"),4786          kValueReplaceAllUsesExceptDocstring)4787      .def(4788          "replace_all_uses_except",4789          [](PyValue &self, PyValue &with, PyOperation &exception) {4790            MlirOperation exceptedUser = exception.get();4791            mlirValueReplaceAllUsesExcept(self, with, 1, &exceptedUser);4792          },4793          nb::arg("with_"), nb::arg("exceptions"),4794          kValueReplaceAllUsesExceptDocstring)4795      .def(4796          "replace_all_uses_except",4797          [](PyValue &self, PyValue &with,4798             std::vector<PyOperation> &exceptions) {4799            // Convert Python list to a SmallVector of MlirOperations4800            llvm::SmallVector<MlirOperation> exceptionOps;4801            for (PyOperation &exception : exceptions)4802              exceptionOps.push_back(exception);4803            mlirValueReplaceAllUsesExcept(4804                self, with, static_cast<intptr_t>(exceptionOps.size()),4805                exceptionOps.data());4806          },4807          nb::arg("with_"), nb::arg("exceptions"),4808          kValueReplaceAllUsesExceptDocstring)4809      .def(4810          MLIR_PYTHON_MAYBE_DOWNCAST_ATTR,4811          [](PyValue &self) -> nb::typed<nb::object, PyValue> {4812            return self.maybeDownCast();4813          },4814          "Downcasts the `Value` to a more specific kind if possible.")4815      .def_prop_ro(4816          "location",4817          [](MlirValue self) {4818            return PyLocation(4819                PyMlirContext::forContext(mlirValueGetContext(self)),4820                mlirValueGetLocation(self));4821          },4822          "Returns the source location of the value.");4823 4824  PyBlockArgument::bind(m);4825  PyOpResult::bind(m);4826  PyOpOperand::bind(m);4827 4828  nb::class_<PyAsmState>(m, "AsmState")4829      .def(nb::init<PyValue &, bool>(), nb::arg("value"),4830           nb::arg("use_local_scope") = false,4831           R"(4832             Creates an `AsmState` for consistent SSA value naming.4833 4834             Args:4835               value: The value to create state for.4836               use_local_scope: Whether to use local scope for naming.)")4837      .def(nb::init<PyOperationBase &, bool>(), nb::arg("op"),4838           nb::arg("use_local_scope") = false,4839           R"(4840             Creates an AsmState for consistent SSA value naming.4841 4842             Args:4843               op: The operation to create state for.4844               use_local_scope: Whether to use local scope for naming.)");4845 4846  //----------------------------------------------------------------------------4847  // Mapping of SymbolTable.4848  //----------------------------------------------------------------------------4849  nb::class_<PySymbolTable>(m, "SymbolTable")4850      .def(nb::init<PyOperationBase &>(),4851           R"(4852             Creates a symbol table for an operation.4853 4854             Args:4855               operation: The `Operation` that defines a symbol table (e.g., a `ModuleOp`).4856 4857             Raises:4858               TypeError: If the operation is not a symbol table.)")4859      .def(4860          "__getitem__",4861          [](PySymbolTable &self,4862             const std::string &name) -> nb::typed<nb::object, PyOpView> {4863            return self.dunderGetItem(name);4864          },4865          R"(4866            Looks up a symbol by name in the symbol table.4867 4868            Args:4869              name: The name of the symbol to look up.4870 4871            Returns:4872              The operation defining the symbol.4873 4874            Raises:4875              KeyError: If the symbol is not found.)")4876      .def("insert", &PySymbolTable::insert, nb::arg("operation"),4877           R"(4878             Inserts a symbol operation into the symbol table.4879 4880             Args:4881               operation: An operation with a symbol name to insert.4882 4883             Returns:4884               The symbol name attribute of the inserted operation.4885 4886             Raises:4887               ValueError: If the operation does not have a symbol name.)")4888      .def("erase", &PySymbolTable::erase, nb::arg("operation"),4889           R"(4890             Erases a symbol operation from the symbol table.4891 4892             Args:4893               operation: The symbol operation to erase.4894 4895             Note:4896               The operation is also erased from the IR and invalidated.)")4897      .def("__delitem__", &PySymbolTable::dunderDel,4898           "Deletes a symbol by name from the symbol table.")4899      .def(4900          "__contains__",4901          [](PySymbolTable &table, const std::string &name) {4902            return !mlirOperationIsNull(mlirSymbolTableLookup(4903                table, mlirStringRefCreate(name.data(), name.length())));4904          },4905          "Checks if a symbol with the given name exists in the table.")4906      // Static helpers.4907      .def_static("set_symbol_name", &PySymbolTable::setSymbolName,4908                  nb::arg("symbol"), nb::arg("name"),4909                  "Sets the symbol name for a symbol operation.")4910      .def_static("get_symbol_name", &PySymbolTable::getSymbolName,4911                  nb::arg("symbol"),4912                  "Gets the symbol name from a symbol operation.")4913      .def_static("get_visibility", &PySymbolTable::getVisibility,4914                  nb::arg("symbol"),4915                  "Gets the visibility attribute of a symbol operation.")4916      .def_static("set_visibility", &PySymbolTable::setVisibility,4917                  nb::arg("symbol"), nb::arg("visibility"),4918                  "Sets the visibility attribute of a symbol operation.")4919      .def_static("replace_all_symbol_uses",4920                  &PySymbolTable::replaceAllSymbolUses, nb::arg("old_symbol"),4921                  nb::arg("new_symbol"), nb::arg("from_op"),4922                  "Replaces all uses of a symbol with a new symbol name within "4923                  "the given operation.")4924      .def_static("walk_symbol_tables", &PySymbolTable::walkSymbolTables,4925                  nb::arg("from_op"), nb::arg("all_sym_uses_visible"),4926                  nb::arg("callback"),4927                  "Walks symbol tables starting from an operation with a "4928                  "callback function.");4929 4930  // Container bindings.4931  PyBlockArgumentList::bind(m);4932  PyBlockIterator::bind(m);4933  PyBlockList::bind(m);4934  PyBlockSuccessors::bind(m);4935  PyBlockPredecessors::bind(m);4936  PyOperationIterator::bind(m);4937  PyOperationList::bind(m);4938  PyOpAttributeMap::bind(m);4939  PyOpOperandIterator::bind(m);4940  PyOpOperandList::bind(m);4941  PyOpResultList::bind(m);4942  PyOpSuccessors::bind(m);4943  PyRegionIterator::bind(m);4944  PyRegionList::bind(m);4945 4946  // Debug bindings.4947  PyGlobalDebugFlag::bind(m);4948 4949  // Attribute builder getter.4950  PyAttrBuilderMap::bind(m);4951 4952  nb::register_exception_translator([](const std::exception_ptr &p,4953                                       void *payload) {4954    // We can't define exceptions with custom fields through pybind, so instead4955    // the exception class is defined in python and imported here.4956    try {4957      if (p)4958        std::rethrow_exception(p);4959    } catch (const MLIRError &e) {4960      nb::object obj = nb::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))4961                           .attr("MLIRError")(e.message, e.errorDiagnostics);4962      PyErr_SetObject(PyExc_Exception, obj.ptr());4963    }4964  });4965}4966