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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Federico Bruzzone (FedericoBruzzone)

<details>
<summary>Changes</summary>

Follow-up to #<!-- -->203323.

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

---
Full diff: https://github.com/llvm/llvm-project/pull/203916.diff


1 Files Affected:

- (modified) mlir/lib/CAPI/Dialect/Linalg.cpp (+16-28) 


``````````diff
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) {

``````````

</details>


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