[mlir][python] Better support for variadic regions in Python bindings

[ftynse]

Summary

Improve support for variadic regions in ODS-generated operation view classes.
In particular, make generated constructors take an extra argument that
specifies the number of variadic regions if the operation has them. Previously,
there was no mechanism to specify a non-zero number of variadic regions. Also
generate named accessors to regions.

[gysit]

LGTM!

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

 // CHECK-LABEL: OPERATION_NAME = "test.attr_sized_operands"
 // CHECK: _ODS_OPERAND_SEGMENTS = [-1,1,-1,]
 def AttrSizedOperandsOp : TestOp<"attr_sized_operands",
                                  [AttrSizedOperandSegments]> {
   // CHECK: def __init__(self, variadic1, non_variadic, variadic2, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
+  // CHECK:   regions = None
   // CHECK:   operands.append(_get_op_results_or_values(variadic1))
   // CHECK:   operands.append(_get_op_result_or_value(non_variadic))
   // CHECK:   if variadic2 is not None: operands.append(_get_op_result_or_value(variadic2))
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def variadic1(self):
   // CHECK:   operand_range = _ods_segmented_accessor(
   // CHECK:       self.operation.operands,
   // CHECK:       self.operation.attributes["operand_segment_sizes"], 0)
   // CHECK:   return operand_range
   //

 // CHECK-LABEL: OPERATION_NAME = "test.attr_sized_results"
 // CHECK: _ODS_RESULT_SEGMENTS = [-1,1,-1,]
 def AttrSizedResultsOp : TestOp<"attr_sized_results",
                                [AttrSizedResultSegments]> {
   // CHECK: def __init__(self, variadic1, non_variadic, variadic2, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
+  // CHECK:   regions = None
   // CHECK:   if variadic1 is not None: results.append(variadic1)
   // CHECK:   results.append(non_variadic)
   // CHECK:   if variadic2 is not None: results.append(variadic2)
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def variadic1(self):
   // CHECK:   result_range = _ods_segmented_accessor(
   // CHECK:       self.operation.results,
   // CHECK:       self.operation.attributes["result_segment_sizes"], 0)
   // CHECK:   return result_range[0] if len(result_range) > 0 else None
   //

 // CHECK-LABEL: OPERATION_NAME = "test.attributed_op"
 // CHECK-NOT: _ODS_OPERAND_SEGMENTS
 // CHECK-NOT: _ODS_RESULT_SEGMENTS
 def AttributedOp : TestOp<"attributed_op"> {
   // CHECK: def __init__(self, i32attr, optionalF32Attr, unitAttr, in_, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
+  // CHECK:   regions = None
   // CHECK:   attributes["i32attr"] = i32attr
   // CHECK:   if optionalF32Attr is not None: attributes["optionalF32Attr"] = optionalF32Attr
   // CHECK:   if bool(unitAttr): attributes["unitAttr"] = _ods_ir.UnitAttr.get(
   // CHECK:     _ods_get_default_loc_context(loc))
   // CHECK:   attributes["in"] = in_
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def i32attr(self):
   // CHECK:   return _ods_ir.IntegerAttr(self.operation.attributes["i32attr"])
 
   // CHECK: @builtins.property
   // CHECK: def optionalF32Attr(self):
   // CHECK:   if "optionalF32Attr" not in self.operation.attributes:

 // CHECK-LABEL: OPERATION_NAME = "test.attributed_op_with_operands"
 // CHECK-NOT: _ODS_OPERAND_SEGMENTS
 // CHECK-NOT: _ODS_RESULT_SEGMENTS
 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:   regions = None
   // CHECK:   operands.append(_get_op_result_or_value(_gen_arg_0))
   // CHECK:   operands.append(_get_op_result_or_value(_gen_arg_2))
   // CHECK:   if bool(in_): attributes["in"] = _ods_ir.UnitAttr.get(
   // CHECK:     _ods_get_default_loc_context(loc))
   // CHECK:   if is_ is not None: attributes["is"] = is_
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def in_(self):
   // CHECK:   return "in" in self.operation.attributes
 
   // CHECK: @builtins.property
   // CHECK: def is_(self):
   // CHECK:   if "is" not in self.operation.attributes:
   // CHECK:     return None
   // CHECK:   return _ods_ir.FloatAttr(self.operation.attributes["is"])
   let arguments = (ins I32, UnitAttr:$in, F32, OptionalAttr<F32Attr>:$is);
 }
 
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class EmptyOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.empty"
 def EmptyOp : TestOp<"empty">;
   // CHECK: def __init__(self, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
+  // CHECK:   regions = None
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class MissingNamesOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.missing_names"
 def MissingNamesOp : TestOp<"missing_names"> {
   // CHECK: def __init__(self, i32, _gen_res_1, i64, _gen_arg_0, f32, _gen_arg_2, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
+  // CHECK:   regions = None
   // CHECK:   results.append(i32)
   // CHECK:   results.append(_gen_res_1)
   // CHECK:   results.append(i64)
   // CHECK:   operands.append(_get_op_result_or_value(_gen_arg_0))
   // CHECK:   operands.append(_get_op_result_or_value(f32))
   // CHECK:   operands.append(_get_op_result_or_value(_gen_arg_2))
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def f32(self):
   // CHECK:   return self.operation.operands[1]
   let arguments = (ins I32, F32:$f32, I64);
 
   // CHECK: @builtins.property
   // CHECK: def i32(self):

 // CHECK-LABEL: OPERATION_NAME = "test.one_variadic_operand"
 // CHECK-NOT: _ODS_OPERAND_SEGMENTS
 // CHECK-NOT: _ODS_RESULT_SEGMENTS
 def OneVariadicOperandOp : TestOp<"one_variadic_operand"> {
   // CHECK: def __init__(self, non_variadic, variadic, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
+  // CHECK:   regions = None
   // CHECK:   operands.append(_get_op_result_or_value(non_variadic))
   // CHECK:   operands.extend(_get_op_results_or_values(variadic))
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def non_variadic(self):
   // CHECK:   return self.operation.operands[0]
   //
   // CHECK: @builtins.property
   // CHECK: def variadic(self):
   // CHECK:   _ods_variadic_group_length = len(self.operation.operands) - 2 + 1
   // CHECK:   return self.operation.operands[1:1 + _ods_variadic_group_length]
   let arguments = (ins AnyType:$non_variadic, Variadic<AnyType>:$variadic);
 }
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class OneVariadicResultOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.one_variadic_result"
 // CHECK-NOT: _ODS_OPERAND_SEGMENTS
 // CHECK-NOT: _ODS_RESULT_SEGMENTS
 def OneVariadicResultOp : TestOp<"one_variadic_result"> {
   // CHECK: def __init__(self, variadic, non_variadic, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
+  // CHECK:   regions = None
   // CHECK:   results.extend(variadic)
   // CHECK:   results.append(non_variadic)
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def variadic(self):
   // CHECK:   _ods_variadic_group_length = len(self.operation.results) - 2 + 1
   // CHECK:   return self.operation.results[0:0 + _ods_variadic_group_length]
   //
   // CHECK: @builtins.property
   // CHECK: def non_variadic(self):
   // CHECK:   _ods_variadic_group_length = len(self.operation.results) - 2 + 1
   // CHECK:   return self.operation.results[1 + _ods_variadic_group_length - 1]
   let results = (outs Variadic<AnyType>:$variadic, AnyType:$non_variadic);
 }
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class PythonKeywordOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.python_keyword"
 def PythonKeywordOp : TestOp<"python_keyword"> {
   // CHECK: def __init__(self, in_, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
+  // CHECK:   regions = None
   // CHECK:   operands.append(_get_op_result_or_value(in_))
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def in_(self):
   // CHECK:   return self.operation.operands[0]
   let arguments = (ins AnyType:$in);
 }
 
 // CHECK: @_ods_cext.register_operation(_Dialect)

 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class SimpleOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.simple"
 def SimpleOp : TestOp<"simple"> {
   // CHECK: def __init__(self, i64, f64, i32, f32, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
+  // CHECK:   regions = None
   // CHECK:   results.append(i64)
   // CHECK:   results.append(f64)
   // CHECK:   operands.append(_get_op_result_or_value(i32))
   // CHECK:   operands.append(_get_op_result_or_value(f32))
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def i32(self):
   // CHECK:   return self.operation.operands[0]
   //
   // CHECK: @builtins.property
   // CHECK: def f32(self):
   // CHECK:   return self.operation.operands[1]
   let arguments = (ins I32:$i32, F32:$f32);
 
   // CHECK: @builtins.property
   // CHECK: def i64(self):
   // CHECK:   return self.operation.results[0]
   //
   // CHECK: @builtins.property
   // CHECK: def f64(self):
   // CHECK:   return self.operation.results[1]
   let results = (outs I64:$i64, F64:$f64);
 }
 
+// CHECK: class VariadicAndNormalRegionOp(_ods_ir.OpView):
+// CHECK-LABEL: OPERATION_NAME = "test.variadic_and_normal_region"
+def VariadicAndNormalRegionOp : TestOp<"variadic_and_normal_region"> {
+  // CHECK:  def __init__(self, num_variadic, *, loc=None, ip=None):
+  // CHECK:    operands = []
+  // CHECK:    results = []
+  // CHECK:    attributes = {}
+  // CHECK:    regions = None
+  // CHECK:    _ods_successors = None
+  // CHECK:    regions = 2 + num_variadic
+  // CHECK:    super().__init__(self.build_generic(
+  // CHECK:      attributes=attributes, results=results, operands=operands,
+  // CHECK:      successors=_ods_successors, regions=regions, loc=loc, ip=ip))
+  let regions = (region AnyRegion:$region, AnyRegion, VariadicRegion<AnyRegion>:$variadic);
+
+  // CHECK:  @builtins.property
+  // CHECK:  def region():
+  // CHECK:    return self.regions[0]
+
+  // CHECK:  @builtins.property
+  // CHECK:  def variadic():
+  // CHECK:    return self.regions[2:]
+}
+
+// CHECK: class VariadicRegionOp(_ods_ir.OpView):
+// CHECK-LABEL: OPERATION_NAME = "test.variadic_region"
+def VariadicRegionOp : TestOp<"variadic_region"> {
+  // CHECK:  def __init__(self, num_variadic, *, loc=None, ip=None):
+  // CHECK:    operands = []
+  // CHECK:    results = []
+  // CHECK:    attributes = {}
+  // CHECK:    regions = None
+  // CHECK:    _ods_successors = None
+  // CHECK:    regions = 0 + num_variadic
+  // CHECK:    super().__init__(self.build_generic(
+  // CHECK:      attributes=attributes, results=results, operands=operands,
+  // CHECK:      successors=_ods_successors, regions=regions, loc=loc, ip=ip))
+  let regions = (region VariadicRegion<AnyRegion>:$Variadic);
+
+  // CHECK:  @builtins.property
+  // CHECK:  def Variadic():
+  // CHECK:    return self.regions[0:]
+}
+
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class WithSuccessorsOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.with_successors"
 def WithSuccessorsOp : TestOp<"with_successors"> {
   // CHECK-NOT:  _ods_successors = None
   // CHECK:      _ods_successors = []
   // CHECK-NEXT: _ods_successors.append(successor)
   // CHECK-NEXT: _ods_successors.extend(successors)
   let successors = (successor AnySuccessor:$successor,
                               VariadicSuccessor<AnySuccessor>:$successors);
 }
+

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

 ///    {0} is the name of the attribute sanitized for Python;
 ///    {1} is the original name of the attribute.
 constexpr const char *attributeDeleterTemplate = R"Py(
   @{0}.deleter
   def {0}(self):
     del self.operation.attributes["{1}"]
 )Py";
 
+constexpr const char *regionAccessorTemplate = R"PY(
+  @builtins.property
+  def {0}():
+    return self.regions[{1}]
+)PY";
+
 static llvm::cl::OptionCategory
     clOpPythonBindingCat("Options for -gen-python-op-bindings");
 
 static llvm::cl::opt<std::string>
     clDialectName("bind-dialect",
                   llvm::cl::desc("The dialect to run the generator for"),
                   llvm::cl::init(""), llvm::cl::cat(clOpPythonBindingCat));
 

 ///       `loc` and `ip`;
 ///   {1} is the code populating `operands`, `results` and `attributes`,
 ///       `successors` fields.
 constexpr const char *initTemplate = R"Py(
   def __init__(self, {0}):
     operands = []
     results = []
     attributes = {{}
+    regions = None
     {1}
     super().__init__(self.build_generic(
       attributes=attributes, results=results, operands=operands,
-      successors=_ods_successors, loc=loc, ip=ip))
+      successors=_ods_successors, regions=regions, loc=loc, ip=ip))
 )Py";
 
 /// Template for appending a single element to the operand/result list.
 ///   {0} is the field name.
 constexpr const char *singleOperandAppendTemplate =
     "operands.append(_get_op_result_or_value({0}))";
 constexpr const char *singleResultAppendTemplate = "results.append({0})";
 
