[MLIR][CAPI][python] expose the python binding for linalgOp.getIndexingMaps (#136054)

This PR is mainly about exposing the python bindings for
`linalgOp.getIndexingMaps`.

---------

Signed-off-by: Bangtian Liu <[email protected]>

Author: bangtianliu

mlir/include/mlir-c/Dialect/Linalg.h

@@ -50,6 +50,9 @@ typedef struct MlirLinalgConvolutionDimensions {
 MLIR_CAPI_EXPORTED MlirLinalgConvolutionDimensions
 mlirLinalgInferConvolutionDimensions(MlirOperation op);
 
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirLinalgGetIndexingMapsAttribute(MlirOperation op);
+
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Linalg, linalg);
 
 #ifdef __cplusplus

mlir/lib/Bindings/Python/DialectLinalg.cpp

@@ -120,6 +120,16 @@ static void populateDialectLinalgSubmodule(nb::module_ m) {
 
   m.def("infer_convolution_dimensions", &InferConvolutionDimensions,
         "Infers convolution dimensions", nb::arg("op"));
+
+  m.def(
+      "get_indexing_maps",
+      [](MlirOperation op) -> std::optional<MlirAttribute> {
+        MlirAttribute attr = mlirLinalgGetIndexingMapsAttribute(op);
+        if (mlirAttributeIsNull(attr))
+          return std::nullopt;
+        return attr;
+      },
+      "Returns the indexing_maps attribute for a linalg op.");
 }
 
 NB_MODULE(_mlirDialectsLinalg, m) {

mlir/lib/CAPI/Dialect/Linalg.cpp

@@ -120,4 +120,14 @@ mlirLinalgInferConvolutionDimensions(MlirOperation op) {
   return result;
 }
 
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirLinalgGetIndexingMapsAttribute(MlirOperation op) {
+  auto linalgOp = llvm::dyn_cast<mlir::linalg::LinalgOp>(unwrap(op));
+  if (!linalgOp)
+    return MlirAttribute{nullptr};
+
+  ArrayAttr attr = linalgOp.getIndexingMaps();
+  return wrap(attr);
+}
+
 MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Linalg, linalg, LinalgDialect)