Skip to content

[mlir][python] source line info #149166

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 21 additions & 7 deletions mlir/lib/Bindings/Python/IRCore.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include "Globals.h"
#include "IRModule.h"
#include "NanobindUtils.h"
#include "Traceback.h"
#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
#include "mlir-c/BuiltinAttributes.h"
#include "mlir-c/Debug.h"
Expand Down Expand Up @@ -1523,7 +1524,7 @@ nb::object PyOperation::create(std::string_view name,
llvm::ArrayRef<MlirValue> operands,
std::optional<nb::dict> attributes,
std::optional<std::vector<PyBlock *>> successors,
int regions, DefaultingPyLocation location,
int regions, PyLocation location,
const nb::object &maybeIp, bool inferType) {
llvm::SmallVector<MlirType, 4> mlirResults;
llvm::SmallVector<MlirBlock, 4> mlirSuccessors;
Expand Down Expand Up @@ -1627,7 +1628,7 @@ nb::object PyOperation::create(std::string_view name,
if (!operation.ptr)
throw nb::value_error("Operation creation failed");
PyOperationRef created =
PyOperation::createDetached(location->getContext(), operation);
PyOperation::createDetached(location.getContext(), operation);
maybeInsertOperation(created, maybeIp);

return created.getObject();
Expand Down Expand Up @@ -1937,9 +1938,9 @@ nb::object PyOpView::buildGeneric(
std::optional<nb::list> resultTypeList, nb::list operandList,
std::optional<nb::dict> attributes,
std::optional<std::vector<PyBlock *>> successors,
std::optional<int> regions, DefaultingPyLocation location,
std::optional<int> regions, PyLocation location,
const nb::object &maybeIp) {
PyMlirContextRef context = location->getContext();
PyMlirContextRef context = location.getContext();

// Class level operation construction metadata.
// Operand and result segment specs are either none, which does no
Expand Down Expand Up @@ -3456,6 +3457,13 @@ void mlir::python::populateIRCore(nb::module_ &m) {
std::optional<std::vector<PyBlock *>> successors, int regions,
DefaultingPyLocation location, const nb::object &maybeIp,
bool inferType) {
//////////////
std::optional<Traceback> tb = Traceback::Get();
PyMlirContextRef ctx = location->getContext();
auto loc = tb->tracebackToLocation(ctx->get());
PyLocation pyLoc{ctx, loc};
//////////////

// Unpack/validate operands.
llvm::SmallVector<MlirValue, 4> mlirOperands;
if (operands) {
Expand All @@ -3468,7 +3476,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
}

return PyOperation::create(name, results, mlirOperands, attributes,
successors, regions, location, maybeIp,
successors, regions, pyLoc, maybeIp,
inferType);
},
nb::arg("name"), nb::arg("results").none() = nb::none(),
Expand Down Expand Up @@ -3517,7 +3525,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
new (self) PyOpView(PyOpView::buildGeneric(
name, opRegionSpec, operandSegmentSpecObj,
resultSegmentSpecObj, resultTypeList, operandList,
attributes, successors, regions, location, maybeIp));
attributes, successors, regions, *location.get(), maybeIp));
},
nb::arg("name"), nb::arg("opRegionSpec"),
nb::arg("operandSegmentSpecObj").none() = nb::none(),
Expand Down Expand Up @@ -3553,6 +3561,12 @@ void mlir::python::populateIRCore(nb::module_ &m) {
std::optional<std::vector<PyBlock *>> successors,
std::optional<int> regions, DefaultingPyLocation location,
const nb::object &maybeIp) {
//////////////
std::optional<Traceback> tb = Traceback::Get();
PyMlirContextRef ctx = location->getContext();
auto loc = tb->tracebackToLocation(ctx->get());
PyLocation pyLoc{ctx, loc};
//////////////
std::string name = nb::cast<std::string>(cls.attr("OPERATION_NAME"));
std::tuple<int, bool> opRegionSpec =
nb::cast<std::tuple<int, bool>>(cls.attr("_ODS_REGIONS"));
Expand All @@ -3561,7 +3575,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
return PyOpView::buildGeneric(name, opRegionSpec, operandSegmentSpec,
resultSegmentSpec, resultTypeList,
operandList, attributes, successors,
regions, location, maybeIp);
regions, pyLoc, maybeIp);
},
nb::arg("cls"), nb::arg("results").none() = nb::none(),
nb::arg("operands").none() = nb::none(),
Expand Down
4 changes: 2 additions & 2 deletions mlir/lib/Bindings/Python/IRModule.h
Original file line number Diff line number Diff line change
Expand Up @@ -722,7 +722,7 @@ class PyOperation : public PyOperationBase, public BaseContextObject {
llvm::ArrayRef<MlirValue> operands,
std::optional<nanobind::dict> attributes,
std::optional<std::vector<PyBlock *>> successors, int regions,
DefaultingPyLocation location, const nanobind::object &ip,
PyLocation location, const nanobind::object &ip,
bool inferType);

/// Creates an OpView suitable for this operation.
Expand Down Expand Up @@ -781,7 +781,7 @@ class PyOpView : public PyOperationBase {
nanobind::list operandList,
std::optional<nanobind::dict> attributes,
std::optional<std::vector<PyBlock *>> successors,
std::optional<int> regions, DefaultingPyLocation location,
std::optional<int> regions, PyLocation location,
const nanobind::object &maybeIp);

/// Construct an instance of a class deriving from OpView, bypassing its
Expand Down
2 changes: 2 additions & 0 deletions mlir/lib/Bindings/Python/MainModule.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
#include "NanobindUtils.h"
#include "Pass.h"
#include "Rewrite.h"
#include "Traceback.h"
#include "mlir/Bindings/Python/Nanobind.h"

namespace nb = nanobind;
Expand Down Expand Up @@ -105,6 +106,7 @@ NB_MODULE(_mlir, m) {
"typeid"_a, nb::kw_only(), "replace"_a = false,
"Register a value caster for casting MLIR values to custom user values.");

BuildTracebackSubmodule(m);
// Define and populate IR submodule.
auto irModule = m.def_submodule("ir", "MLIR IR Bindings");
populateIRCore(irModule);
Expand Down
Loading
Loading