if (!element.isVariableLength()) {
         formatString = multiResultAppendTemplate;
       }
     }
 
     builderLines.push_back(llvm::formatv(formatString.data(), name));
   }
 }
 
+/// If the operation has variadic regions, adds a builder argument to specify
+/// the number of those regions and builder lines to forward it to the generic
+/// constructor.
+static void
+populateBuilderRegions(const Operator &op,
+                       llvm::SmallVectorImpl<std::string> &builderArgs,
+                       llvm::SmallVectorImpl<std::string> &builderLines) {
+  if (op.hasNoVariadicRegions())
+    return;
+
+  // This is currently enforced when Operator is constructed.
+  assert(op.getNumVariadicRegions() == 1 &&
+         op.getRegion(op.getNumRegions() - 1).isVariadic() &&
+         "expected the last region to be varidic");
+
+  const NamedRegion &region = op.getRegion(op.getNumRegions() - 1);
+  std::string name =
+      ("num_" + region.name.take_front().lower() + region.name.drop_front())
+          .str();
+  builderArgs.push_back(name);
+  builderLines.push_back(
+      llvm::formatv("regions = {0} + {1}", op.getNumRegions() - 1, name));
+}
+
 /// Emits a default builder constructing an operation from the list of its
 /// result types, followed by a list of its operands.
 static void emitDefaultOpBuilder(const Operator &op, raw_ostream &os) {
   // If we are asked to skip default builders, comply.
   if (op.skipDefaultBuilders())
     return;
 
   llvm::SmallVector<std::string> builderArgs;
   llvm::SmallVector<std::string> builderLines;
   llvm::SmallVector<std::string> operandArgNames;
   llvm::SmallVector<std::string> successorArgNames;
   builderArgs.reserve(op.getNumOperands() + op.getNumResults() +
                       op.getNumNativeAttributes() + op.getNumSuccessors());
   populateBuilderArgs(op, builderArgs, operandArgNames, successorArgNames);
 
   populateBuilderLinesResult(
       op, llvm::makeArrayRef(builderArgs).take_front(op.getNumResults()),
       builderLines);
   populateBuilderLinesOperand(op, operandArgNames, builderLines);
   populateBuilderLinesAttr(
       op, llvm::makeArrayRef(builderArgs).drop_front(op.getNumResults()),
       builderLines);
   populateBuilderLinesSuccessors(op, successorArgNames, builderLines);
+  populateBuilderRegions(op, builderArgs, builderLines);
 
   builderArgs.push_back("*");
   builderArgs.push_back("loc=None");
   builderArgs.push_back("ip=None");
   os << llvm::formatv(initTemplate, llvm::join(builderArgs, ", "),
                       llvm::join(builderLines, "\n    "));
 }
 

 static void emitRegionAttributes(const Operator &op, raw_ostream &os) {
   // Emit _ODS_REGIONS = (min_region_count, has_no_variadic_regions).
   // Note that the base OpView class defines this as (0, True).
   unsigned minRegionCount = op.getNumRegions() - op.getNumVariadicRegions();
   os << llvm::formatv(opClassRegionSpecTemplate, minRegionCount,
                       op.hasNoVariadicRegions() ? "True" : "False");
 }
 
