From a40f47c9725f3aed6752f050f3010f3701d0cff7 Mon Sep 17 00:00:00 2001 From: Maksim Levental Date: Tue, 12 Aug 2025 17:59:59 -0400 Subject: [PATCH] [mlir][python] automatic location inference (#151246) This PR implements "automatic" location inference in the bindings. The way it works is it walks the frame stack collecting source locations (Python captures these in the frame itself). It is inspired by JAX's [implementation](https://github.com/jax-ml/jax/blob/523ddcfbcad005deab5a7d542df4c706f5ee5e9c/jax/_src/interpreters/mlir.py#L462) but moves the frame stack traversal into the bindings for better performance. The system supports registering "included" and "excluded" filenames; frames originating from functions in included filenames **will not** be filtered and frames originating from functions in excluded filenames **will** be filtered (in that order). This allows excluding all the generated `*_ops_gen.py` files. The system is also "toggleable" and off by default to save people who have their own systems (such as JAX) from the added cost. Note, the system stores the entire stacktrace (subject to `locTracebackFramesLimit`) in the `Location` using specifically a `CallSiteLoc`. This can be useful for profiling tools (flamegraphs etc.). Shoutout to the folks at JAX for coming up with a good system. --------- Co-authored-by: Jacques Pienaar --- mlir/lib/Bindings/Python/Globals.h | 39 ++++++ mlir/lib/Bindings/Python/IRCore.cpp | 122 +++++++++++++++--- mlir/lib/Bindings/Python/IRModule.cpp | 70 +++++++++- mlir/lib/Bindings/Python/IRModule.h | 15 +-- mlir/lib/Bindings/Python/MainModule.cpp | 23 +++- mlir/python/mlir/dialects/_ods_common.py | 8 +- mlir/test/lit.cfg.py | 3 + mlir/test/python/ir/auto_location.py | 101 +++++++++++++++ mlir/test/python/ir/context_managers.py | 22 +--- mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp | 1 + 10 files changed, 347 insertions(+), 57 deletions(-) create mode 100644 mlir/test/python/ir/auto_location.py diff --git a/mlir/lib/Bindings/Python/Globals.h b/mlir/lib/Bindings/Python/Globals.h index 826a34a53517..71a051cb3d9f 100644 --- a/mlir/lib/Bindings/Python/Globals.h +++ b/mlir/lib/Bindings/Python/Globals.h @@ -10,15 +10,19 @@ #define MLIR_BINDINGS_PYTHON_GLOBALS_H #include +#include #include +#include #include #include "NanobindUtils.h" #include "mlir-c/IR.h" #include "mlir/CAPI/Support.h" #include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/StringExtras.h" #include "llvm/ADT/StringRef.h" #include "llvm/ADT/StringSet.h" +#include "llvm/Support/Regex.h" namespace mlir { namespace python { @@ -114,6 +118,39 @@ public: std::optional lookupOperationClass(llvm::StringRef operationName); + class TracebackLoc { + public: + bool locTracebacksEnabled(); + + void setLocTracebacksEnabled(bool value); + + size_t locTracebackFramesLimit(); + + void setLocTracebackFramesLimit(size_t value); + + void registerTracebackFileInclusion(const std::string &file); + + void registerTracebackFileExclusion(const std::string &file); + + bool isUserTracebackFilename(llvm::StringRef file); + + static constexpr size_t kMaxFrames = 512; + + private: + nanobind::ft_mutex mutex; + bool locTracebackEnabled_ = false; + size_t locTracebackFramesLimit_ = 10; + std::unordered_set userTracebackIncludeFiles; + std::unordered_set userTracebackExcludeFiles; + std::regex userTracebackIncludeRegex; + bool rebuildUserTracebackIncludeRegex = false; + std::regex userTracebackExcludeRegex; + bool rebuildUserTracebackExcludeRegex = false; + llvm::StringMap isUserTracebackFilenameCache; + }; + + TracebackLoc &getTracebackLoc() { return tracebackLoc; } + private: static PyGlobals *instance; @@ -134,6 +171,8 @@ private: /// Set of dialect namespaces that we have attempted to import implementation /// modules for. llvm::StringSet<> loadedDialectModules; + + TracebackLoc tracebackLoc; }; } // namespace python diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp index ee88aa4753d0..1aec7ffe8e08 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -20,11 +20,8 @@ #include "nanobind/nanobind.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/SmallVector.h" -#include "llvm/Support/raw_ostream.h" #include -#include -#include namespace nb = nanobind; using namespace nb::literals; @@ -1523,7 +1520,7 @@ nb::object PyOperation::create(std::string_view name, llvm::ArrayRef operands, std::optional attributes, std::optional> successors, - int regions, DefaultingPyLocation location, + int regions, PyLocation &location, const nb::object &maybeIp, bool inferType) { llvm::SmallVector mlirResults; llvm::SmallVector mlirSuccessors; @@ -1627,7 +1624,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(); @@ -1937,9 +1934,9 @@ nb::object PyOpView::buildGeneric( std::optional resultTypeList, nb::list operandList, std::optional attributes, std::optional> successors, - std::optional regions, DefaultingPyLocation location, + std::optional 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 @@ -2789,6 +2786,90 @@ private: PyOperationRef operation; }; +MlirLocation tracebackToLocation(MlirContext ctx) { + size_t framesLimit = + PyGlobals::get().getTracebackLoc().locTracebackFramesLimit(); + // Use a thread_local here to avoid requiring a large amount of space. + thread_local std::array + frames; + size_t count = 0; + + nb::gil_scoped_acquire acquire; + PyThreadState *tstate = PyThreadState_GET(); + PyFrameObject *next; + PyFrameObject *pyFrame = PyThreadState_GetFrame(tstate); + // In the increment expression: + // 1. get the next prev frame; + // 2. decrement the ref count on the current frame (in order that it can get + // gc'd, along with any objects in its closure and etc); + // 3. set current = next. + for (; pyFrame != nullptr && count < framesLimit; + next = PyFrame_GetBack(pyFrame), Py_XDECREF(pyFrame), pyFrame = next) { + PyCodeObject *code = PyFrame_GetCode(pyFrame); + auto fileNameStr = + nb::cast(nb::borrow(code->co_filename)); + llvm::StringRef fileName(fileNameStr); + if (!PyGlobals::get().getTracebackLoc().isUserTracebackFilename(fileName)) + continue; + +#if PY_VERSION_HEX < 0x030b00f0 + std::string name = + nb::cast(nb::borrow(code->co_name)); + llvm::StringRef funcName(name); + int startLine = PyFrame_GetLineNumber(pyFrame); + MlirLocation loc = + mlirLocationFileLineColGet(ctx, wrap(fileName), startLine, 0); +#else + // co_qualname and PyCode_Addr2Location added in py3.11 + std::string name = + nb::cast(nb::borrow(code->co_qualname)); + llvm::StringRef funcName(name); + int startLine, startCol, endLine, endCol; + int lasti = PyFrame_GetLasti(pyFrame); + if (!PyCode_Addr2Location(code, lasti, &startLine, &startCol, &endLine, + &endCol)) { + throw nb::python_error(); + } + MlirLocation loc = mlirLocationFileLineColRangeGet( + ctx, wrap(fileName), startLine, startCol, endLine, endCol); +#endif + + frames[count] = mlirLocationNameGet(ctx, wrap(funcName), loc); + ++count; + } + // When the loop breaks (after the last iter), current frame (if non-null) + // is leaked without this. + Py_XDECREF(pyFrame); + + if (count == 0) + return mlirLocationUnknownGet(ctx); + + MlirLocation callee = frames[0]; + assert(!mlirLocationIsNull(callee) && "expected non-null callee location"); + if (count == 1) + return callee; + + MlirLocation caller = frames[count - 1]; + assert(!mlirLocationIsNull(caller) && "expected non-null caller location"); + for (int i = count - 2; i >= 1; i--) + caller = mlirLocationCallSiteGet(frames[i], caller); + + return mlirLocationCallSiteGet(callee, caller); +} + +PyLocation +maybeGetTracebackLocation(const std::optional &location) { + if (location.has_value()) + return location.value(); + if (!PyGlobals::get().getTracebackLoc().locTracebacksEnabled()) + return DefaultingPyLocation::resolve(); + + PyMlirContext &ctx = DefaultingPyMlirContext::resolve(); + MlirLocation mlirLoc = tracebackToLocation(ctx.get()); + PyMlirContextRef ref = PyMlirContext::forContext(ctx.get()); + return {ref, mlirLoc}; +} + } // namespace //------------------------------------------------------------------------------ @@ -3052,10 +3133,10 @@ void mlir::python::populateIRCore(nb::module_ &m) { .def("__eq__", [](PyLocation &self, nb::object other) { return false; }) .def_prop_ro_static( "current", - [](nb::object & /*class*/) { + [](nb::object & /*class*/) -> std::optional { auto *loc = PyThreadContextEntry::getDefaultLocation(); if (!loc) - throw nb::value_error("No current Location"); + return std::nullopt; return loc; }, "Gets the Location bound to the current thread or raises ValueError") @@ -3240,8 +3321,9 @@ void mlir::python::populateIRCore(nb::module_ &m) { kModuleParseDocstring) .def_static( "create", - [](DefaultingPyLocation loc) { - MlirModule module = mlirModuleCreateEmpty(loc); + [](const std::optional &loc) { + PyLocation pyLoc = maybeGetTracebackLocation(loc); + MlirModule module = mlirModuleCreateEmpty(pyLoc.get()); return PyModule::forModule(module).releaseObject(); }, nb::arg("loc").none() = nb::none(), "Creates an empty module") @@ -3462,8 +3544,8 @@ void mlir::python::populateIRCore(nb::module_ &m) { std::optional> operands, std::optional attributes, std::optional> successors, int regions, - DefaultingPyLocation location, const nb::object &maybeIp, - bool inferType) { + const std::optional &location, + const nb::object &maybeIp, bool inferType) { // Unpack/validate operands. llvm::SmallVector mlirOperands; if (operands) { @@ -3475,8 +3557,9 @@ void mlir::python::populateIRCore(nb::module_ &m) { } } + PyLocation pyLoc = maybeGetTracebackLocation(location); 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(), @@ -3520,12 +3603,14 @@ void mlir::python::populateIRCore(nb::module_ &m) { std::optional resultTypeList, nb::list operandList, std::optional attributes, std::optional> successors, - std::optional regions, DefaultingPyLocation location, + std::optional regions, + const std::optional &location, const nb::object &maybeIp) { + PyLocation pyLoc = maybeGetTracebackLocation(location); new (self) PyOpView(PyOpView::buildGeneric( name, opRegionSpec, operandSegmentSpecObj, resultSegmentSpecObj, resultTypeList, operandList, - attributes, successors, regions, location, maybeIp)); + attributes, successors, regions, pyLoc, maybeIp)); }, nb::arg("name"), nb::arg("opRegionSpec"), nb::arg("operandSegmentSpecObj").none() = nb::none(), @@ -3559,17 +3644,18 @@ void mlir::python::populateIRCore(nb::module_ &m) { [](nb::handle cls, std::optional resultTypeList, nb::list operandList, std::optional attributes, std::optional> successors, - std::optional regions, DefaultingPyLocation location, + std::optional regions, std::optional location, const nb::object &maybeIp) { std::string name = nb::cast(cls.attr("OPERATION_NAME")); std::tuple opRegionSpec = nb::cast>(cls.attr("_ODS_REGIONS")); nb::object operandSegmentSpec = cls.attr("_ODS_OPERAND_SEGMENTS"); nb::object resultSegmentSpec = cls.attr("_ODS_RESULT_SEGMENTS"); + PyLocation pyLoc = maybeGetTracebackLocation(location); 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(), diff --git a/mlir/lib/Bindings/Python/IRModule.cpp b/mlir/lib/Bindings/Python/IRModule.cpp index e600f1bbd449..0de2f1711829 100644 --- a/mlir/lib/Bindings/Python/IRModule.cpp +++ b/mlir/lib/Bindings/Python/IRModule.cpp @@ -13,9 +13,9 @@ #include "Globals.h" #include "NanobindUtils.h" +#include "mlir-c/Bindings/Python/Interop.h" #include "mlir-c/Support.h" #include "mlir/Bindings/Python/Nanobind.h" -#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind. namespace nb = nanobind; using namespace mlir; @@ -197,3 +197,71 @@ PyGlobals::lookupOperationClass(llvm::StringRef operationName) { // Not found and loading did not yield a registration. return std::nullopt; } + +bool PyGlobals::TracebackLoc::locTracebacksEnabled() { + nanobind::ft_lock_guard lock(mutex); + return locTracebackEnabled_; +} + +void PyGlobals::TracebackLoc::setLocTracebacksEnabled(bool value) { + nanobind::ft_lock_guard lock(mutex); + locTracebackEnabled_ = value; +} + +size_t PyGlobals::TracebackLoc::locTracebackFramesLimit() { + nanobind::ft_lock_guard lock(mutex); + return locTracebackFramesLimit_; +} + +void PyGlobals::TracebackLoc::setLocTracebackFramesLimit(size_t value) { + nanobind::ft_lock_guard lock(mutex); + locTracebackFramesLimit_ = std::min(value, kMaxFrames); +} + +void PyGlobals::TracebackLoc::registerTracebackFileInclusion( + const std::string &file) { + nanobind::ft_lock_guard lock(mutex); + auto reg = "^" + llvm::Regex::escape(file); + if (userTracebackIncludeFiles.insert(reg).second) + rebuildUserTracebackIncludeRegex = true; + if (userTracebackExcludeFiles.count(reg)) { + if (userTracebackExcludeFiles.erase(reg)) + rebuildUserTracebackExcludeRegex = true; + } +} + +void PyGlobals::TracebackLoc::registerTracebackFileExclusion( + const std::string &file) { + nanobind::ft_lock_guard lock(mutex); + auto reg = "^" + llvm::Regex::escape(file); + if (userTracebackExcludeFiles.insert(reg).second) + rebuildUserTracebackExcludeRegex = true; + if (userTracebackIncludeFiles.count(reg)) { + if (userTracebackIncludeFiles.erase(reg)) + rebuildUserTracebackIncludeRegex = true; + } +} + +bool PyGlobals::TracebackLoc::isUserTracebackFilename( + const llvm::StringRef file) { + nanobind::ft_lock_guard lock(mutex); + if (rebuildUserTracebackIncludeRegex) { + userTracebackIncludeRegex.assign( + llvm::join(userTracebackIncludeFiles, "|")); + rebuildUserTracebackIncludeRegex = false; + isUserTracebackFilenameCache.clear(); + } + if (rebuildUserTracebackExcludeRegex) { + userTracebackExcludeRegex.assign( + llvm::join(userTracebackExcludeFiles, "|")); + rebuildUserTracebackExcludeRegex = false; + isUserTracebackFilenameCache.clear(); + } + if (!isUserTracebackFilenameCache.contains(file)) { + std::string fileStr = file.str(); + bool include = std::regex_search(fileStr, userTracebackIncludeRegex); + bool exclude = std::regex_search(fileStr, userTracebackExcludeRegex); + isUserTracebackFilenameCache[file] = include || !exclude; + } + return isUserTracebackFilenameCache[file]; +} diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h index 9c22dea157c0..fa16ae3ce329 100644 --- a/mlir/lib/Bindings/Python/IRModule.h +++ b/mlir/lib/Bindings/Python/IRModule.h @@ -192,16 +192,6 @@ public: PyMlirContext(const PyMlirContext &) = delete; PyMlirContext(PyMlirContext &&) = delete; - /// For the case of a python __init__ (nanobind::init) method, pybind11 is - /// quite strict about needing to return a pointer that is not yet associated - /// to an nanobind::object. Since the forContext() method acts like a pool, - /// possibly returning a recycled context, it does not satisfy this need. The - /// usual way in python to accomplish such a thing is to override __new__, but - /// that is also not supported by pybind11. Instead, we use this entry - /// point which always constructs a fresh context (which cannot alias an - /// existing one because it is fresh). - static PyMlirContext *createNewContextForInit(); - /// Returns a context reference for the singleton PyMlirContext wrapper for /// the given context. static PyMlirContextRef forContext(MlirContext context); @@ -722,8 +712,7 @@ public: llvm::ArrayRef operands, std::optional attributes, std::optional> successors, int regions, - DefaultingPyLocation location, const nanobind::object &ip, - bool inferType); + PyLocation &location, const nanobind::object &ip, bool inferType); /// Creates an OpView suitable for this operation. nanobind::object createOpView(); @@ -781,7 +770,7 @@ public: nanobind::list operandList, std::optional attributes, std::optional> successors, - std::optional regions, DefaultingPyLocation location, + std::optional regions, PyLocation &location, const nanobind::object &maybeIp); /// Construct an instance of a class deriving from OpView, bypassing its diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp index 6f4943100660..278847e7ac7f 100644 --- a/mlir/lib/Bindings/Python/MainModule.cpp +++ b/mlir/lib/Bindings/Python/MainModule.cpp @@ -6,7 +6,6 @@ // //===----------------------------------------------------------------------===// - #include "Globals.h" #include "IRModule.h" #include "NanobindUtils.h" @@ -44,7 +43,27 @@ NB_MODULE(_mlir, m) { .def("_register_operation_impl", &PyGlobals::registerOperationImpl, "operation_name"_a, "operation_class"_a, nb::kw_only(), "replace"_a = false, - "Testing hook for directly registering an operation"); + "Testing hook for directly registering an operation") + .def("loc_tracebacks_enabled", + [](PyGlobals &self) { + return self.getTracebackLoc().locTracebacksEnabled(); + }) + .def("set_loc_tracebacks_enabled", + [](PyGlobals &self, bool enabled) { + self.getTracebackLoc().setLocTracebacksEnabled(enabled); + }) + .def("set_loc_tracebacks_frame_limit", + [](PyGlobals &self, int n) { + self.getTracebackLoc().setLocTracebackFramesLimit(n); + }) + .def("register_traceback_file_inclusion", + [](PyGlobals &self, const std::string &filename) { + self.getTracebackLoc().registerTracebackFileInclusion(filename); + }) + .def("register_traceback_file_exclusion", + [](PyGlobals &self, const std::string &filename) { + self.getTracebackLoc().registerTracebackFileExclusion(filename); + }); // Aside from making the globals accessible to python, having python manage // it is necessary to make sure it is destroyed (and releases its python diff --git a/mlir/python/mlir/dialects/_ods_common.py b/mlir/python/mlir/dialects/_ods_common.py index a5efa057cfb6..10abd06ff266 100644 --- a/mlir/python/mlir/dialects/_ods_common.py +++ b/mlir/python/mlir/dialects/_ods_common.py @@ -78,12 +78,12 @@ def equally_sized_accessor( def get_default_loc_context(location=None): """ Returns a context in which the defaulted location is created. If the location - is None, takes the current location from the stack, raises ValueError if there - is no location on the stack. + is None, takes the current location from the stack. """ if location is None: - # Location.current raises ValueError if there is no current location. - return _cext.ir.Location.current.context + if _cext.ir.Location.current: + return _cext.ir.Location.current.context + return None return location.context diff --git a/mlir/test/lit.cfg.py b/mlir/test/lit.cfg.py index f392bdacadd3..ba7eeeed8ef3 100644 --- a/mlir/test/lit.cfg.py +++ b/mlir/test/lit.cfg.py @@ -378,3 +378,6 @@ if config.run_rocm_tests: if config.arm_emulator_executable: config.available_features.add("arm-emulator") + +if sys.version_info >= (3, 11): + config.available_features.add("python-ge-311") diff --git a/mlir/test/python/ir/auto_location.py b/mlir/test/python/ir/auto_location.py new file mode 100644 index 000000000000..c2d51083c137 --- /dev/null +++ b/mlir/test/python/ir/auto_location.py @@ -0,0 +1,101 @@ +# RUN: %PYTHON %s | FileCheck %s +# REQUIRES: python-ge-311 +import gc +from contextlib import contextmanager + +from mlir.ir import * +from mlir.dialects._ods_common import _cext +from mlir.dialects import arith, _arith_ops_gen + + +def run(f): + print("\nTEST:", f.__name__) + f() + gc.collect() + assert Context._get_live_count() == 0 + + +@contextmanager +def with_infer_location(): + _cext.globals.set_loc_tracebacks_enabled(True) + yield + _cext.globals.set_loc_tracebacks_enabled(False) + + +# CHECK-LABEL: TEST: testInferLocations +@run +def testInferLocations(): + with Context() as ctx, with_infer_location(): + ctx.allow_unregistered_dialects = True + + op = Operation.create("custom.op1") + one = arith.constant(IndexType.get(), 1) + _cext.globals.register_traceback_file_exclusion(arith.__file__) + two = arith.constant(IndexType.get(), 2) + + # fmt: off + # CHECK: loc(callsite("testInferLocations"("{{.*}}[[SEP:[/\\]]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":31:13 to :43) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at ""("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4)))) + # fmt: on + print(op.location) + + # fmt: off + # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]arith.py":65:12 to :76) at callsite("constant"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]arith.py":110:40 to :81) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":32:14 to :48) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at ""("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4)))))) + # fmt: on + print(one.location) + + # fmt: off + # CHECK: loc(callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":34:14 to :48) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at ""("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4)))) + # fmt: on + print(two.location) + + _cext.globals.register_traceback_file_inclusion(_arith_ops_gen.__file__) + three = arith.constant(IndexType.get(), 3) + # fmt: off + # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":405:4 to :218) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":52:16 to :50) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at ""("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4))))) + # fmt: on + print(three.location) + + def foo(): + four = arith.constant(IndexType.get(), 4) + print(four.location) + + # fmt: off + # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":405:4 to :218) at callsite("testInferLocations..foo"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":59:19 to :53) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":65:8 to :13) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at ""("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4)))))) + # fmt: on + foo() + + _cext.globals.register_traceback_file_exclusion(__file__) + + # fmt: off + # CHECK: loc("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":405:4 to :218)) + # fmt: on + foo() + + def bar1(): + def bar2(): + def bar3(): + five = arith.constant(IndexType.get(), 5) + print(five.location) + + bar3() + + bar2() + + _cext.globals.register_traceback_file_inclusion(__file__) + _cext.globals.register_traceback_file_exclusion(_arith_ops_gen.__file__) + + _cext.globals.set_loc_tracebacks_frame_limit(2) + # fmt: off + # CHECK: loc(callsite("testInferLocations..bar1..bar2..bar3"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":77:27 to :61) at "testInferLocations..bar1..bar2"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":80:16 to :22))) + # fmt: on + bar1() + + _cext.globals.set_loc_tracebacks_frame_limit(1) + # fmt: off + # CHECK: loc("testInferLocations..bar1..bar2..bar3"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":77:27 to :61)) + # fmt: on + bar1() + + _cext.globals.set_loc_tracebacks_frame_limit(0) + # CHECK: loc(unknown) + bar1() diff --git a/mlir/test/python/ir/context_managers.py b/mlir/test/python/ir/context_managers.py index 8091687f7f08..5d9f9ceee97f 100644 --- a/mlir/test/python/ir/context_managers.py +++ b/mlir/test/python/ir/context_managers.py @@ -35,25 +35,14 @@ def testLocationEnterExit(): # Asserting a different context should clear it. with Context() as ctx2: assert Context.current is ctx2 - try: - _ = Location.current - except ValueError: - pass - else: - assert False, "Expected exception" + assert Location.current is None # And should restore. assert Context.current is ctx1 assert Location.current is loc1 # All should clear. - try: - _ = Location.current - except ValueError as e: - # CHECK: No current Location - print(e) - else: - assert False, "Expected exception" + assert Location.current is None run(testLocationEnterExit) @@ -72,12 +61,7 @@ def testInsertionPointEnterExit(): assert InsertionPoint.current is ip assert Location.current is loc1 # Location should clear. - try: - _ = Location.current - except ValueError: - pass - else: - assert False, "Expected exception" + assert Location.current is None # Asserting the same Context should preserve. with ctx1: diff --git a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp index d2e38e9d2319..038f56d5a215 100644 --- a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp +++ b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp @@ -41,6 +41,7 @@ from ._ods_common import ( segmented_accessor as _ods_segmented_accessor, ) _ods_ir = _ods_cext.ir +_ods_cext.globals.register_traceback_file_exclusion(__file__) import builtins from typing import Sequence as _Sequence, Union as _Union