From e5960dcfefa7cc459d195c2cb4baeee5e97b5a64 Mon Sep 17 00:00:00 2001
From: max <maksim.levental@gmail.com>
Date: Tue, 29 Jul 2025 18:21:53 -0400
Subject: [PATCH 01/10] [mlir][python] auto-locs

---
 mlir/lib/Bindings/Python/Globals.h            |  37 +++++
 mlir/lib/Bindings/Python/IRCore.cpp           | 136 ++++++++++++++++--
 mlir/lib/Bindings/Python/IRModule.cpp         |  68 ++++++++-
 mlir/lib/Bindings/Python/IRModule.h           |   5 +-
 mlir/lib/Bindings/Python/MainModule.cpp       |  23 ++-
 mlir/test/python/ir/auto_location.py          |  72 ++++++++++
 mlir/test/python/ir/lit.local.cfg             |   2 +
 mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp |   1 +
 8 files changed, 323 insertions(+), 21 deletions(-)
 create mode 100644 mlir/test/python/ir/auto_location.py
 create mode 100644 mlir/test/python/ir/lit.local.cfg

diff --git a/mlir/lib/Bindings/Python/Globals.h b/mlir/lib/Bindings/Python/Globals.h
index 826a34a535176..2071ba92f5236 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 <optional>
+#include <regex>
 #include <string>
+#include <unordered_set>
 #include <vector>
 
 #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,37 @@ class PyGlobals {
   std::optional<nanobind::object>
   lookupOperationClass(llvm::StringRef operationName);
 
+  class TracebackLoc {
+  public:
+    bool locTracebacksEnabled() const;
+
+    void setLocTracebacksEnabled(bool value);
+
+    size_t locTracebackFramesLimit() const;
+
+    void setLocTracebackFramesLimit(size_t value);
+
+    void registerTracebackFileInclusion(const std::string &file);
+
+    void registerTracebackFileExclusion(const std::string &file);
+
+    bool isUserTracebackFilename(llvm::StringRef file);
+
+  private:
+    nanobind::ft_mutex mutex;
+    bool locTracebackEnabled_ = false;
+    size_t locTracebackFramesLimit_ = 10;
+    std::unordered_set<std::string> userTracebackIncludeFiles;
+    std::unordered_set<std::string> userTracebackExcludeFiles;
+    std::regex userTracebackIncludeRegex;
+    bool rebuildUserTracebackIncludeRegex = false;
+    std::regex userTracebackExcludeRegex;
+    bool rebuildUserTracebackExcludeRegex = false;
+    llvm::StringMap<bool> isUserTracebackFilenameCache;
+  };
+
+  TracebackLoc &getTracebackLoc() { return tracebackLoc; }
+
 private:
   static PyGlobals *instance;
 
