[MLIR] Add InferShapedTypeOpInterface bindings

[ataheridezfouli-groq]

Summary

Add C and python bindings for InferShapedTypeOpInterface
and ShapedTypeComponents. This allows users to invoke
InferShapedTypeOpInterface for ops that implement it.

[ftynse]

I have a design suggestion, please see inline comments.

mlir/include/mlir-c/Interfaces.h

44

This type isn't used in C++ in the same way as the other C_API_STRUCT classes this approach was intended for, that is, it is not a wrapper around a pointer (unlike Attribute/Type/*Op). Using a pointer for it would require C++ allocation and deallocation just for API purposes.

I'd consider a different approach here: "unpack" the structure into components (dimensions given as rank + pointer to the contiguous list of the first, MlirType elementType, MlirAttr attribute, int ranked), and pass them as separate attributes to the callback below. The user will then copy them over to any storage they have without having to deal with C++ memory.

96

Nit.

112

See above, this callback can take more than one argument.

115

Nit: canonical what?

mlir/lib/Bindings/Python/IRInterfaces.cpp

304–470

This class can still exist, but contain fields for scalars (type, attribute, boolean) and store contents of the shape array in a py::list so it is managed by Python.

340

Nit: please be systematic about ending the documentation with a period.

426–478

Please refactor the existing code to reuse it instead of copying.

mlir/lib/CAPI/Interfaces/Interfaces.cpp

149

Nit: unsigned->size_t.

171–198

Please refactor existing code into a function and call it instead of copying.

mlir/test/python/python_test_ops.td

112

