[mlir][py] Reuse more of CAPI build time inference.

[jpienaar]

Summary

This reduces code generated for type inference and instead reuses
facilities CAPI side that performed same role.

[mehdi_amini]

It's interesting that mlirOperationCreate in the C API is failable: it makes it a discrepancy with the C++ model right?

mlir/lib/Bindings/Python/IRCore.cpp

1483

Can we outline this?

2871

Can you update the kOperationCreateDocstring?

[jpienaar]

Outline result unpacking & update doc string.

[jpienaar]

Update one comment string that got changed.

Also the diff shown is unfortunate as the changes to buildGeneric seems large due to outlining.

[jpienaar]

In D156041#4525474, @mehdi_amini wrote:

It's interesting that mlirOperationCreate in the C API is failable: it makes it a discrepancy with the C++ model right?

I think this only happens in invalid cases and in those cases one is in UB C++ side too (well mostly it crashes). But yes pre-existing difference here.

[mehdi_amini]

You have things from the subsequent revision to squash here fist

[jpienaar]

In D156041#4526510, @mehdi_amini wrote:

You have things from the subsequent revision to squash here fist

Will fix, thanks

[jpienaar]

Woops I forgot to submit this

[jpienaar]

Seems this was submitted but review was just not linked https://github.com/llvm/llvm-project/commit/f573bc24d4ea61d68ea61e49eeefbb9a84fa4f34

mlir/lib/Bindings/Python/IRCore.cpp

Args:
   results: Sequence of Type representing op result types.
   attributes: Dict of str:Attribute.
   successors: List of Block for the operation's successors.
   regions: Number of regions to create.
   location: A Location object (defaults to resolve from context manager).
   ip: An InsertionPoint (defaults to resolve from context manager or set to
     False to disable insertion, even with an insertion point set in the
     context manager).