@@ -134,6 +169,8 @@ class PyGlobals {
   /// 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 5feed95f96f53..8e976c2ddad17 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 <optional>
-#include <system_error>
-#include <utility>
 
 namespace nb = nanobind;
 using namespace nb::literals;
@@ -1523,7 +1520,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;
@@ -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<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
@@ -2789,6 +2786,111 @@ class PyOpAttributeMap {
   PyOperationRef operation;
 };
 
+// copied/borrow from
+// https://github.com/python/pythoncapi-compat/blob/b541b98df1e3e5aabb5def27422a75c876f5a88a/pythoncapi_compat.h#L222
+// bpo-40421 added PyFrame_GetLasti() to Python 3.11.0b1
+#if PY_VERSION_HEX < 0x030b00b1 && !defined(PYPY_VERSION)
+int PyFrame_GetLasti(PyFrameObject *frame) {
+#if PY_VERSION_HEX >= 0x030a00a7
+  // bpo-27129: Since Python 3.10.0a7, f_lasti is an instruction offset,
+  // not a bytes offset anymore. Python uses 16-bit "wordcode" (2 bytes)
+  // instructions.
+  if (frame->f_lasti < 0) {
+    return -1;
+  }
+  return frame->f_lasti * 2;
+#else
+  return frame->f_lasti;
+#endif
+}
+#endif
+
+constexpr size_t kMaxFrames = 512;
+
+MlirLocation tracebackToLocation(MlirContext ctx) {
+  size_t framesLimit =
+      PyGlobals::get().getTracebackLoc().locTracebackFramesLimit();
+  // We use a thread_local here mostly to avoid requiring a large amount of
+  // space.
+  thread_local std::array<MlirLocation, kMaxFrames> frames;
+  size_t count = 0;
+
+  assert(PyGILState_Check());
+
+  PyThreadState *tstate = PyThreadState_GET();
+
+  PyFrameObject *next;
+  for (PyFrameObject *pyFrame = PyThreadState_GetFrame(tstate);
+       pyFrame != nullptr && count < framesLimit;
+       next = PyFrame_GetBack(pyFrame), Py_XDECREF(pyFrame), pyFrame = next) {
+    PyCodeObject *code = PyFrame_GetCode(pyFrame);
+    auto fileNameStr =
+        nb::cast<std::string>(nb::borrow<nb::str>(code->co_filename));
+    llvm::StringRef fileName(fileNameStr);
+    if (!PyGlobals::get().getTracebackLoc().isUserTracebackFilename(fileName))
+      continue;
+
+#if PY_VERSION_HEX < 0x030b00f0
+    std::string name =
+        nb::cast<std::string>(nb::borrow<nb::str>(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<std::string>(nb::borrow<nb::str>(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;
+    if (count > framesLimit)
+      break;
+  }
+
+  if (count == 0)
+    return mlirLocationUnknownGet(ctx);
+  if (count == 1)
+    return frames[0];
+
+  MlirLocation callee = frames[0];
+  MlirLocation caller = frames[count - 1];
+  for (int i = count - 2; i >= 1; i--)
+    caller = mlirLocationCallSiteGet(frames[i], caller);
+
+  return mlirLocationCallSiteGet(callee, caller);
+}
+
+PyLocation
+maybeGetTracebackLocation(const std::optional<PyLocation> &location) {
+  MlirLocation mlirLoc;
+  MlirContext mlirCtx;
+  if (!location.has_value() &&
+      PyGlobals::get().getTracebackLoc().locTracebacksEnabled()) {
+    mlirCtx = DefaultingPyMlirContext::resolve().get();
+    mlirLoc = tracebackToLocation(mlirCtx);
+  } else if (!location.has_value()) {
+    mlirLoc = DefaultingPyLocation::resolve();
+    mlirCtx = mlirLocationGetContext(mlirLoc);
+  } else {
+    mlirLoc = *location;
+    mlirCtx = mlirLocationGetContext(mlirLoc);
+  }
+  assert(!mlirLocationIsNull(mlirLoc) && "expected non-null mlirLoc");
+  PyMlirContextRef ctx = PyMlirContext::forContext(mlirCtx);
+  return {ctx, mlirLoc};
+}
+
 } // namespace
 
 //------------------------------------------------------------------------------
@@ -3240,8 +3342,9 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           kModuleParseDocstring)
       .def_static(
           "create",
-          [](DefaultingPyLocation loc) {
-            MlirModule module = mlirModuleCreateEmpty(loc);
+          [](std::optional<PyLocation> 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")
@@ -3454,7 +3557,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
              std::optional<std::vector<PyValue *>> operands,
              std::optional<nb::dict> attributes,
              std::optional<std::vector<PyBlock *>> successors, int regions,
-             DefaultingPyLocation location, const nb::object &maybeIp,
+             std::optional<PyLocation> location, const nb::object &maybeIp,
              bool inferType) {
             // Unpack/validate operands.
             llvm::SmallVector<MlirValue, 4> mlirOperands;
@@ -3467,8 +3570,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(),
@@ -3512,12 +3616,13 @@ void mlir::python::populateIRCore(nb::module_ &m) {
                  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, std::optional<PyLocation> 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(),
@@ -3551,17 +3656,18 @@ void mlir::python::populateIRCore(nb::module_ &m) {
       [](nb::handle cls, 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, std::optional<PyLocation> location,
          const nb::object &maybeIp) {
         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"));
         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 e600f1bbd4493..09d51ee6904e7 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" // This is expected after nanobind.
 #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,69 @@ PyGlobals::lookupOperationClass(llvm::StringRef operationName) {
   // Not found and loading did not yield a registration.
   return std::nullopt;
 }
+
+bool PyGlobals::TracebackLoc::locTracebacksEnabled() const {
+  return locTracebackEnabled_;
+}
+
+void PyGlobals::TracebackLoc::setLocTracebacksEnabled(bool value) {
+  nanobind::ft_lock_guard lock(mutex);
+  locTracebackEnabled_ = value;
+}
+
+size_t PyGlobals::TracebackLoc::locTracebackFramesLimit() const {
+  return locTracebackFramesLimit_;
+}
+
+void PyGlobals::TracebackLoc::setLocTracebackFramesLimit(size_t value) {
+  nanobind::ft_lock_guard lock(mutex);
+  locTracebackFramesLimit_ = value;
+}
+
+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 9c22dea157c06..87e1a0b12da00 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/mlir/lib/Bindings/Python/IRModule.h
@@ -722,8 +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,
-         bool inferType);
+         PyLocation location, const nanobind::object &ip, bool inferType);
 
   /// Creates an OpView suitable for this operation.
   nanobind::object createOpView();
@@ -781,7 +780,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
diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp
index 6f49431006605..278847e7ac7f5 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/test/python/ir/auto_location.py b/mlir/test/python/ir/auto_location.py
new file mode 100644
index 0000000000000..1cb838e834c82
--- /dev/null
+++ b/mlir/test/python/ir/auto_location.py
@@ -0,0 +1,72 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+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, Location.unknown(), 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"("{{.*}}/test/python/ir/auto_location.py":31:13 to :43) at callsite("run"("{{.*}}/test/python/ir/auto_location.py":13:4 to :7) at "<module>"("{{.*}}/test/python/ir/auto_location.py":26:1 to :4))))
+        # fmt: on
+        print(op.location)
+
+        # fmt: off
+        # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}/mlir/dialects/arith.py":65:12 to :76) at callsite("constant"("{{.*}}/mlir/dialects/arith.py":110:40 to :81) at callsite("testInferLocations"("{{.*}}/test/python/ir/auto_location.py":32:14 to :48) at callsite("run"("{{.*}}/test/python/ir/auto_location.py":13:4 to :7) at "<module>"("{{.*}}/test/python/ir/auto_location.py":26:1 to :4))))))
+        # fmt: on
+        print(one.location)
+
+        # fmt: off
+        # CHECK: loc(callsite("testInferLocations"("{{.*}}/test/python/ir/auto_location.py":34:14 to :48) at callsite("run"("{{.*}}/test/python/ir/auto_location.py":13:4 to :7) at "<module>"("{{.*}}/test/python/ir/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__"("{{.*}}/mlir/dialects/_arith_ops_gen.py":405:4 to :218) at callsite("testInferLocations"("{{.*}}/test/python/ir/auto_location.py":52:16 to :50) at callsite("run"("{{.*}}/test/python/ir/auto_location.py":13:4 to :7) at "<module>"("{{.*}}/test/python/ir/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__"("{{.*}}/mlir/dialects/_arith_ops_gen.py":405:4 to :218) at callsite("testInferLocations.<locals>.foo"("{{.*}}/test/python/ir/auto_location.py":59:19 to :53) at callsite("testInferLocations"("{{.*}}/test/python/ir/auto_location.py":65:8 to :13) at callsite("run"("{{.*}}/test/python/ir/auto_location.py":13:4 to :7) at "<module>"("{{.*}}/test/python/ir/auto_location.py":26:1 to :4))))))
+        # fmt: on
+        foo()
+
+        _cext.globals.register_traceback_file_exclusion(__file__)
+
+        # fmt: off
+        # CHECK: loc("ConstantOp.__init__"("{{.*}}/mlir/dialects/_arith_ops_gen.py":405:4 to :218))
+        # fmt: on
+        foo()
diff --git a/mlir/test/python/ir/lit.local.cfg b/mlir/test/python/ir/lit.local.cfg
new file mode 100644
index 0000000000000..f59e80bc93ab8
--- /dev/null
+++ b/mlir/test/python/ir/lit.local.cfg
@@ -0,0 +1,2 @@
+if "Windows" in config.host_os:
+    config.excludes.add("auto_location.py")
diff --git a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
index d2e38e9d23198..038f56d5a2150 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

From 105e6ddff55dbcb2e85e3ebb5d401be08efbfdc8 Mon Sep 17 00:00:00 2001
From: max <maksim.levental@gmail.com>
Date: Thu, 31 Jul 2025 11:23:53 -0400
Subject: [PATCH 02/10] guard reading as well

---
 mlir/lib/Bindings/Python/Globals.h    |  4 ++--
 mlir/lib/Bindings/Python/IRModule.cpp |  6 ++++--
 mlir/test/python/ir/auto_location.py  | 12 ++++++------
 mlir/test/python/ir/lit.local.cfg     |  2 --
 mlir/test/python/lit.local.cfg        |  4 ++++
 5 files changed, 16 insertions(+), 12 deletions(-)
 delete mode 100644 mlir/test/python/ir/lit.local.cfg

diff --git a/mlir/lib/Bindings/Python/Globals.h b/mlir/lib/Bindings/Python/Globals.h
index 2071ba92f5236..c28b771fabc21 100644
--- a/mlir/lib/Bindings/Python/Globals.h
+++ b/mlir/lib/Bindings/Python/Globals.h
@@ -120,11 +120,11 @@ class PyGlobals {
 
   class TracebackLoc {
   public:
-    bool locTracebacksEnabled() const;
+    bool locTracebacksEnabled();
 
     void setLocTracebacksEnabled(bool value);
 
-    size_t locTracebackFramesLimit() const;
+    size_t locTracebackFramesLimit();
 
     void setLocTracebackFramesLimit(size_t value);
 
diff --git a/mlir/lib/Bindings/Python/IRModule.cpp b/mlir/lib/Bindings/Python/IRModule.cpp
index 09d51ee6904e7..b6d7b862b6b2f 100644
--- a/mlir/lib/Bindings/Python/IRModule.cpp
+++ b/mlir/lib/Bindings/Python/IRModule.cpp
@@ -198,7 +198,8 @@ PyGlobals::lookupOperationClass(llvm::StringRef operationName) {
   return std::nullopt;
 }
 
-bool PyGlobals::TracebackLoc::locTracebacksEnabled() const {
+bool PyGlobals::TracebackLoc::locTracebacksEnabled() {
+  nanobind::ft_lock_guard lock(mutex);
   return locTracebackEnabled_;
 }
 
@@ -207,7 +208,8 @@ void PyGlobals::TracebackLoc::setLocTracebacksEnabled(bool value) {
   locTracebackEnabled_ = value;
 }
 
-size_t PyGlobals::TracebackLoc::locTracebackFramesLimit() const {
+size_t PyGlobals::TracebackLoc::locTracebackFramesLimit() {
+  nanobind::ft_lock_guard lock(mutex);
   return locTracebackFramesLimit_;
 }
 
diff --git a/mlir/test/python/ir/auto_location.py b/mlir/test/python/ir/auto_location.py
index 1cb838e834c82..fc8f3a5d11cd3 100644
--- a/mlir/test/python/ir/auto_location.py
+++ b/mlir/test/python/ir/auto_location.py
@@ -34,24 +34,24 @@ def testInferLocations():
         two = arith.constant(IndexType.get(), 2)
 
         # fmt: off
-        # CHECK: loc(callsite("testInferLocations"("{{.*}}/test/python/ir/auto_location.py":31:13 to :43) at callsite("run"("{{.*}}/test/python/ir/auto_location.py":13:4 to :7) at "<module>"("{{.*}}/test/python/ir/auto_location.py":26:1 to :4))))
+        # 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 "<module>"("{{.*}}[[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__"("{{.*}}/mlir/dialects/arith.py":65:12 to :76) at callsite("constant"("{{.*}}/mlir/dialects/arith.py":110:40 to :81) at callsite("testInferLocations"("{{.*}}/test/python/ir/auto_location.py":32:14 to :48) at callsite("run"("{{.*}}/test/python/ir/auto_location.py":13:4 to :7) at "<module>"("{{.*}}/test/python/ir/auto_location.py":26:1 to :4))))))
+        # 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 "<module>"("{{.*}}[[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"("{{.*}}/test/python/ir/auto_location.py":34:14 to :48) at callsite("run"("{{.*}}/test/python/ir/auto_location.py":13:4 to :7) at "<module>"("{{.*}}/test/python/ir/auto_location.py":26:1 to :4))))
+        # 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 "<module>"("{{.*}}[[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__"("{{.*}}/mlir/dialects/_arith_ops_gen.py":405:4 to :218) at callsite("testInferLocations"("{{.*}}/test/python/ir/auto_location.py":52:16 to :50) at callsite("run"("{{.*}}/test/python/ir/auto_location.py":13:4 to :7) at "<module>"("{{.*}}/test/python/ir/auto_location.py":26:1 to :4)))))
+        # 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 "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4)))))
         # fmt: on
         print(three.location)
 
@@ -60,13 +60,13 @@ def foo():
             print(four.location)
 
         # fmt: off
-        # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}/mlir/dialects/_arith_ops_gen.py":405:4 to :218) at callsite("testInferLocations.<locals>.foo"("{{.*}}/test/python/ir/auto_location.py":59:19 to :53) at callsite("testInferLocations"("{{.*}}/test/python/ir/auto_location.py":65:8 to :13) at callsite("run"("{{.*}}/test/python/ir/auto_location.py":13:4 to :7) at "<module>"("{{.*}}/test/python/ir/auto_location.py":26:1 to :4))))))
+        # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":405:4 to :218) at callsite("testInferLocations.<locals>.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 "<module>"("{{.*}}[[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__"("{{.*}}/mlir/dialects/_arith_ops_gen.py":405:4 to :218))
+        # CHECK: loc("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":405:4 to :218))
         # fmt: on
         foo()
diff --git a/mlir/test/python/ir/lit.local.cfg b/mlir/test/python/ir/lit.local.cfg
deleted file mode 100644
index f59e80bc93ab8..0000000000000
--- a/mlir/test/python/ir/lit.local.cfg
+++ /dev/null
@@ -1,2 +0,0 @@
-if "Windows" in config.host_os:
-    config.excludes.add("auto_location.py")
diff --git a/mlir/test/python/lit.local.cfg b/mlir/test/python/lit.local.cfg
index 12d6e1f22744a..a4d0fb6c10968 100644
--- a/mlir/test/python/lit.local.cfg
+++ b/mlir/test/python/lit.local.cfg
@@ -2,3 +2,7 @@ config.environment["ASAN_OPTIONS"] = "detect_leaks=0"
 if not config.enable_bindings_python:
     config.unsupported = True
 config.excludes.add("python_test_ops.td")
+
+import sys
+if sys.version_info >= (3, 11):
+    config.available_features.add("python-ge-311")

From eb48702dfe497c0e16ee862cb53699ebd5f29770 Mon Sep 17 00:00:00 2001
From: max <maksim.levental@gmail.com>
Date: Thu, 31 Jul 2025 15:02:42 -0400
Subject: [PATCH 03/10] fix comments

---
 mlir/lib/Bindings/Python/IRCore.cpp   | 22 +---------------------
 mlir/lib/Bindings/Python/IRModule.cpp |  2 +-
 mlir/test/lit.cfg.py                  |  3 +++
 mlir/test/python/ir/auto_location.py  |  2 +-
 mlir/test/python/lit.local.cfg        |  4 ----
 5 files changed, 6 insertions(+), 27 deletions(-)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 8e976c2ddad17..43096613b0daa 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2786,32 +2786,12 @@ class PyOpAttributeMap {
   PyOperationRef operation;
 };
 
-// copied/borrow from
-// https://github.com/python/pythoncapi-compat/blob/b541b98df1e3e5aabb5def27422a75c876f5a88a/pythoncapi_compat.h#L222
-// bpo-40421 added PyFrame_GetLasti() to Python 3.11.0b1
-#if PY_VERSION_HEX < 0x030b00b1 && !defined(PYPY_VERSION)
-int PyFrame_GetLasti(PyFrameObject *frame) {
-#if PY_VERSION_HEX >= 0x030a00a7
-  // bpo-27129: Since Python 3.10.0a7, f_lasti is an instruction offset,
-  // not a bytes offset anymore. Python uses 16-bit "wordcode" (2 bytes)
-  // instructions.
-  if (frame->f_lasti < 0) {
-    return -1;
-  }
-  return frame->f_lasti * 2;
-#else
-  return frame->f_lasti;
-#endif
-}
-#endif
-
 constexpr size_t kMaxFrames = 512;
 
 MlirLocation tracebackToLocation(MlirContext ctx) {
   size_t framesLimit =
       PyGlobals::get().getTracebackLoc().locTracebackFramesLimit();
-  // We use a thread_local here mostly to avoid requiring a large amount of
-  // space.
+  // Use a thread_local here to avoid requiring a large amount of space.
   thread_local std::array<MlirLocation, kMaxFrames> frames;
   size_t count = 0;
 
diff --git a/mlir/lib/Bindings/Python/IRModule.cpp b/mlir/lib/Bindings/Python/IRModule.cpp
index b6d7b862b6b2f..21c314ee81419 100644
--- a/mlir/lib/Bindings/Python/IRModule.cpp
+++ b/mlir/lib/Bindings/Python/IRModule.cpp
@@ -13,7 +13,7 @@
 
 #include "Globals.h"
 #include "NanobindUtils.h"
-#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
+#include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/Support.h"
 #include "mlir/Bindings/Python/Nanobind.h"
 
diff --git a/mlir/test/lit.cfg.py b/mlir/test/lit.cfg.py
index feaf5fb852a1d..f4df21d235bfe 100644
--- a/mlir/test/lit.cfg.py
+++ b/mlir/test/lit.cfg.py
@@ -375,3 +375,6 @@ def have_host_jit_feature_support(feature_name):
 
 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
index fc8f3a5d11cd3..95841e3a9ca55 100644
--- a/mlir/test/python/ir/auto_location.py
+++ b/mlir/test/python/ir/auto_location.py
@@ -1,5 +1,5 @@
 # RUN: %PYTHON %s | FileCheck %s
-
+# REQUIRES: python-ge-311
 import gc
 from contextlib import contextmanager
 
diff --git a/mlir/test/python/lit.local.cfg b/mlir/test/python/lit.local.cfg
index a4d0fb6c10968..12d6e1f22744a 100644
--- a/mlir/test/python/lit.local.cfg
+++ b/mlir/test/python/lit.local.cfg
@@ -2,7 +2,3 @@ config.environment["ASAN_OPTIONS"] = "detect_leaks=0"
 if not config.enable_bindings_python:
     config.unsupported = True
 config.excludes.add("python_test_ops.td")
-
-import sys
-if sys.version_info >= (3, 11):
-    config.available_features.add("python-ge-311")

From 2942fc2ec0d7c17bd1a921ea179630f3b69ffba6 Mon Sep 17 00:00:00 2001
From: max <maksim.levental@gmail.com>
Date: Thu, 31 Jul 2025 15:20:46 -0400
Subject: [PATCH 04/10] add frame limit test

---
 mlir/test/python/ir/auto_location.py | 31 +++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/mlir/test/python/ir/auto_location.py b/mlir/test/python/ir/auto_location.py
index 95841e3a9ca55..6a31d84a7c227 100644
--- a/mlir/test/python/ir/auto_location.py
+++ b/mlir/test/python/ir/auto_location.py
@@ -12,7 +12,7 @@ def run(f):
     print("\nTEST:", f.__name__)
     f()
     gc.collect()
-    assert Context._get_live_count() == 0
+    # assert Context._get_live_count() == 0
 
 
 @contextmanager
@@ -70,3 +70,32 @@ def foo():
         # 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.<locals>.bar1.<locals>.bar2.<locals>.bar3"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":77:27 to :61) at "testInferLocations.<locals>.bar1.<locals>.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.<locals>.bar1.<locals>.bar2.<locals>.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()

From 83091aa240226429299463bc5f8cf3912ee1dbec Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental@gmail.com>
Date: Mon, 11 Aug 2025 16:55:09 -0400
Subject: [PATCH 05/10] Update auto_location.py

restore live_count == 0
---
 mlir/test/python/ir/auto_location.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/python/ir/auto_location.py b/mlir/test/python/ir/auto_location.py
index 6a31d84a7c227..538ac2476c38e 100644
--- a/mlir/test/python/ir/auto_location.py
+++ b/mlir/test/python/ir/auto_location.py
@@ -12,7 +12,7 @@ def run(f):
     print("\nTEST:", f.__name__)
     f()
     gc.collect()
-    # assert Context._get_live_count() == 0
+    assert Context._get_live_count() == 0
 
 
 @contextmanager

From 6e5cd758c82ebbb6231869222f8434e362c4468e Mon Sep 17 00:00:00 2001
From: max <maksim.levental@gmail.com>
Date: Mon, 11 Aug 2025 21:52:45 -0400
Subject: [PATCH 06/10] make get_default_loc_context return None

---
 mlir/lib/Bindings/Python/IRCore.cpp      | 11 +++++++----
 mlir/python/mlir/dialects/_ods_common.py |  8 ++++----
 mlir/test/python/ir/auto_location.py     |  2 +-
 3 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 43096613b0daa..0945ce95d6976 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2840,11 +2840,14 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
 
   if (count == 0)
     return mlirLocationUnknownGet(ctx);
-  if (count == 1)
-    return frames[0];
 
   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);
 
@@ -3134,10 +3137,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<PyLocation *> {
             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")
diff --git a/mlir/python/mlir/dialects/_ods_common.py b/mlir/python/mlir/dialects/_ods_common.py
index a5efa057cfb64..10abd06ff266e 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/python/ir/auto_location.py b/mlir/test/python/ir/auto_location.py
index 538ac2476c38e..c2d51083c1379 100644
--- a/mlir/test/python/ir/auto_location.py
+++ b/mlir/test/python/ir/auto_location.py
@@ -25,7 +25,7 @@ def with_infer_location():
 # CHECK-LABEL: TEST: testInferLocations
 @run
 def testInferLocations():
-    with Context() as ctx, Location.unknown(), with_infer_location():
+    with Context() as ctx, with_infer_location():
         ctx.allow_unregistered_dialects = True
 
         op = Operation.create("custom.op1")

From a4bc05ad3a7325b72f51b75cb78408e7046ac18e Mon Sep 17 00:00:00 2001
From: max <maksim.levental@gmail.com>
Date: Mon, 11 Aug 2025 22:01:24 -0400
Subject: [PATCH 07/10] fix live contexts

---
 mlir/lib/Bindings/Python/IRCore.cpp | 48 ++++++++++++-----------------
 mlir/lib/Bindings/Python/IRModule.h | 14 ++-------
 2 files changed, 22 insertions(+), 40 deletions(-)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 0945ce95d6976..a7665a445ef44 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -1520,7 +1520,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, PyLocation location,
+                               int regions, PyLocation &location,
                                const nb::object &maybeIp, bool inferType) {
   llvm::SmallVector<MlirType, 4> mlirResults;
   llvm::SmallVector<MlirBlock, 4> mlirSuccessors;
@@ -1934,7 +1934,7 @@ 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, PyLocation location,
+    std::optional<int> regions, PyLocation &location,
     const nb::object &maybeIp) {
   PyMlirContextRef context = location.getContext();
 
@@ -2795,13 +2795,12 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
   thread_local std::array<MlirLocation, kMaxFrames> frames;
   size_t count = 0;
 
-  assert(PyGILState_Check());
-
+  nb::gil_scoped_acquire acquire;
   PyThreadState *tstate = PyThreadState_GET();
 
   PyFrameObject *next;
-  for (PyFrameObject *pyFrame = PyThreadState_GetFrame(tstate);
-       pyFrame != nullptr && count < framesLimit;
+  PyFrameObject *pyFrame = PyThreadState_GetFrame(tstate);
+  for (; pyFrame != nullptr && count < framesLimit;
        next = PyFrame_GetBack(pyFrame), Py_XDECREF(pyFrame), pyFrame = next) {
     PyCodeObject *code = PyFrame_GetCode(pyFrame);
     auto fileNameStr =
@@ -2834,9 +2833,8 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
 
     frames[count] = mlirLocationNameGet(ctx, wrap(funcName), loc);
     ++count;
-    if (count > framesLimit)
-      break;
   }
+  Py_XDECREF(pyFrame);
 
   if (count == 0)
     return mlirLocationUnknownGet(ctx);
@@ -2856,22 +2854,15 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
 
 PyLocation
 maybeGetTracebackLocation(const std::optional<PyLocation> &location) {
-  MlirLocation mlirLoc;
-  MlirContext mlirCtx;
-  if (!location.has_value() &&
-      PyGlobals::get().getTracebackLoc().locTracebacksEnabled()) {
-    mlirCtx = DefaultingPyMlirContext::resolve().get();
-    mlirLoc = tracebackToLocation(mlirCtx);
-  } else if (!location.has_value()) {
-    mlirLoc = DefaultingPyLocation::resolve();
-    mlirCtx = mlirLocationGetContext(mlirLoc);
-  } else {
-    mlirLoc = *location;
-    mlirCtx = mlirLocationGetContext(mlirLoc);
-  }
-  assert(!mlirLocationIsNull(mlirLoc) && "expected non-null mlirLoc");
-  PyMlirContextRef ctx = PyMlirContext::forContext(mlirCtx);
-  return {ctx, mlirLoc};
+  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
@@ -3325,7 +3316,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           kModuleParseDocstring)
       .def_static(
           "create",
-          [](std::optional<PyLocation> loc) {
+          [](const std::optional<PyLocation> &loc) {
             PyLocation pyLoc = maybeGetTracebackLocation(loc);
             MlirModule module = mlirModuleCreateEmpty(pyLoc.get());
             return PyModule::forModule(module).releaseObject();
@@ -3540,8 +3531,8 @@ void mlir::python::populateIRCore(nb::module_ &m) {
              std::optional<std::vector<PyValue *>> operands,
              std::optional<nb::dict> attributes,
              std::optional<std::vector<PyBlock *>> successors, int regions,
-             std::optional<PyLocation> location, const nb::object &maybeIp,
-             bool inferType) {
+             const std::optional<PyLocation> &location,
+             const nb::object &maybeIp, bool inferType) {
             // Unpack/validate operands.
             llvm::SmallVector<MlirValue, 4> mlirOperands;
             if (operands) {
@@ -3599,7 +3590,8 @@ void mlir::python::populateIRCore(nb::module_ &m) {
                  std::optional<nb::list> resultTypeList, nb::list operandList,
                  std::optional<nb::dict> attributes,
                  std::optional<std::vector<PyBlock *>> successors,
-                 std::optional<int> regions, std::optional<PyLocation> location,
+                 std::optional<int> regions,
+                 const std::optional<PyLocation> &location,
                  const nb::object &maybeIp) {
                 PyLocation pyLoc = maybeGetTracebackLocation(location);
                 new (self) PyOpView(PyOpView::buildGeneric(
diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h
index 87e1a0b12da00..fa16ae3ce3294 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/mlir/lib/Bindings/Python/IRModule.h
@@ -192,16 +192,6 @@ class PyMlirContext {
   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,7 +712,7 @@ class PyOperation : public PyOperationBase, public BaseContextObject {
          llvm::ArrayRef<MlirValue> operands,
          std::optional<nanobind::dict> attributes,
          std::optional<std::vector<PyBlock *>> successors, int regions,
-         PyLocation 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();
@@ -780,7 +770,7 @@ class PyOpView : public PyOperationBase {
                nanobind::list operandList,
                std::optional<nanobind::dict> attributes,
                std::optional<std::vector<PyBlock *>> successors,
-               std::optional<int> regions, PyLocation location,
+               std::optional<int> regions, PyLocation &location,
                const nanobind::object &maybeIp);
 
   /// Construct an instance of a class deriving from OpView, bypassing its

From 21d68e6827bbb3bb02f65c74d8630478d58ffb19 Mon Sep 17 00:00:00 2001
From: max <maksim.levental@gmail.com>
Date: Tue, 12 Aug 2025 01:13:20 -0400
Subject: [PATCH 08/10] fix location is None test

---
 mlir/test/python/ir/context_managers.py | 22 +++-------------------
 1 file changed, 3 insertions(+), 19 deletions(-)

diff --git a/mlir/test/python/ir/context_managers.py b/mlir/test/python/ir/context_managers.py
index 8091687f7f082..5d9f9ceee97f3 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:

From 70244f353a4c98d38157fa33431768a2b223f494 Mon Sep 17 00:00:00 2001
From: max <maksim.levental@gmail.com>
Date: Tue, 12 Aug 2025 01:23:25 -0400
Subject: [PATCH 09/10] upper bound kMaxFrames

---
 mlir/lib/Bindings/Python/Globals.h    | 2 ++
 mlir/lib/Bindings/Python/IRCore.cpp   | 6 ++----
 mlir/lib/Bindings/Python/IRModule.cpp | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Bindings/Python/Globals.h b/mlir/lib/Bindings/Python/Globals.h
index c28b771fabc21..71a051cb3d9f5 100644
--- a/mlir/lib/Bindings/Python/Globals.h
+++ b/mlir/lib/Bindings/Python/Globals.h
@@ -134,6 +134,8 @@ class PyGlobals {
 
     bool isUserTracebackFilename(llvm::StringRef file);
 
+    static constexpr size_t kMaxFrames = 512;
+
   private:
     nanobind::ft_mutex mutex;
     bool locTracebackEnabled_ = false;
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index a7665a445ef44..ee18e7e104115 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2786,18 +2786,16 @@ class PyOpAttributeMap {
   PyOperationRef operation;
 };
 
-constexpr size_t kMaxFrames = 512;
-
 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<MlirLocation, kMaxFrames> frames;
+  thread_local std::array<MlirLocation, PyGlobals::TracebackLoc::kMaxFrames>
+      frames;
   size_t count = 0;
 
   nb::gil_scoped_acquire acquire;
   PyThreadState *tstate = PyThreadState_GET();
-
   PyFrameObject *next;
   PyFrameObject *pyFrame = PyThreadState_GetFrame(tstate);
   for (; pyFrame != nullptr && count < framesLimit;
diff --git a/mlir/lib/Bindings/Python/IRModule.cpp b/mlir/lib/Bindings/Python/IRModule.cpp
index 21c314ee81419..0de2f1711829b 100644
--- a/mlir/lib/Bindings/Python/IRModule.cpp
+++ b/mlir/lib/Bindings/Python/IRModule.cpp
@@ -215,7 +215,7 @@ size_t PyGlobals::TracebackLoc::locTracebackFramesLimit() {
 
 void PyGlobals::TracebackLoc::setLocTracebackFramesLimit(size_t value) {
   nanobind::ft_lock_guard lock(mutex);
-  locTracebackFramesLimit_ = value;
+  locTracebackFramesLimit_ = std::min(value, kMaxFrames);
 }
 
 void PyGlobals::TracebackLoc::registerTracebackFileInclusion(

From 2e63eeee62483888cd963b6f91fe5aaccedcd957 Mon Sep 17 00:00:00 2001
From: max <maksim.levental@gmail.com>
Date: Tue, 12 Aug 2025 14:58:44 -0400
Subject: [PATCH 10/10] add comment

---
 mlir/lib/Bindings/Python/IRCore.cpp | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index ee18e7e104115..6946303012626 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2798,6 +2798,11 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
   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);
@@ -2832,6 +2837,8 @@ MlirLocation tracebackToLocation(MlirContext ctx) {
     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)