Ultra-nit. (That's why it is better to have code in .cpp fliles).

[ataheridezfouli-groq]

mlir/include/mlir-c/Interfaces.h

44

Hey Alex!
Thanks for taking the time to review this code! I just wanted to clarify something before I go ahead and make the changes. Are you saying that we should still have the MlirShapedTypeComponents structure, but just define it as something like this:

struct MlirShapedTypeComponents {
    intptr_t rank;
    int64_t *shape;
    MlirType elementType;
    MlirAttr attribute;
    int ranked;
};

Or are you saying to get rid of MlirShapedTypeComponents altogether, as well as any of the C-APIs, as just pass around the components?

Thanks again for all your comments!

[ataheridezfouli-groq]

Addressed review comments

[ataheridezfouli-groq]

@ftynse

Thanks again for the review. I addressed your comments. My concern now is that now with the changes, PyShapedTypeComponents is completely separated from mlir::ShapedTypeComponents. Let me know what you think! Thanks!

[ftynse]

LGTM with a couple of nits.

mlir/include/mlir-c/Interfaces.h

44

Either way is fine with me as long as there is no allocation for individual structs (it's okay to allocate an array thereof).

mlir/lib/Bindings/Python/IRInterfaces.cpp

44

Nit: please document top-level entities.

47

Nit: now that this is a function, we can invert the condition here and do exit early (https://llvm.org/docs/CodingStandards.html#use-early-exits-and-continue-to-simplify-code).

50

Nit: use auto && with enumerate.

304–470

Nit: ShapedTypeComponents reads like a class, but we don't actually have one, spell it as shaped type components. Also, the ownership part of the comment doesn't seem relevant anymore.

327–329

I think the casters defined in PybindAdaptors.h should allow us to just return self.elementType here.

340–353

Optional: we may consider passing the shape and attribute as kwargs.

354–361

Nit: a more Pythonic way could be to have rank property return None when there is no rank instead of having a separate has_rank.

[ataheridezfouli-groq]

mlir/lib/Bindings/Python/IRInterfaces.cpp

327–329

Seems to compile but during runtime throws this type error:

TypeError: Unable to convert function return value to a Python type!

354–361

I added the changes for rank and shape properties to return None when there is no rank. However I will still keep has_rank to be consistent with the mlir::ShapedTypeComponents interface as well as other python classes such as ShapedType. Let me know if you really want me to remove it.

[ataheridezfouli-groq]

Address various minor comments

[ataheridezfouli-groq]

Add hasRank to CAPI callback

[makslevental]

mlir/lib/Bindings/Python/IRInterfaces.cpp

330–340

nit: I believe all of these def_statics can be def_classmethods like this

cpp
      .def_classmethod(
          "get",
          [](const py::object &cls, MlirType elementType) {
            return cls(type);
          },

[ataheridezfouli-groq]

mlir/lib/Bindings/Python/IRInterfaces.cpp

330–340

py::class_ does not have a def_classmethod as seen here: https://github.com/pybind/pybind11/blob/cca4c51ca463ea02fa504331ff21bc313c80c7f3/include/pybind11/pybind11.h#L1496

In the MLIR python bindings, there is a hacky way to define classmethod in lib/Bindings/Python/IRCore.cpp::classmethod. But I believe this method just sets the attribute of an object, which in my case I can't have 3 attributes with the same name.

In general I agree with you that factory functions in python should use classmethod but this seems to be a limitation of pybind11. And I believe that's the reason why all the other py::class_ objects in MLIR use a static get method.

[makslevental]

mlir/lib/Bindings/Python/IRInterfaces.cpp

330–340

El Oh El I thought those uses of def_classmethod (in mlir_*_subclass) were via a pybind feature - I didn't realize they were actually members of pure_subclass. Alright carry on!

[ataheridezfouli-groq]

mlir/lib/Bindings/Python/IRInterfaces.cpp

330–340

Haha no worries! I learned something new as well!

[ataheridezfouli-groq]

Rebase

mlir/include/mlir-c/Interfaces.h

 mlirOperationImplementsInterfaceStatic(MlirStringRef operationName,
                                        MlirContext context,
                                        MlirTypeID interfaceTypeID);
 
 //===----------------------------------------------------------------------===//
 // InferTypeOpInterface.
 //===----------------------------------------------------------------------===//
 
 /// Returns the interface TypeID of the InferTypeOpInterface.
 MLIR_CAPI_EXPORTED MlirTypeID mlirInferTypeOpInterfaceTypeID();
 
 /// These callbacks are used to return multiple types from functions while
 /// transferring ownership to the caller. The first argument is the number of
 /// consecutive elements pointed to by the second argument. The third argument
 /// is an opaque pointer forwarded to the callback by the caller.
 typedef void (*MlirTypesCallback)(intptr_t, MlirType *, void *);
 
 /// Infers the return types of the operation identified by its canonical given
 /// the arguments that will be supplied to its generic builder. Calls `callback`
 /// with the types of inferred arguments, potentially several times, on success.
 /// Returns failure otherwise.
 MLIR_CAPI_EXPORTED MlirLogicalResult mlirInferTypeOpInterfaceInferReturnTypes(
     MlirStringRef opName, MlirContext context, MlirLocation location,
     intptr_t nOperands, MlirValue *operands, MlirAttribute attributes,
     void *properties, intptr_t nRegions, MlirRegion *regions,
     MlirTypesCallback callback, void *userData);
 
+//===----------------------------------------------------------------------===//
+// InferShapedTypeOpInterface.
+//===----------------------------------------------------------------------===//
+
+/// Returns the interface TypeID of the InferShapedTypeOpInterface.
+MLIR_CAPI_EXPORTED MlirTypeID mlirInferShapedTypeOpInterfaceTypeID();
+
+/// These callbacks are used to return multiple shaped type components from
+/// functions while transferring ownership to the caller. The first argument is
+/// the has rank boolean followed by the the rank and a pointer to the shape
+/// (if applicable). The next argument is the element type, then the attribute.
+/// The last argument is an opaque pointer forwarded to the callback by the
+/// caller. This callback will be called potentially multiple times for each
+/// shaped type components.
+typedef void (*MlirShapedTypeComponentsCallback)(bool, intptr_t,
+                                                 const int64_t *, MlirType,
+                                                 MlirAttribute, void *);
+
+/// Infers the return shaped type components of the operation. Calls `callback`
+/// with the types of inferred arguments on success. Returns failure otherwise.
+MLIR_CAPI_EXPORTED MlirLogicalResult
+mlirInferShapedTypeOpInterfaceInferReturnTypes(
+    MlirStringRef opName, MlirContext context, MlirLocation location,
+    intptr_t nOperands, MlirValue *operands, MlirAttribute attributes,
+    void *properties, intptr_t nRegions, MlirRegion *regions,
+    MlirShapedTypeComponentsCallback callback, void *userData);
+
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_INTERFACES_H

mlir/lib/Bindings/Python/IRInterfaces.cpp

 //===- IRInterfaces.cpp - MLIR IR interfaces pybind -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#include <utility>
 #include <optional>
+#include <utility>
 
 #include "IRModule.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/Interfaces.h"
 #include "llvm/ADT/STLExtras.h"
 
 namespace py = pybind11;
 

 constexpr static const char *opviewDoc =
     R"(Returns an OpView subclass _instance_ for which the interface was
 constructed)";
 
 constexpr static const char *inferReturnTypesDoc =
     R"(Given the arguments required to build an operation, attempts to infer
 its return types. Raises ValueError on failure.)";
 
+constexpr static const char *inferReturnTypeComponentsDoc =
+    R"(Given the arguments required to build an operation, attempts to infer
+its return shaped type components. Raises ValueError on failure.)";
+
+namespace {
+
+/// Takes in an optional ist of operands and converts them into a SmallVector
+/// of MlirVlaues. Returns an empty SmallVector if the list is empty.
+llvm::SmallVector<MlirValue> wrapOperands(std::optional<py::list> operandList) {
+  llvm::SmallVector<MlirValue> mlirOperands;
+
+  if (!operandList || operandList->empty()) {
+    return mlirOperands;
+  }
+
+  // Note: as the list may contain other lists this may not be final size.
+  mlirOperands.reserve(operandList->size());
+  for (const auto &&it : llvm::enumerate(*operandList)) {
+    PyValue *val;
+    try {
+      val = py::cast<PyValue *>(it.value());
+      if (!val)
+        throw py::cast_error();
+      mlirOperands.push_back(val->get());
+      continue;
+    } catch (py::cast_error &err) {
+      // Intentionally unhandled to try sequence below first.
+      (void)err;
+    }
+
+    try {
+      auto vals = py::cast<py::sequence>(it.value());
+      for (py::object v : vals) {
+        try {
+          val = py::cast<PyValue *>(v);
+          if (!val)
+            throw py::cast_error();
+          mlirOperands.push_back(val->get());
+        } catch (py::cast_error &err) {
+          throw py::value_error(
+              (llvm::Twine("Operand ") + llvm::Twine(it.index()) +
+               " must be a Value or Sequence of Values (" + err.what() + ")")
+                  .str());
+        }
+      }
+      continue;
+    } catch (py::cast_error &err) {
+      throw py::value_error((llvm::Twine("Operand ") + llvm::Twine(it.index()) +
+                             " must be a Value or Sequence of Values (" +
+                             err.what() + ")")
+                                .str());
+    }
+
+    throw py::cast_error();
+  }
+
+  return mlirOperands;
+}
+
+/// Takes in an optional vector of PyRegions and returns a SmallVector of
+/// MlirRegion. Returns an empty SmallVector if the list is empty.
+llvm::SmallVector<MlirRegion>
+wrapRegions(std::optional<std::vector<PyRegion>> regions) {
+  llvm::SmallVector<MlirRegion> mlirRegions;
+
+  if (regions) {
+    mlirRegions.reserve(regions->size());
+    for (PyRegion &region : *regions) {
+      mlirRegions.push_back(region);
+    }
+  }
+
+  return mlirRegions;
+}
+
+} // namespace
+
 /// CRTP base class for Python classes representing MLIR Op interfaces.
 /// Interface hierarchies are flat so no base class is expected here. The
 /// derived class is expected to define the following static fields:
 ///  - `const char *pyClassName` - the name of the Python class to create;
 ///  - `GetTypeIDFunctionTy getInterfaceID` - the function producing the TypeID
 ///    of the interface.
 /// Derived classes may redefine the `bindDerived(ClassTy &)` method to bind
 /// interface-specific methods.
if (operation != nullptr) {
         std::string msg = "the operation does not implement ";
         throw py::value_error(msg + ConcreteIface::pyClassName);
       }
     }
   }
 
   /// Creates the Python bindings for this class in the given module.
   static void bind(py::module &m) {
-    py::class_<ConcreteIface> cls(m, "InferTypeOpInterface",
+    py::class_<ConcreteIface> cls(m, ConcreteIface::pyClassName,
                                   py::module_local());
     cls.def(py::init<py::object, DefaultingPyMlirContext>(), py::arg("object"),
             py::arg("context") = py::none(), constructorDoc)
         .def_property_readonly("operation",
                                &PyConcreteOpInterface::getOperationObject,
                                operationDoc)
         .def_property_readonly("opview", &PyConcreteOpInterface::getOpView,
                                opviewDoc);
public:
   const std::string &getOpName() { return opName; }
 
 private:
   PyOperation *operation = nullptr;
   std::string opName;
   py::object obj;
 };
 
-/// Python wrapper for InterTypeOpInterface. This interface has only static
+/// Python wrapper for InferTypeOpInterface. This interface has only static
 /// methods.
 class PyInferTypeOpInterface
     : public PyConcreteOpInterface<PyInferTypeOpInterface> {
 public:
   using PyConcreteOpInterface<PyInferTypeOpInterface>::PyConcreteOpInterface;
 
   constexpr static const char *pyClassName = "InferTypeOpInterface";
   constexpr static GetTypeIDFunctionTy getInterfaceID =
public:
   /// Given the arguments required to build an operation, attempts to infer its
   /// return types. Throws value_error on failure.
   std::vector<PyType>
   inferReturnTypes(std::optional<py::list> operandList,
                    std::optional<PyAttribute> attributes, void *properties,
                    std::optional<std::vector<PyRegion>> regions,
                    DefaultingPyMlirContext context,
                    DefaultingPyLocation location) {
-    llvm::SmallVector<MlirValue> mlirOperands;
-    llvm::SmallVector<MlirRegion> mlirRegions;
-
-    if (operandList && !operandList->empty()) {
-      // Note: as the list may contain other lists this may not be final size.
-      mlirOperands.reserve(operandList->size());
-      for (const auto& it : llvm::enumerate(*operandList)) {
-        PyValue* val;
-        try {
-          val = py::cast<PyValue *>(it.value());
-          if (!val)
-            throw py::cast_error();
-          mlirOperands.push_back(val->get());
-          continue;
-        } catch (py::cast_error &err) {
-          // Intentionally unhandled to try sequence below first.
-          (void)err;
-        }
-
-        try {
-          auto vals = py::cast<py::sequence>(it.value());
-          for (py::object v : vals) {
-            try {
-              val = py::cast<PyValue *>(v);
-              if (!val)
-                throw py::cast_error();
-              mlirOperands.push_back(val->get());
-            } catch (py::cast_error &err) {
-              throw py::value_error(
-                  (llvm::Twine("Operand ") + llvm::Twine(it.index()) +
-                   " must be a Value or Sequence of Values (" + err.what() +
-                   ")")
-                      .str());
-            }
-          }
-          continue;
-        } catch (py::cast_error &err) {
-          throw py::value_error(
-              (llvm::Twine("Operand ") + llvm::Twine(it.index()) +
-               " must be a Value or Sequence of Values (" + err.what() + ")")
-                  .str());
-        }
-
-        throw py::cast_error();
-      }
-    }
-
-    if (regions) {
-      mlirRegions.reserve(regions->size());
-      for (PyRegion &region : *regions) {
-        mlirRegions.push_back(region);
-      }
-    }
+    llvm::SmallVector<MlirValue> mlirOperands = wrapOperands(operandList);
+    llvm::SmallVector<MlirRegion> mlirRegions = wrapRegions(regions);
 
     std::vector<PyType> inferredTypes;
     PyMlirContext &pyContext = context.resolve();
     AppendResultsCallbackData data{inferredTypes, pyContext};
     MlirStringRef opNameRef =
         mlirStringRefCreate(getOpName().data(), getOpName().length());
     MlirAttribute attributeDict =
         attributes ? attributes->get() : mlirAttributeGetNull();
cls.def("inferReturnTypes", &PyInferTypeOpInterface::inferReturnTypes,
             py::arg("operands") = py::none(),
             py::arg("attributes") = py::none(),
             py::arg("properties") = py::none(), py::arg("regions") = py::none(),
             py::arg("context") = py::none(), py::arg("loc") = py::none(),
             inferReturnTypesDoc);
   }
 };
 
-void populateIRInterfaces(py::module &m) { PyInferTypeOpInterface::bind(m); }
+/// Wrapper around an shaped type components.
+class PyShapedTypeComponents {
+public:
+  PyShapedTypeComponents(MlirType elementType) : elementType(elementType) {}
+  PyShapedTypeComponents(py::list shape, MlirType elementType)
+      : shape(shape), elementType(elementType), ranked(true) {}
+  PyShapedTypeComponents(py::list shape, MlirType elementType,
+                         MlirAttribute attribute)
+      : shape(shape), elementType(elementType), attribute(attribute),
+        ranked(true) {}
+  PyShapedTypeComponents(PyShapedTypeComponents &) = delete;
+  PyShapedTypeComponents(PyShapedTypeComponents &&other)
+      : shape(other.shape), elementType(other.elementType),
+        attribute(other.attribute), ranked(other.ranked) {}
+
+  static void bind(py::module &m) {
+    py::class_<PyShapedTypeComponents>(m, "ShapedTypeComponents",
+                                       py::module_local())
+        .def_property_readonly(
+            "element_type",
+            [](PyShapedTypeComponents &self) {
+              return PyType(PyMlirContext::forContext(
+                                mlirTypeGetContext(self.elementType)),
+                            self.elementType);
+            },
+            "Returns the element type of the shaped type components.")
+        .def_static(
+            "get",
+            [](PyType &elementType) {
+              return PyShapedTypeComponents(elementType);
+            },
+            py::arg("element_type"),
+            "Create an shaped type components object with only the element "
+            "type.")
+        .def_static(
+            "get",
+            [](py::list shape, PyType &elementType) {
+              return PyShapedTypeComponents(shape, elementType);
+            },
+            py::arg("shape"), py::arg("element_type"),
+            "Create a ranked shaped type components object.")
+        .def_static(
+            "get",
+            [](py::list shape, PyType &elementType, PyAttribute &attribute) {
+              return PyShapedTypeComponents(shape, elementType, attribute);
+            },
+            py::arg("shape"), py::arg("element_type"), py::arg("attribute"),
+            "Create a ranked shaped type components object with attribute.")
+        .def_property_readonly(
+            "has_rank",
+            [](PyShapedTypeComponents &self) -> bool { return self.ranked; },
+            "Returns whether the given shaped type component is ranked.")
+        .def_property_readonly(
+            "rank",
+            [](PyShapedTypeComponents &self) -> py::object {
+              if (!self.ranked) {
+                return py::none();
+              }
+              return py::int_(self.shape.size());
+            },
+            "Returns the rank of the given ranked shaped type components. If "
+            "the shaped type components does not have a rank, None is "
+            "returned.")
+        .def_property_readonly(
+            "shape",
+            [](PyShapedTypeComponents &self) -> py::object {
+              if (!self.ranked) {
+                return py::none();
+              }
+              return py::list(self.shape);
+            },
+            "Returns the shape of the ranked shaped type components as a list "
+            "of integers. Returns none if the shaped type component does not "
+            "have a rank.");
+  }
+
+  pybind11::object getCapsule();
+  static PyShapedTypeComponents createFromCapsule(pybind11::object capsule);
+
+private:
+  py::list shape;
+  MlirType elementType;
+  MlirAttribute attribute;
+  bool ranked{false};
+};
+
+/// Python wrapper for InferShapedTypeOpInterface. This interface has only
+/// static methods.
+class PyInferShapedTypeOpInterface
+    : public PyConcreteOpInterface<PyInferShapedTypeOpInterface> {
+public:
+  using PyConcreteOpInterface<
+      PyInferShapedTypeOpInterface>::PyConcreteOpInterface;
+
+  constexpr static const char *pyClassName = "InferShapedTypeOpInterface";
+  constexpr static GetTypeIDFunctionTy getInterfaceID =
+      &mlirInferShapedTypeOpInterfaceTypeID;
+
+  /// C-style user-data structure for type appending callback.
+  struct AppendResultsCallbackData {
+    std::vector<PyShapedTypeComponents> &inferredShapedTypeComponents;
+  };
+
+  /// Appends the shaped type components provided as unpacked shape, element
+  /// type, attribute to the user-data.
+  static void appendResultsCallback(bool hasRank, intptr_t rank,
+                                    const int64_t *shape, MlirType elementType,
+                                    MlirAttribute attribute, void *userData) {
+    auto *data = static_cast<AppendResultsCallbackData *>(userData);
+    if (!hasRank) {
+      data->inferredShapedTypeComponents.emplace_back(elementType);
+    } else {
+      py::list shapeList;
+      for (intptr_t i = 0; i < rank; ++i) {
+        shapeList.append(shape[i]);
+      }
+      data->inferredShapedTypeComponents.emplace_back(shapeList, elementType,
+                                                      attribute);
+    }
+  }
+
+  /// Given the arguments required to build an operation, attempts to infer the
+  /// shaped type components. Throws value_error on failure.
+  std::vector<PyShapedTypeComponents> inferReturnTypeComponents(
+      std::optional<py::list> operandList,
+      std::optional<PyAttribute> attributes, void *properties,
+      std::optional<std::vector<PyRegion>> regions,
+      DefaultingPyMlirContext context, DefaultingPyLocation location) {
+    llvm::SmallVector<MlirValue> mlirOperands = wrapOperands(operandList);
+    llvm::SmallVector<MlirRegion> mlirRegions = wrapRegions(regions);
+
+    std::vector<PyShapedTypeComponents> inferredShapedTypeComponents;
+    PyMlirContext &pyContext = context.resolve();
+    AppendResultsCallbackData data{inferredShapedTypeComponents};
+    MlirStringRef opNameRef =
+        mlirStringRefCreate(getOpName().data(), getOpName().length());
+    MlirAttribute attributeDict =
+        attributes ? attributes->get() : mlirAttributeGetNull();
+
+    MlirLogicalResult result = mlirInferShapedTypeOpInterfaceInferReturnTypes(
+        opNameRef, pyContext.get(), location.resolve(), mlirOperands.size(),
+        mlirOperands.data(), attributeDict, properties, mlirRegions.size(),
+        mlirRegions.data(), &appendResultsCallback, &data);
+
+    if (mlirLogicalResultIsFailure(result)) {
+      throw py::value_error("Failed to infer result shape type components");
+    }
+
+    return inferredShapedTypeComponents;
+  }
+
+  static void bindDerived(ClassTy &cls) {
+    cls.def("inferReturnTypeComponents",
+            &PyInferShapedTypeOpInterface::inferReturnTypeComponents,
+            py::arg("operands") = py::none(),
+            py::arg("attributes") = py::none(), py::arg("regions") = py::none(),
+            py::arg("properties") = py::none(), py::arg("context") = py::none(),
+            py::arg("loc") = py::none(), inferReturnTypeComponentsDoc);
+  }
+};
+
+void populateIRInterfaces(py::module &m) {
+  PyInferTypeOpInterface::bind(m);
+  PyShapedTypeComponents::bind(m);
+  PyInferShapedTypeOpInterface::bind(m);
+}
 
 } // namespace python
 } // namespace mlir