+  infer_type: Whether to infer result types.
 Returns:
   A new "detached" Operation object. Detached operations can be added
   to blocks, which causes them to become "attached."
 )";
 
 static const char kOperationPrintDocstring[] =
     R"(Prints the assembly form of the operation to a file like object.
 

 }
 
 py::object PyOperation::create(const std::string &name,
                                std::optional<std::vector<PyType *>> results,
                                std::optional<std::vector<PyValue *>> operands,
                                std::optional<py::dict> attributes,
                                std::optional<std::vector<PyBlock *>> successors,
                                int regions, DefaultingPyLocation location,
-                               const py::object &maybeIp) {
+                               const py::object &maybeIp, bool inferType) {
   llvm::SmallVector<MlirValue, 4> mlirOperands;
   llvm::SmallVector<MlirType, 4> mlirResults;
   llvm::SmallVector<MlirBlock, 4> mlirSuccessors;
   llvm::SmallVector<std::pair<std::string, MlirAttribute>, 4> mlirAttributes;
 
   // General parameter validation.
   if (regions < 0)
     throw py::value_error("number of regions must be >= 0");
py::object PyOperation::create(const std::string &name,
 
   // Apply unpacked/validated to the operation state. Beyond this
   // point, exceptions cannot be thrown or else the state will leak.
   MlirOperationState state =
       mlirOperationStateGet(toMlirStringRef(name), location);
   if (!mlirOperands.empty())
     mlirOperationStateAddOperands(&state, mlirOperands.size(),
                                   mlirOperands.data());
+  state.enableResultTypeInference = inferType;
   if (!mlirResults.empty())
     mlirOperationStateAddResults(&state, mlirResults.size(),
                                  mlirResults.data());
   if (!mlirAttributes.empty()) {
     // Note that the attribute names directly reference bytes in
     // mlirAttributes, so that vector must not be changed from here
     // on.
     llvm::SmallVector<MlirNamedAttribute, 4> mlirNamedAttributes;
if (regions) {
     for (int i = 0; i < regions; ++i)
       mlirRegions[i] = mlirRegionCreate();
     mlirOperationStateAddOwnedRegions(&state, mlirRegions.size(),
                                       mlirRegions.data());
   }
 
   // Construct the operation.
   MlirOperation operation = mlirOperationCreate(&state);
+  if (!operation.ptr)
+    throw py::value_error("Operation creation failed");
   PyOperationRef created =
       PyOperation::createDetached(location->getContext(), operation);
   maybeInsertOperation(created, maybeIp);
 
   return created->createOpView();
 }
 
 py::object PyOperation::clone(const py::object &maybeIp) {
void PyOperation::erase() {
   mlirOperationDestroy(operation);
   valid = false;
 }
 
 //------------------------------------------------------------------------------
 // PyOpView
 //------------------------------------------------------------------------------
 
-py::object
-PyOpView::buildGeneric(const py::object &cls, py::list resultTypeList,
-                       py::list operandList, std::optional<py::dict> attributes,
-                       std::optional<std::vector<PyBlock *>> successors,
-                       std::optional<int> regions,
-                       DefaultingPyLocation location,
-                       const py::object &maybeIp) {
-  PyMlirContextRef context = location->getContext();
-  // Class level operation construction metadata.
-  std::string name = py::cast<std::string>(cls.attr("OPERATION_NAME"));
-  // Operand and result segment specs are either none, which does no
-  // variadic unpacking, or a list of ints with segment sizes, where each
-  // element is either a positive number (typically 1 for a scalar) or -1 to
-  // indicate that it is derived from the length of the same-indexed operand
-  // or result (implying that it is a list at that position).
-  py::object operandSegmentSpecObj = cls.attr("_ODS_OPERAND_SEGMENTS");
-  py::object resultSegmentSpecObj = cls.attr("_ODS_RESULT_SEGMENTS");
-
-  std::vector<int32_t> operandSegmentLengths;
-  std::vector<int32_t> resultSegmentLengths;
-
-  // Validate/determine region count.
-  auto opRegionSpec = py::cast<std::tuple<int, bool>>(cls.attr("_ODS_REGIONS"));
-  int opMinRegionCount = std::get<0>(opRegionSpec);
-  bool opHasNoVariadicRegions = std::get<1>(opRegionSpec);
-  if (!regions) {
-    regions = opMinRegionCount;
-  }
-  if (*regions < opMinRegionCount) {
-    throw py::value_error(
-        (llvm::Twine("Operation \"") + name + "\" requires a minimum of " +
-         llvm::Twine(opMinRegionCount) +
-         " regions but was built with regions=" + llvm::Twine(*regions))
-            .str());
-  }
-  if (opHasNoVariadicRegions && *regions > opMinRegionCount) {
-    throw py::value_error(
-        (llvm::Twine("Operation \"") + name + "\" requires a maximum of " +
-         llvm::Twine(opMinRegionCount) +
-         " regions but was built with regions=" + llvm::Twine(*regions))
-            .str());
-  }
-
-  // Unpack results.
-  std::vector<PyType *> resultTypes;
+static void populateResultTypes(StringRef name, py::list resultTypeList,
+                                const py::object &resultSegmentSpecObj,
+                                std::vector<int32_t> &resultSegmentLengths,
+                                std::vector<PyType *> &resultTypes) {
   resultTypes.reserve(resultTypeList.size());
   if (resultSegmentSpecObj.is_none()) {
     // Non-variadic result unpacking.
     for (const auto &it : llvm::enumerate(resultTypeList)) {
       try {
         resultTypes.push_back(py::cast<PyType *>(it.value()));
         if (!resultTypes.back())
           throw py::cast_error();
if (resultSegmentSpec.size() != resultTypeList.size()) {
                              llvm::Twine(resultTypeList.size()))
                                 .str());
     }
     resultSegmentLengths.reserve(resultTypeList.size());
     for (const auto &it :
          llvm::enumerate(llvm::zip(resultTypeList, resultSegmentSpec))) {
       int segmentSpec = std::get<1>(it.value());
       if (segmentSpec == 1 || segmentSpec == 0) {
         // Unpack unary element.
         try {
           auto *resultType = py::cast<PyType *>(std::get<0>(it.value()));
           if (resultType) {
             resultTypes.push_back(resultType);
             resultSegmentLengths.push_back(1);
           } else if (segmentSpec == 0) {
             // Allowed to be optional.
             resultSegmentLengths.push_back(0);
for (const auto &it :
                                  err.what() + ")")
                                     .str());
         }
       } else {
         throw py::value_error("Unexpected segment spec");
       }
     }
   }
+}
+
+py::object PyOpView::buildGeneric(
+    const py::object &cls, std::optional<py::list> resultTypeList,
+    py::list operandList, std::optional<py::dict> attributes,
+    std::optional<std::vector<PyBlock *>> successors,
+    std::optional<int> regions, DefaultingPyLocation location,
+    const py::object &maybeIp) {
+  PyMlirContextRef context = location->getContext();
+  // Class level operation construction metadata.
+  std::string name = py::cast<std::string>(cls.attr("OPERATION_NAME"));
+  // Operand and result segment specs are either none, which does no
+  // variadic unpacking, or a list of ints with segment sizes, where each
+  // element is either a positive number (typically 1 for a scalar) or -1 to
+  // indicate that it is derived from the length of the same-indexed operand
+  // or result (implying that it is a list at that position).
+  py::object operandSegmentSpecObj = cls.attr("_ODS_OPERAND_SEGMENTS");
+  py::object resultSegmentSpecObj = cls.attr("_ODS_RESULT_SEGMENTS");
+
+  std::vector<int32_t> operandSegmentLengths;
+  std::vector<int32_t> resultSegmentLengths;
+
+  // Validate/determine region count.
+  auto opRegionSpec = py::cast<std::tuple<int, bool>>(cls.attr("_ODS_REGIONS"));
+  int opMinRegionCount = std::get<0>(opRegionSpec);
+  bool opHasNoVariadicRegions = std::get<1>(opRegionSpec);
+  if (!regions) {
+    regions = opMinRegionCount;
+  }
+  if (*regions < opMinRegionCount) {
+    throw py::value_error(
+        (llvm::Twine("Operation \"") + name + "\" requires a minimum of " +
+         llvm::Twine(opMinRegionCount) +
+         " regions but was built with regions=" + llvm::Twine(*regions))
+            .str());
+  }
+  if (opHasNoVariadicRegions && *regions > opMinRegionCount) {
+    throw py::value_error(
+        (llvm::Twine("Operation \"") + name + "\" requires a maximum of " +
+         llvm::Twine(opMinRegionCount) +
+         " regions but was built with regions=" + llvm::Twine(*regions))
+            .str());
+  }
+
+  // Unpack results.
+  std::vector<PyType *> resultTypes;
+  if (resultTypeList.has_value()) {
+    populateResultTypes(name, *resultTypeList, resultSegmentSpecObj,
+                        resultSegmentLengths, resultTypes);
+  }
 
   // Unpack operands.
   std::vector<PyValue *> operands;
   operands.reserve(operands.size());
   if (operandSegmentSpecObj.is_none()) {
     // Non-sized operand unpacking.
     for (const auto &it : llvm::enumerate(operandList)) {
       try {
py::object PyOpView::buildGeneric(
   }
 
   // Delegate to create.
   return PyOperation::create(name,
                              /*results=*/std::move(resultTypes),
                              /*operands=*/std::move(operands),
                              /*attributes=*/std::move(attributes),
                              /*successors=*/std::move(successors),
-                             /*regions=*/*regions, location, maybeIp);
+                             /*regions=*/*regions, location, maybeIp,
+                             !resultTypeList);
 }
 
 pybind11::object PyOpView::constructDerived(const pybind11::object &cls,
                                             const PyOperation &operation) {
   // TODO: pybind11 2.6 supports a more direct form.
   // Upgrade many years from now.
   //   auto opViewType = py::type::of<PyOpView>();
   py::handle opViewType = py::detail::get_type_handle(typeid(PyOpView), true);
void mlir::python::populateIRCore(py::module &m) {
 
   py::class_<PyOperation, PyOperationBase>(m, "Operation", py::module_local())
       .def_static("create", &PyOperation::create, py::arg("name"),
                   py::arg("results") = py::none(),
                   py::arg("operands") = py::none(),
                   py::arg("attributes") = py::none(),
                   py::arg("successors") = py::none(), py::arg("regions") = 0,
                   py::arg("loc") = py::none(), py::arg("ip") = py::none(),
-                  kOperationCreateDocstring)
+                  py::arg("infer_type") = false, kOperationCreateDocstring)
       .def_static(
           "parse",
           [](const std::string &sourceStr, const std::string &sourceName,
              DefaultingPyMlirContext context) {
             return PyOperation::parse(context->getRef(), sourceStr, sourceName)
                 ->createOpView();
           },
           py::arg("source"), py::kw_only(), py::arg("source_name") = "",

mlir/lib/Bindings/Python/IRModule.h

public:
   static pybind11::object createFromCapsule(pybind11::object capsule);
 
   /// Creates an operation. See corresponding python docstring.
   static pybind11::object
   create(const std::string &name, std::optional<std::vector<PyType *>> results,
          std::optional<std::vector<PyValue *>> operands,
          std::optional<pybind11::dict> attributes,
          std::optional<std::vector<PyBlock *>> successors, int regions,
-         DefaultingPyLocation location, const pybind11::object &ip);
+         DefaultingPyLocation location, const pybind11::object &ip,
+         bool inferType = false);
 
   /// Creates an OpView suitable for this operation.
   pybind11::object createOpView();
 
   /// Erases the underlying MlirOperation, removes its pointer from the
   /// parent context's live operations map, and sets the valid bit false.
   void erase();
 

 /// python types.
 class PyOpView : public PyOperationBase {
 public:
   PyOpView(const pybind11::object &operationObject);
   PyOperation &getOperation() override { return operation; }
 
   pybind11::object getOperationObject() { return operationObject; }
 
-  static pybind11::object
-  buildGeneric(const pybind11::object &cls, pybind11::list resultTypeList,
-               pybind11::list operandList,
-               std::optional<pybind11::dict> attributes,
-               std::optional<std::vector<PyBlock *>> successors,
-               std::optional<int> regions, DefaultingPyLocation location,
-               const pybind11::object &maybeIp);
+  static pybind11::object buildGeneric(
+      const pybind11::object &cls, std::optional<pybind11::list> resultTypeList,
+      pybind11::list operandList, std::optional<pybind11::dict> attributes,
 
   /// Construct an instance of a class deriving from OpView, bypassing its
   /// `__init__` method. The derived class will typically define a constructor
   /// that provides a convenient builder, but we need to side-step this when
   /// constructing an `OpView` for an already-built operation.
   ///
   /// The caller is responsible for verifying that `operation` is a valid
   /// operation to construct `cls` with.

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

def EmptyOp : TestOp<"empty">;
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
   // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
 
 
 // CHECK-LABEL: OPERATION_NAME = "test.infer_result_types_implied_op"
 def InferResultTypesImpliedOp : TestOp<"infer_result_types_implied_op"> {
   // CHECK:  def __init__(self, *, loc=None, ip=None):
-  // CHECK:    operands = []
-  // CHECK:    results = []
   // CHECK:    _ods_context = _ods_get_default_loc_context(loc)
-  // CHECK:    results = _ods_ir.InferTypeOpInterface(InferResultTypesImpliedOp).inferReturnTypes(
-  // CHECK:        operands=operands,
-  // CHECK:        attributes=_ods_ir.DictAttr.get(attributes, context=_ods_context),
-  // CHECK:        context=_ods_context,
-  // CHECK:        loc=loc)
+  // CHECK:   super().__init__(self.build_generic(
+  // CHECK:     attributes=attributes, operands=operands,
+  // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
   let results = (outs I32:$i32, F32:$f32);
 }
 
 // CHECK-LABEL: OPERATION_NAME = "test.infer_result_types_op"
 def InferResultTypesOp : TestOp<"infer_result_types_op", [InferTypeOpInterface]> {
   // CHECK:  def __init__(self, *, loc=None, ip=None):
   // CHECK:    operands = []
-  // CHECK:    results = []
-  // CHECK:    _ods_context = _ods_get_default_loc_context(loc)
-  // CHECK:    results = _ods_ir.InferTypeOpInterface(InferResultTypesOp).inferReturnTypes(
-  // CHECK:        operands=operands,
-  // CHECK:        attributes=_ods_ir.DictAttr.get(attributes, context=_ods_context),
-  // CHECK:        context=_ods_context,
-  // CHECK:        loc=loc)
+  // CHECK:   super().__init__(self.build_generic(
+  // CHECK:     attributes=attributes, operands=operands,
+  // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
   let results = (outs AnyType, AnyType, AnyType);
 }
 
 // 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):

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

 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,
+      attributes=attributes,{2} operands=operands,
       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})";

 _ods_derived_result_type = (
     _ods_ir.TypeAttr(_ods_result_type_source_attr).value
     if _ods_ir.TypeAttr.isinstance(_ods_result_type_source_attr) else
     _ods_result_type_source_attr.type))PY";
 
 /// Python code template appending {0} type {1} times to the results list.
 constexpr const char *appendSameResultsTemplate = "results.extend([{0}] * {1})";
 
-/// Python code template for inferring the operation results using the
-/// corresponding interface:
-///   - {0} is the name of the class for which the types are inferred.
-constexpr const char *inferTypeInterfaceTemplate =
-    R"PY(results = _ods_ir.InferTypeOpInterface({0}).inferReturnTypes(
-    operands=operands,
-    attributes=_ods_ir.DictAttr.get(attributes, context=_ods_context),
-    context=_ods_context,
-    loc=loc)
-)PY";
-
 /// Appends the given multiline string as individual strings into
 /// `builderLines`.
 static void appendLineByLine(StringRef string,
                              llvm::SmallVectorImpl<std::string> &builderLines) {
 
   std::pair<StringRef, StringRef> split = std::make_pair(string, string);
   do {
     split = split.second.split('\n');
appendLineByLine(
         llvm::formatv(deriveTypeFromAttrTemplate, firstAttr.name).str(),
         builderLines);
     builderLines.push_back(llvm::formatv(appendSameResultsTemplate,
                                          "_ods_derived_result_type",
                                          op.getNumResults()));
     return;
   }
 
-  if (hasInferTypeInterface(op)) {
-    appendLineByLine(
-        llvm::formatv(inferTypeInterfaceTemplate, op.getCppClassName()).str(),
-        builderLines);
+  if (hasInferTypeInterface(op))
     return;
-  }
 
   // For each element, find or generate a name.
   for (int i = 0, e = op.getNumResults(); i < e; ++i) {
     const NamedTypeConstraint &element = op.getResult(i);
     std::string name = names[i];
 
     // Choose the formatting string based on the element kind.
     llvm::StringRef formatString;
for (size_t i = 0, cnt = builderArgs.size(); i < cnt; ++i) {
     if (isKeywordArgFn(i)) {
       builderArgs[i].append("=None");
       functionArgs.push_back(builderArgs[i]);
     }
   }
   functionArgs.push_back("loc=None");
   functionArgs.push_back("ip=None");
   os << llvm::formatv(initTemplate, llvm::join(functionArgs, ", "),
-                      llvm::join(builderLines, "\n    "));
+                      llvm::join(builderLines, "\n    "),
+                      hasInferTypeInterface(op) ? "" : " results=results,");
 }
 
 static void emitSegmentSpec(
     const Operator &op, const char *kind,
     llvm::function_ref<int(const Operator &)> getNumElements,
     llvm::function_ref<const NamedTypeConstraint &(const Operator &, int)>
         getElement,
     raw_ostream &os) {