+/// Emits named accessors to regions.
+static void emitRegionAccessors(const Operator &op, raw_ostream &os) {
+  for (auto en : llvm::enumerate(op.getRegions())) {
+    const NamedRegion &region = en.value();
+    if (region.name.empty())
+      continue;
+
+    assert((!region.isVariadic() || en.index() == op.getNumRegions() - 1) &&
+           "expected only the last region to be variadic");
+    os << llvm::formatv(regionAccessorTemplate, sanitizeName(region.name),
+                        std::to_string(en.index()) +
+                            (region.isVariadic() ? ":" : ""));
+  }
+}
+
 /// Emits bindings for a specific Op to the given output stream.
 static void emitOpBindings(const Operator &op,
                            const AttributeClasses &attributeClasses,
                            raw_ostream &os) {
   os << llvm::formatv(opClassTemplate, op.getCppClassName(),
                       op.getOperationName());
 
   // Sized segments.
   if (op.getTrait(attrSizedTraitForKind("operand")) != nullptr) {
     emitSegmentSpec(op, "OPERAND", getNumOperands, getOperand, os);
   }
   if (op.getTrait(attrSizedTraitForKind("result")) != nullptr) {
     emitSegmentSpec(op, "RESULT", getNumResults, getResult, os);
   }
 
   emitRegionAttributes(op, os);
   emitDefaultOpBuilder(op, os);
   emitOperandAccessors(op, os);
   emitAttributeAccessors(op, attributeClasses, os);
   emitResultAccessors(op, os);
+  emitRegionAccessors(op, os);
 }
 
 /// Emits bindings for the dialect specified in the command line, including file
 /// headers and utilities. Returns `false` on success to comply with Tablegen
 /// registration requirements.
 static bool emitAllOps(const llvm::RecordKeeper &records, raw_ostream &os) {
   if (clDialectName.empty())
     llvm::PrintFatalError("dialect name not provided");