mlir/lib/CAPI/Interfaces/Interfaces.cpp

 
 
 //===- Interfaces.cpp - C Interface for MLIR Interfaces -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir-c/Interfaces.h"
 
 #include "mlir/CAPI/IR.h"
+#include "mlir/CAPI/Interfaces.h"
 #include "mlir/CAPI/Support.h"
 #include "mlir/CAPI/Wrap.h"
+#include "mlir/IR/ValueRange.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "llvm/ADT/ScopeExit.h"
 #include <optional>
 
 using namespace mlir;
 
+namespace {
+
+std::optional<RegisteredOperationName>
+getRegisteredOperationName(MlirContext context, MlirStringRef opName) {
+  StringRef name(opName.data, opName.length);
+  std::optional<RegisteredOperationName> info =
+      RegisteredOperationName::lookup(name, unwrap(context));
+  return info;
+}
+
+std::optional<Location> maybeGetLocation(MlirLocation location) {
+  std::optional<Location> maybeLocation;
+  if (!mlirLocationIsNull(location))
+    maybeLocation = unwrap(location);
+  return maybeLocation;
+}
+
+SmallVector<Value> unwrapOperands(intptr_t nOperands, MlirValue *operands) {
+  SmallVector<Value> unwrappedOperands;
+  (void)unwrapList(nOperands, operands, unwrappedOperands);
+  return unwrappedOperands;
+}
+
+DictionaryAttr unwrapAttributes(MlirAttribute attributes) {
+  DictionaryAttr attributeDict;
+  if (!mlirAttributeIsNull(attributes))
+    attributeDict = unwrap(attributes).cast<DictionaryAttr>();
+  return attributeDict;
+}
+
+SmallVector<std::unique_ptr<Region>> unwrapRegions(intptr_t nRegions,
+                                                   MlirRegion *regions) {
+  // Create a vector of unique pointers to regions and make sure they are not
+  // deleted when exiting the scope. This is a hack caused by C++ API expecting
+  // an list of unique pointers to regions (without ownership transfer
+  // semantics) and C API making ownership transfer explicit.
+  SmallVector<std::unique_ptr<Region>> unwrappedRegions;
+  unwrappedRegions.reserve(nRegions);
+  for (intptr_t i = 0; i < nRegions; ++i)
+    unwrappedRegions.emplace_back(unwrap(*(regions + i)));
+  auto cleaner = llvm::make_scope_exit([&]() {
+    for (auto &region : unwrappedRegions)
+      region.release();
+  });
+  return unwrappedRegions;
+}
+
+} // namespace
+
 bool mlirOperationImplementsInterface(MlirOperation operation,
                                       MlirTypeID interfaceTypeID) {
   std::optional<RegisteredOperationName> info =
       unwrap(operation)->getRegisteredInfo();
   return info && info->hasInterface(unwrap(interfaceTypeID));
 }
 
 bool mlirOperationImplementsInterfaceStatic(MlirStringRef operationName,

 
 MlirLogicalResult mlirInferTypeOpInterfaceInferReturnTypes(
     MlirStringRef opName, MlirContext context, MlirLocation location,
     intptr_t nOperands, MlirValue *operands, MlirAttribute attributes,
     void *properties, intptr_t nRegions, MlirRegion *regions,
     MlirTypesCallback callback, void *userData) {
   StringRef name(opName.data, opName.length);
   std::optional<RegisteredOperationName> info =
-      RegisteredOperationName::lookup(name, unwrap(context));
+      getRegisteredOperationName(context, opName);
   if (!info)
     return mlirLogicalResultFailure();
 
-  std::optional<Location> maybeLocation;
-  if (!mlirLocationIsNull(location))
-    maybeLocation = unwrap(location);
-  SmallVector<Value> unwrappedOperands;
-  (void)unwrapList(nOperands, operands, unwrappedOperands);
-  DictionaryAttr attributeDict;
-  if (!mlirAttributeIsNull(attributes))
-    attributeDict = unwrap(attributes).cast<DictionaryAttr>();
-
-  // Create a vector of unique pointers to regions and make sure they are not
-  // deleted when exiting the scope. This is a hack caused by C++ API expecting
-  // an list of unique pointers to regions (without ownership transfer
-  // semantics) and C API making ownership transfer explicit.
-  SmallVector<std::unique_ptr<Region>> unwrappedRegions;
-  unwrappedRegions.reserve(nRegions);
-  for (intptr_t i = 0; i < nRegions; ++i)
-    unwrappedRegions.emplace_back(unwrap(*(regions + i)));
-  auto cleaner = llvm::make_scope_exit([&]() {
-    for (auto &region : unwrappedRegions)
-      region.release();
-  });
+  std::optional<Location> maybeLocation = maybeGetLocation(location);
+  SmallVector<Value> unwrappedOperands = unwrapOperands(nOperands, operands);
+  DictionaryAttr attributeDict = unwrapAttributes(attributes);
+  SmallVector<std::unique_ptr<Region>> unwrappedRegions =
+      unwrapRegions(nRegions, regions);
 
   SmallVector<Type> inferredTypes;
   if (failed(info->getInterface<InferTypeOpInterface>()->inferReturnTypes(
           unwrap(context), maybeLocation, unwrappedOperands, attributeDict,
           properties, unwrappedRegions, inferredTypes)))
     return mlirLogicalResultFailure();
 
   SmallVector<MlirType> wrappedInferredTypes;
   wrappedInferredTypes.reserve(inferredTypes.size());
   for (Type t : inferredTypes)
     wrappedInferredTypes.push_back(wrap(t));
   callback(wrappedInferredTypes.size(), wrappedInferredTypes.data(), userData);
   return mlirLogicalResultSuccess();
 }
+
+MlirTypeID mlirInferShapedTypeOpInterfaceTypeID() {
+  return wrap(InferShapedTypeOpInterface::getInterfaceID());
+}
+
+MlirLogicalResult mlirInferShapedTypeOpInterfaceInferReturnTypes(
+    MlirStringRef opName, MlirContext context, MlirLocation location,
+    intptr_t nOperands, MlirValue *operands, MlirAttribute attributes,
+    void *properties, intptr_t nRegions, MlirRegion *regions,
+    MlirShapedTypeComponentsCallback callback, void *userData) {
+  std::optional<RegisteredOperationName> info =
+      getRegisteredOperationName(context, opName);
+  if (!info)
+    return mlirLogicalResultFailure();
+
+  std::optional<Location> maybeLocation = maybeGetLocation(location);
+  SmallVector<Value> unwrappedOperands = unwrapOperands(nOperands, operands);
+  DictionaryAttr attributeDict = unwrapAttributes(attributes);
+  SmallVector<std::unique_ptr<Region>> unwrappedRegions =
+      unwrapRegions(nRegions, regions);
+
+  SmallVector<ShapedTypeComponents> inferredTypeComponents;
+  if (failed(info->getInterface<InferShapedTypeOpInterface>()
+                 ->inferReturnTypeComponents(
+                     unwrap(context), maybeLocation,
+                     mlir::ValueRange(llvm::ArrayRef(unwrappedOperands)),
+                     attributeDict, properties, unwrappedRegions,
+                     inferredTypeComponents)))
+    return mlirLogicalResultFailure();
+
+  bool hasRank;
+  intptr_t rank;
+  const int64_t *shapeData;
+  for (ShapedTypeComponents t : inferredTypeComponents) {
+    if (t.hasRank()) {
+      hasRank = true;
+      rank = t.getDims().size();
+      shapeData = t.getDims().data();
+    } else {
+      hasRank = false;
+      rank = 0;
+      shapeData = nullptr;
+    }
+    callback(hasRank, rank, shapeData, wrap(t.getElementType()),
+             wrap(t.getAttribute()), userData);
+  }
+  return mlirLogicalResultSuccess();
+}

mlir/python/mlir/_mlir_libs/_mlir/ir.pyi

__all__ = [
     "Float8E5M2FNUZType",
     "F16Type",
     "F32Type",
     "F64Type",
     "FlatSymbolRefAttr",
     "FloatAttr",
     "FunctionType",
     "IndexType",
+    "InferShapedTypeOpInterface",
     "InferTypeOpInterface",
     "InsertionPoint",
     "IntegerAttr",
     "IntegerSet",
     "IntegerSetConstraint",
     "IntegerSetConstraintList",
     "IntegerType",
     "Location",
__all__ = [
     "Operation",
     "OperationIterator",
     "OperationList",
     "RankedTensorType",
     "Region",
     "RegionIterator",
     "RegionSequence",
     "ShapedType",
+    "ShapedTypeComponents",
     "StringAttr",
     "SymbolTable",
     "TupleType",
     "Type",
     "TypeAttr",
     "UnitAttr",
     "UnrankedMemRefType",
     "UnrankedTensorType",

 # TODO: Auto-generated. Audit and fix.
 class IndexType(Type):
     def __init__(self, cast_from_type: Type) -> None: ...
     @staticmethod
     def get(*args, **kwargs) -> IndexType: ...
     @staticmethod
     def isinstance(arg: Any) -> bool: ...
 
+class InferShapedTypeOpInterface:
+    def __init__(self, object: object, context: Optional[Context] = None) -> None: ...
+    def inferReturnTypeComponents(self, operands: Optional[List] = None, attributes: Optional[Attribute] = None, properties = None, regions: Optional[List[Region]] = None, context: Optional[Context] = None, loc: Optional[Location] = None) -> List[ShapedTypeComponents]: ...
+    @property
+    def operation(self) -> Operation: ...
+    @property
+    def opview(self) -> OpView: ...
+
 class InferTypeOpInterface:
     def __init__(self, object: object, context: Optional[Context] = None) -> None: ...
-    def inferReturnTypes(self, operands: Optional[List] = None, attributes: Optional[Attribute] = None, regions: Optional[List[Region]] = None, context: Optional[Context] = None, loc: Optional[Location] = None) -> List[Type]: ...
+    def inferReturnTypes(self, operands: Optional[List] = None, attributes: Optional[Attribute] = None, properties = None, regions: Optional[List[Region]] = None, context: Optional[Context] = None, loc: Optional[Location] = None) -> List[Type]: ...
     @property
     def operation(self) -> Operation: ...
     @property
     def opview(self) -> OpView: ...
 
 class InsertionPoint:
     current: ClassVar[InsertionPoint] = ...  # read-only
     @overload
class ShapedType(Type):
     def has_rank(self) -> bool: ...
     @property
     def has_static_shape(self) -> bool: ...
     @property
     def rank(self) -> int: ...
     @property
     def shape(self) -> List[int]: ...
 
+class ShapedTypeComponents:
+    @property
+    def element_type(self) -> Type: ...
+    @staticmethod
+    def get(*args, **kwargs) -> ShapedTypeComponents: ...
+    @property
+    def has_rank(self) -> bool: ...
+    @property
+    def rank(self) -> int: ...
+    @property
+    def shape(self) -> List[int]: ...
+
 # TODO: Auto-generated. Audit and fix.
 class StringAttr(Attribute):
     def __init__(self, cast_from_attr: Attribute) -> None: ...
     @staticmethod
     def get(*args, **kwargs) -> Any: ...
     @staticmethod
     def get_typed(*args, **kwargs) -> Any: ...
     @staticmethod

mlir/test/python/dialects/python_test.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
+import mlir.dialects.func as func
 import mlir.dialects.python_test as test
 import mlir.dialects.tensor as tensor
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   return f
 
with InsertionPoint(module.body):
       print(Value(t))
 
       tt = Tensor(t)
       # CHECK: Tensor(%{{.*}} = tensor.empty() : tensor<10x10xi8>)
       print(tt)
 
       # CHECK: False
       print(tt.is_null())
+
+
+# CHECK-LABEL: TEST: inferReturnTypeComponents
+@run
+def inferReturnTypeComponents():
+    with Context() as ctx, Location.unknown(ctx):
+        test.register_python_test_dialect(ctx)
+        module = Module.create()
+        i32 = IntegerType.get_signless(32)
+        with InsertionPoint(module.body):
+            resultType = UnrankedTensorType.get(i32)
+            operandTypes = [
+                RankedTensorType.get([1, 3, 10, 10], i32),
+                UnrankedTensorType.get(i32),
+            ]
+            f = func.FuncOp(
+                "test_inferReturnTypeComponents", (operandTypes, [resultType])
+            )
+            entry_block = Block.create_at_start(f.operation.regions[0], operandTypes)
+            with InsertionPoint(entry_block):
+                ranked_op = test.InferShapedTypeComponentsOp(
+                    resultType, entry_block.arguments[0]
+                )
+                unranked_op = test.InferShapedTypeComponentsOp(
+                    resultType, entry_block.arguments[1]
+                )
+
+        # CHECK: has rank: True
+        # CHECK: rank: 4
+        # CHECK: element type: i32
+        # CHECK: shape: [1, 3, 10, 10]
+        iface = InferShapedTypeOpInterface(ranked_op)
+        shaped_type_components = iface.inferReturnTypeComponents(
+            operands=[ranked_op.operand]
+        )[0]
+        print("has rank:", shaped_type_components.has_rank)
+        print("rank:", shaped_type_components.rank)
+        print("element type:", shaped_type_components.element_type)
+        print("shape:", shaped_type_components.shape)
+
+        # CHECK: has rank: False
+        # CHECK: rank: None
+        # CHECK: element type: i32
+        # CHECK: shape: None
+        iface = InferShapedTypeOpInterface(unranked_op)
+        shaped_type_components = iface.inferReturnTypeComponents(
+          operands=[unranked_op.operand]
+        )[0]
+        print("has rank:", shaped_type_components.has_rank)
+        print("rank:", shaped_type_components.rank)
+        print("element type:", shaped_type_components.element_type)
+        print("shape:", shaped_type_components.shape)

mlir/test/python/python_test_ops.td

 }
 
 // If all result types are buildable, the InferTypeOpInterface is implied and is
 // autogenerated by C++ ODS.
 def InferResultsImpliedOp : TestOp<"infer_results_implied_op"> {
   let results = (outs I32:$integer, F64:$flt, Index:$index);
 }
 
+def InferShapedTypeComponentsOp : TestOp<"infer_shaped_type_components_op",
+  [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+                             ["inferReturnTypeComponents"]>]> {
+  let arguments = (ins AnyTensor:$operand);
+  let results = (outs AnyTensor:$result);
+
+  let extraClassDefinition = [{
+    ::mlir::LogicalResult $cppClass::inferReturnTypeComponents(
+      ::mlir::MLIRContext *context, ::std::optional<::mlir::Location> location,
+      ::mlir::ValueShapeRange operands, ::mlir::DictionaryAttr attributes,
+      ::mlir::OpaqueProperties properties, ::mlir::RegionRange regions,
+      ::llvm::SmallVectorImpl<
+        ::mlir::ShapedTypeComponents>& inferredShapedTypeComponents) {
+      $cppClass::Adaptor adaptor(operands, attributes, properties, regions);
+      auto operandType =
+          adaptor.getOperand().getType().cast<::mlir::ShapedType>();
+      if (operandType.hasRank()) {
+        inferredShapedTypeComponents.emplace_back(operandType.getShape(),
+            operandType.getElementType());
+      } else {
+        inferredShapedTypeComponents.emplace_back(operandType.getElementType());
+      }
+      return ::mlir::success();
+    }
+  }];
+}
+
 def SameOperandAndResultTypeOp : TestOp<"same_operand_and_result_type_op",
                                         [SameOperandsAndResultType]> {
   let arguments = (ins Variadic<AnyType>);
   let results = (outs AnyType:$one, AnyType:$two);
 }
 
 def FirstAttrDeriveTypeAttrOp : TestOp<"first_attr_derive_type_attr_op",
                                [FirstAttrDerivedResultType]> {