[Mlir-commits] [mlir] [MLIR][Linalg] Improve morphism semantics (PR #222959)

[Federico Bruzzone]

================
@@ -43,11 +43,54 @@ static LogicalResult generalizeNamedOpPrecondition(LinalgOp linalgOp) {
   return success();
 }
 
-FailureOr<GenericOp> mlir::linalg::generalizeNamedOp(RewriterBase &rewriter,
-                                                     LinalgOp linalgOp) {
+// Converts a named matmul-like op (`matmul`, `batch_matmul`, or
+// `batch_reduce_matmul`) into a `linalg.contract` category op, preserving the
+// operand indexing maps and cast semantics. Returns failure for other ops.
+static FailureOr<LinalgOp> generalizeToContractOp(RewriterBase &rewriter,
+                                                  LinalgOp namedOp) {
+  // These are the ODS-defined matmul-like operations.
+  // For OpDSL declared contractions, please move them to ODS first,
+  // then add them to the isa<> check below + tests.
+  if (!isa<MatmulOp, BatchMatmulOp, BatchReduceMatmulOp>(
+          namedOp.getOperation()))
+    return failure();
+
+  SmallVector<NamedAttribute> attributes;
+
+  // Preserve operand indexing semantics (transposition, batch/reduction dims)
+  // via the named op's indexing maps.
+  SmallVector<Attribute> indexingMaps = llvm::map_to_vector(
+      namedOp.getIndexingMapsArray(),
+      [](AffineMap map) -> Attribute { return AffineMapAttr::get(map); });
+  attributes.push_back(rewriter.getNamedAttr(
+      "indexing_maps", rewriter.getArrayAttr(indexingMaps)));
+
+  // Only the unsigned cast needs to be explicit; signed is the default.
+  if (auto castAttr = namedOp->getAttrOfType<TypeFnAttr>("cast");
+      castAttr && castAttr.getValue() == TypeFn::cast_unsigned)
+    attributes.push_back(rewriter.getNamedAttr("cast", castAttr));
+
+  LinalgOp contractOp = rewriter.replaceOpWithNewOp<ContractOp>(
----------------
FedericoBruzzone wrote:

Should the `contract` path keep discardable attrs like the generic path does?

https://github.com/llvm/llvm-project/blob/e1a4305b37f20140819a1c768cd873f93add61b5/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp#L111-L116

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