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

[llvmlistbot at llvm.org]

Author: Federico Bruzzone
Date: 2026-06-16T16:30:51+02:00
New Revision: bb0eada255b31e341d8843f26757f45fdf48b301

URL: https://github.com/llvm/llvm-project/commit/bb0eada255b31e341d8843f26757f45fdf48b301
DIFF: https://github.com/llvm/llvm-project/commit/bb0eada255b31e341d8843f26757f45fdf48b301.diff

LOG: [NFC][mlir][linalg] add `toContractionDimensions` for healthy code reuse (#203916)

Follow-up to #203323.

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

---------

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>

Added: 
    

Modified: 
    mlir/lib/CAPI/Dialect/Linalg.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/CAPI/Dialect/Linalg.cpp b/mlir/lib/CAPI/Dialect/Linalg.cpp
index 64bf6aacf9905..92ead3eed9a95 100644
--- a/mlir/lib/CAPI/Dialect/Linalg.cpp
+++ b/mlir/lib/CAPI/Dialect/Linalg.cpp
@@ -48,42 +48,38 @@ 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);
   for (size_t i = 0; i < numMaps; ++i) {
     maps.push_back(unwrap(indexingMaps[i]));
   }
@@ -91,20 +87,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) {