[Mlir-commits] [mlir] [mlir][linalg] add `toContractionDimensions` for healthy code reuse (PR #203916)

[Federico Bruzzone]

https://github.com/FedericoBruzzone created https://github.com/llvm/llvm-project/pull/203916

Follow-up to #203323.

This PR adds the `toContractionDimensions` static function, aligning with what we suggested in #203323.

>From e7c4f7ce7950a3003ea1196a9483ee397391a112 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Mon, 15 Jun 2026 16:27:53 +0200
Subject: [PATCH] [mlir][linalg] add `toContractionDimensions` for healthy code
 reuse

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 mlir/lib/CAPI/Dialect/Linalg.cpp | 44 ++++++++++++--------------------
 1 file changed, 16 insertions(+), 28 deletions(-)

diff --git a/mlir/lib/CAPI/Dialect/Linalg.cpp b/mlir/lib/CAPI/Dialect/Linalg.cpp
index 64bf6aacf9905..01b7a87b1b284 100644
--- a/mlir/lib/CAPI/Dialect/Linalg.cpp
+++ b/mlir/lib/CAPI/Dialect/Linalg.cpp
@@ -48,39 +48,36 @@ MLIR_CAPI_EXPORTED bool mlirLinalgIsAContractionOp(MlirOperation op) {
   return linalg::isaContractionOpInterface(linalgOp);
 }
 
+static MlirLinalgContractionDimensions
+toContractionDimensions(MLIRContext *ctx,
+                        const linalg::ContractionDimensions &dims) {
+  auto toAttr = [ctx](ArrayRef<unsigned> vals) -> MlirAttribute {
+    return wrap(DenseI32ArrayAttr::get(ctx, llvm::to_vector_of<int32_t>(vals)));
+  };
+  return {toAttr(dims.batch), toAttr(dims.m), toAttr(dims.n), toAttr(dims.k)};
+}
+
 MLIR_CAPI_EXPORTED MlirLinalgContractionDimensions
 mlirLinalgInferContractionDimensions(MlirOperation op) {
-  MlirLinalgContractionDimensions result{};
   auto linalgOp = dyn_cast<linalg::LinalgOp>(unwrap(op));
   if (!linalgOp)
-    return result;
+    return {};
 
   FailureOr<linalg::ContractionDimensions> maybeDims =
       linalg::inferContractionDims(linalgOp);
   if (failed(maybeDims))
-    return result;
+    return {};
 
   const linalg::ContractionDimensions &contractionDims = *maybeDims;
   MLIRContext *ctx = linalgOp.getContext();
-
-  auto toAttr = [ctx](ArrayRef<unsigned> vals) -> MlirAttribute {
-    return wrap(DenseI32ArrayAttr::get(ctx, llvm::to_vector_of<int32_t>(vals)));
-  };
-
-  result.batch = toAttr(contractionDims.batch);
-  result.m = toAttr(contractionDims.m);
-  result.n = toAttr(contractionDims.n);
-  result.k = toAttr(contractionDims.k);
-
-  return result;
+  return toContractionDimensions(ctx, contractionDims);
 }
 
 MLIR_CAPI_EXPORTED MlirLinalgContractionDimensions
 mlirLinalgInferContractionDimensionsFromMaps(const MlirAffineMap *indexingMaps,
                                              size_t numMaps) {
-  MlirLinalgContractionDimensions result{};
-  if (!indexingMaps || numMaps == 0)
-    return result;
+  if (!indexingMaps || numMaps != 3)
+    return {};
 
   SmallVector<AffineMap, 3> maps;
   maps.reserve(numMaps);
@@ -91,20 +88,11 @@ mlirLinalgInferContractionDimensionsFromMaps(const MlirAffineMap *indexingMaps,
   FailureOr<linalg::ContractionDimensions> maybeDims =
       linalg::inferContractionDims(maps);
   if (failed(maybeDims))
-    return result;
+    return {};
 
   MLIRContext *ctx = maps[0].getContext();
 
-  auto toAttr = [ctx](ArrayRef<unsigned> vals) -> MlirAttribute {
-    return wrap(DenseI32ArrayAttr::get(ctx, llvm::to_vector_of<int32_t>(vals)));
-  };
-
-  result.batch = toAttr(maybeDims->batch);
-  result.m = toAttr(maybeDims->m);
-  result.n = toAttr(maybeDims->n);
-  result.k = toAttr(maybeDims->k);
-
-  return result;
+  return toContractionDimensions(ctx, *maybeDims);
 }
 
 MLIR_CAPI_EXPORTED bool mlirLinalgIsAConvolutionOp(MlirOperation op) {