[Mlir-commits] [mlir] [mlir][linalg] Add maps-based inferConvolutionDims overload (PR #203323)

[Pooja Hemashekar]

================
@@ -115,41 +115,52 @@ MLIR_CAPI_EXPORTED bool mlirLinalgIsAConvolutionOp(MlirOperation op) {
   return linalg::isaConvolutionOpInterface(linalgOp);
 }
 
+static MlirLinalgConvolutionDimensions
+toConvolutionDimensions(MLIRContext *ctx,
+                        const linalg::ConvolutionDimensions &dims) {
+  auto toI32Attr = [ctx](ArrayRef<unsigned> vals) -> MlirAttribute {
+    return wrap(DenseI32ArrayAttr::get(ctx, llvm::to_vector_of<int32_t>(vals)));
+  };
+  auto toI64Attr = [ctx](ArrayRef<int64_t> vals) -> MlirAttribute {
+    return wrap(DenseI64ArrayAttr::get(ctx, vals));
+  };
+  return {toI32Attr(dims.batch),         toI32Attr(dims.outputImage),
+          toI32Attr(dims.outputChannel), toI32Attr(dims.filterLoop),
+          toI32Attr(dims.inputChannel),  toI32Attr(dims.depth),
+          toI64Attr(dims.strides),       toI64Attr(dims.dilations)};
+}
+
 MLIR_CAPI_EXPORTED MlirLinalgConvolutionDimensions
 mlirLinalgInferConvolutionDimensions(MlirOperation op) {
-  MlirLinalgConvolutionDimensions result{};
   auto linalgOp = llvm::dyn_cast<mlir::linalg::LinalgOp>(unwrap(op));
   if (!linalgOp)
-    return result;
+    return MlirLinalgConvolutionDimensions{};
 
   FailureOr<linalg::ConvolutionDimensions> maybeDims =
       linalg::inferConvolutionDims(linalgOp);
   if (failed(maybeDims))
-    return result;
+    return MlirLinalgConvolutionDimensions{};
 
-  const linalg::ConvolutionDimensions &dims = *maybeDims;
-  MLIRContext *ctx = linalgOp.getContext();
+  return toConvolutionDimensions(linalgOp.getContext(), *maybeDims);
+}
 
-  auto toI32Attr =
-      [&ctx](const SmallVector<unsigned, 2> &vals) -> MlirAttribute {
-    return wrap(DenseI32ArrayAttr::get(ctx, llvm::to_vector_of<int32_t>(vals)));
-  };
+MLIR_CAPI_EXPORTED MlirLinalgConvolutionDimensions
+mlirLinalgInferConvolutionDimensionsFromMaps(const MlirAffineMap *indexingMaps,
+                                             size_t numMaps) {
+  if (!indexingMaps || numMaps == 0)
+    return MlirLinalgConvolutionDimensions{};
 
-  auto toI64Attr =
-      [&ctx](const SmallVector<int64_t, 2> &vals) -> MlirAttribute {
-    return wrap(DenseI64ArrayAttr::get(ctx, vals));
-  };
+  SmallVector<AffineMap, 3> maps;
+  maps.reserve(numMaps);
----------------
phemashekar wrote:

you;re right, reserve is redundant here. On checking `numMaps == 3`, I'd prefer to leave that to `inferConvolutionDims`, which already enforces it. That keeps the CAPI a thin wrapper with a single source of truth for validation, just in case if things were to differ in future.

https://github.com/llvm/llvm-project/pull/203323