make get_default_loc_context return None

Author: makslevental

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")

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
 
 

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")