[mlir][CAPI][python] expose the python bindings for linalg::isaConvolutionOpInterface and linalg::inferConvolutionDims (#135253)

This PR is mainly about exposing the python bindings for
`linalg::isaConvolutionOpInterface` and `linalg::inferConvolutionDims`.

---------

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

Author: bangtianliu

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

@@ -22,7 +22,7 @@ extern "C" {
 MLIR_CAPI_EXPORTED void
 mlirLinalgFillBuiltinNamedOpRegion(MlirOperation mlirOp);
 
-MLIR_CAPI_EXPORTED bool mlirLinalgIsContractionOp(MlirOperation op);
+MLIR_CAPI_EXPORTED bool mlirLinalgIsAContractionOp(MlirOperation op);
 
 struct MlirLinalgContractionDimensions {
   MlirAttribute batch;
@@ -34,6 +34,22 @@ struct MlirLinalgContractionDimensions {
 MLIR_CAPI_EXPORTED MlirLinalgContractionDimensions
 mlirLinalgInferContractionDimensions(MlirOperation op);
 
+MLIR_CAPI_EXPORTED bool mlirLinalgIsAConvolutionOp(MlirOperation op);
+
+struct MlirLinalgConvolutionDimensions {
+  MlirAttribute batch;
+  MlirAttribute outputImage;
+  MlirAttribute outputChannel;
+  MlirAttribute filterLoop;
+  MlirAttribute inputChannel;
+  MlirAttribute depth;
+  MlirAttribute strides;
+  MlirAttribute dilations;
+};
+
+MLIR_CAPI_EXPORTED MlirLinalgConvolutionDimensions
+mlirLinalgInferConvolutionDimensions(MlirOperation op);
+
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Linalg, linalg);
 
 #ifdef __cplusplus

mlir/lib/Bindings/Python/DialectLinalg.cpp

@@ -28,6 +28,26 @@ InferContractionDimensions(MlirOperation op) {
   return dims;
 }
 
+static std::optional<MlirLinalgConvolutionDimensions>
+InferConvolutionDimensions(MlirOperation op) {
+  MlirLinalgConvolutionDimensions dims =
+      mlirLinalgInferConvolutionDimensions(op);
+
+  // Detect "empty" result. This occurs when `op` is not a convolution op,
+  // or when `linalg::inferConvolutionDims` fails.
+  if (mlirAttributeIsNull(dims.batch) &&
+      mlirAttributeIsNull(dims.outputImage) &&
+      mlirAttributeIsNull(dims.outputChannel) &&
+      mlirAttributeIsNull(dims.filterLoop) &&
+      mlirAttributeIsNull(dims.inputChannel) &&
+      mlirAttributeIsNull(dims.depth) && mlirAttributeIsNull(dims.strides) &&
+      mlirAttributeIsNull(dims.dilations)) {
+    return std::nullopt;
+  }
+
+  return dims;
+}
+
 static void populateDialectLinalgSubmodule(nb::module_ m) {
   m.def(
       "fill_builtin_region",
@@ -36,7 +56,7 @@ static void populateDialectLinalgSubmodule(nb::module_ m) {
       "Fill the region for `op`, which is assumed to be a builtin named Linalg "
       "op.");
 
-  m.def("isa_contraction_op", &mlirLinalgIsContractionOp,
+  m.def("isa_contraction_op", &mlirLinalgIsAContractionOp,
         "Checks if the given operation is a Linalg contraction operation.",
         nb::arg("op"));
 
@@ -59,6 +79,47 @@ static void populateDialectLinalgSubmodule(nb::module_ m) {
         "Infers contraction dimensions (batch/m/n/k) for a Linalg contraction "
         "op.",
         nb::arg("op"));
+
+  m.def("isa_convolution_op", &mlirLinalgIsAConvolutionOp,
+        "Checks if the given operation is a Linalg convolution operation.",
+        nb::arg("op"));
+
+  nb::class_<MlirLinalgConvolutionDimensions>(m, "ConvolutionDimensions")
+      .def_prop_ro("batch",
+                   [](const MlirLinalgConvolutionDimensions &self) {
+                     return self.batch;
+                   })
+      .def_prop_ro("output_image",
+                   [](const MlirLinalgConvolutionDimensions &self) {
+                     return self.outputImage;
+                   })
+      .def_prop_ro("output_channel",
+                   [](const MlirLinalgConvolutionDimensions &self) {
+                     return self.outputChannel;
+                   })
+      .def_prop_ro("filter_loop",
+                   [](const MlirLinalgConvolutionDimensions &self) {
+                     return self.filterLoop;
+                   })
+      .def_prop_ro("input_channel",
+                   [](const MlirLinalgConvolutionDimensions &self) {
+                     return self.inputChannel;
+                   })
+      .def_prop_ro("depth",
+                   [](const MlirLinalgConvolutionDimensions &self) {
+                     return self.depth;
+                   })
+      .def_prop_ro("strides",
+                   [](const MlirLinalgConvolutionDimensions &self) {
+                     return self.strides;
+                   })
+      .def_prop_ro("dilations",
+                   [](const MlirLinalgConvolutionDimensions &self) {
+                     return self.dilations;
+                   });
+
+  m.def("infer_convolution_dimensions", &InferConvolutionDimensions,
+        "Infers convolution dimensions", nb::arg("op"));
 }
 
 NB_MODULE(_mlirDialectsLinalg, m) {

mlir/lib/CAPI/Dialect/Linalg.cpp

@@ -41,7 +41,7 @@ void mlirLinalgFillBuiltinNamedOpRegion(MlirOperation mlirOp) {
   fun(b, *body, op->getAttrs());
 }
 
-MLIR_CAPI_EXPORTED bool mlirLinalgIsContractionOp(MlirOperation op) {
+MLIR_CAPI_EXPORTED bool mlirLinalgIsAContractionOp(MlirOperation op) {
   auto linalgOp = llvm::dyn_cast<mlir::linalg::LinalgOp>(unwrap(op));
   // isaContractionOpInterface handles null linalgOp internally.
   return linalg::isaContractionOpInterface(linalgOp);
@@ -75,4 +75,49 @@ mlirLinalgInferContractionDimensions(MlirOperation op) {
   return result;
 }
 
+MLIR_CAPI_EXPORTED bool mlirLinalgIsAConvolutionOp(MlirOperation op) {
+  auto linalgOp = llvm::dyn_cast<mlir::linalg::LinalgOp>(unwrap(op));
+  if (!linalgOp)
+    return false;
+
+  return linalg::isaConvolutionOpInterface(linalgOp);
+}
+
+MLIR_CAPI_EXPORTED MlirLinalgConvolutionDimensions
+mlirLinalgInferConvolutionDimensions(MlirOperation op) {
+  MlirLinalgConvolutionDimensions result{};
+  auto linalgOp = llvm::dyn_cast<mlir::linalg::LinalgOp>(unwrap(op));
+  if (!linalgOp)
+    return result;
+
+  FailureOr<linalg::ConvolutionDimensions> maybeDims =
+      linalg::inferConvolutionDims(linalgOp);
+  if (failed(maybeDims))
+    return result;
+
+  linalg::ConvolutionDimensions dims = *maybeDims;
+  MLIRContext *ctx = linalgOp.getContext();
+
+  auto toI32Attr =
+      [&ctx](const SmallVector<unsigned, 2> &vals) -> MlirAttribute {
+    return wrap(DenseI32ArrayAttr::get(ctx, llvm::to_vector_of<int32_t>(vals)));
+  };
+
+  auto toI64Attr =
+      [&ctx](const SmallVector<int64_t, 2> &vals) -> MlirAttribute {
+    return wrap(DenseI64ArrayAttr::get(ctx, vals));
+  };
+
+  result.batch = toI32Attr(dims.batch);
+  result.outputImage = toI32Attr(dims.outputImage);
+  result.outputChannel = toI32Attr(dims.outputChannel);
+  result.filterLoop = toI32Attr(dims.filterLoop);
+  result.inputChannel = toI32Attr(dims.inputChannel);
+  result.depth = toI32Attr(dims.depth);
+  result.strides = toI64Attr(dims.strides);
+  result.dilations = toI64Attr(dims.dilations);
+
+  return result;
+}
+
 MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Linalg, linalg, LinalgDialect)