[mlir] Add a _get_default_loc_context utility to Python bindings

[ftynse]

Summary

This utility function is helpful for dialect-specific builders that need
to access the context through location, and the location itself may be
either provided as an argument or expected to be recovered from the
implicit location stack.

mlir/lib/Bindings/Python/mlir/dialects/__init__.py

def _equally_sized_accessor(elements, n_variadic, n_preceding_simple,
 
   total_variadic_length = len(elements) - n_variadic + 1
   # This should be enforced by the C++-side trait verifier.
   assert total_variadic_length % n_variadic == 0
 
   elements_per_group = total_variadic_length // n_variadic
   start = n_preceding_simple + n_preceding_variadic * elements_per_group
   return start, elements_per_group
+
+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.
+  """
+  if location is None:
+    # Location.current raises ValueError if there is no current location.
+    return _cext.ir.Location.current.context
+  return location.context

mlir/test/mlir-tblgen/op-python-bindings.td

def AttrSizedOperandsOp : TestOp<"attr_sized_operands",
   // CHECK:   operand_segment_sizes = array.array('L')
   // CHECK:   operands += [*variadic1]
   // CHECK:   operand_segment_sizes.append(len(variadic1))
   // CHECK:   operands.append(non_variadic)
   // CHECK:   operand_segment_sizes.append(1)
   // CHECK:   if variadic2 is not None: operands.append(variadic2)
   // CHECK:   operand_segment_sizes.append(0 if variadic2 is None else 1)
   // CHECK:   attributes["operand_segment_sizes"] = _ir.DenseElementsAttr.get(operand_segment_sizes,
-  // CHECK:       context=Location.current.context if loc is None else loc.context)
+  // CHECK:       context=_get_default_loc_context(loc))
   // CHECK:   super().__init__(_ir.Operation.create(
   // CHECK:     "test.attr_sized_operands", attributes=attributes, operands=operands, results=results,
   // CHECK:     loc=loc, ip=ip))
 
   // CHECK: @property
   // CHECK: def variadic1(self):
   // CHECK:   operand_range = _segmented_accessor(
   // CHECK:       self.operation.operands,
def AttrSizedResultsOp : TestOp<"attr_sized_results",
   // CHECK:   result_segment_sizes = array.array('L')
   // CHECK:   if variadic1 is not None: results.append(variadic1)
   // CHECK:   result_segment_sizes.append(0 if variadic1 is None else 1)
   // CHECK:   results.append(non_variadic)
   // CHECK:   result_segment_sizes.append(1) # non_variadic
   // CHECK:   if variadic2 is not None: results.append(variadic2)
   // CHECK:   result_segment_sizes.append(0 if variadic2 is None else 1)
   // CHECK:   attributes["result_segment_sizes"] = _ir.DenseElementsAttr.get(result_segment_sizes,
-  // CHECK:       context=Location.current.context if loc is None else loc.context)
+  // CHECK:       context=_get_default_loc_context(loc))
   // CHECK:   super().__init__(_ir.Operation.create(
   // CHECK:     "test.attr_sized_results", attributes=attributes, operands=operands, results=results,
   // CHECK:     loc=loc, ip=ip))
 
   // CHECK: @property
   // CHECK: def variadic1(self):
   // CHECK:   result_range = _segmented_accessor(
   // CHECK:       self.operation.results,

 def AttributedOp : TestOp<"attributed_op"> {
   // CHECK: def __init__(self, i32attr, optionalF32Attr, unitAttr, in_, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   attributes["i32attr"] = i32attr
   // CHECK:   if optionalF32Attr is not None: attributes["optionalF32Attr"] = optionalF32Attr
   // CHECK:   if bool(unitAttr): attributes["unitAttr"] = _ir.UnitAttr.get(
-  // CHECK:     _ir.Location.current.context if loc is None else loc.context)
+  // CHECK:     _get_default_loc_context(loc))
   // CHECK:   attributes["in"] = in_
   // CHECK:   super().__init__(_ir.Operation.create(
   // CHECK:     "test.attributed_op", attributes=attributes, operands=operands, results=results,
   // CHECK:     loc=loc, ip=ip))
 
   // CHECK: @property
   // CHECK: def i32attr(self):
   // CHECK:   return _ir.IntegerAttr(self.operation.attributes["i32attr"])

 def AttributedOpWithOperands : TestOp<"attributed_op_with_operands"> {
   // CHECK: def __init__(self, _gen_arg_0, in_, _gen_arg_2, is_, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   operands.append(_gen_arg_0)
   // CHECK:   operands.append(_gen_arg_2)
   // CHECK:   if bool(in_): attributes["in"] = _ir.UnitAttr.get(
-  // CHECK:     _ir.Location.current.context if loc is None else loc.context)
+  // CHECK:     _get_default_loc_context(loc))
   // CHECK:   if is_ is not None: attributes["is"] = is_
   // CHECK:   super().__init__(_ir.Operation.create(
   // CHECK:     "test.attributed_op_with_operands", attributes=attributes, operands=operands, results=results,
   // CHECK:     loc=loc, ip=ip))
 
   // CHECK: @property
   // CHECK: def in_(self):
   // CHECK:   return "in" in self.operation.attributes

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

 using namespace mlir::tblgen;
 
 /// File header and includes.
 constexpr const char *fileHeader = R"Py(
 # Autogenerated by mlir-tblgen; don't manually edit.
 
 import array
 from . import _cext
-from . import _segmented_accessor, _equally_sized_accessor
+from . import _segmented_accessor, _equally_sized_accessor, _get_default_loc_context
 _ir = _cext.ir
 )Py";
 
 /// Template for dialect class:
 ///   {0} is the dialect namespace.
 constexpr const char *dialectClassTemplate = R"Py(
 @_cext.register_dialect
 class _Dialect(_ir.Dialect):

 
 /// Template for setting up the segment sizes buffer.
 constexpr const char *segmentDeclarationTemplate =
     "{0}_segment_sizes = array.array('L')";
 
 /// Template for attaching segment sizes to the attribute list.
 constexpr const char *segmentAttributeTemplate =
     R"Py(attributes["{0}_segment_sizes"] = _ir.DenseElementsAttr.get({0}_segment_sizes,
-      context=Location.current.context if loc is None else loc.context))Py";
+      context=_get_default_loc_context(loc)))Py";
 
 /// Template for appending the unit size to the segment sizes.
 ///   {0} is either 'operand' or 'result';
 ///   {1} is the field name.
 constexpr const char *singleElementSegmentTemplate =
     "{0}_segment_sizes.append(1) # {1}";
 
 /// Template for appending 0/1 for an optional element to the segment sizes.

 /// Template for setting an optional attribute in the operation builder.
 ///   {0} is the attribute name;
 ///   {1} is the builder argument name.
 constexpr const char *initOptionalAttributeTemplate =
     R"Py(if {1} is not None: attributes["{0}"] = {1})Py";
 
 constexpr const char *initUnitAttributeTemplate =
     R"Py(if bool({1}): attributes["{0}"] = _ir.UnitAttr.get(
-      _ir.Location.current.context if loc is None else loc.context))Py";
+      _get_default_loc_context(loc)))Py";
 
 /// Populates `builderArgs` with the Python-compatible names of builder function
 /// arguments, first the results, then the intermixed attributes and operands in
 /// the same order as they appear in the `arguments` field of the op definition.
 /// Additionally, `operandNames` is populated with names of operands in their
 /// order of appearance.
 static void
 populateBuilderArgs(const Operator &op,