[Mlir-commits] [mlir] [mlir][linalg] Restrict linalg.contract results (PR #205988)

[llvmlistbot at llvm.org]

https://github.com/mygitljf updated https://github.com/llvm/llvm-project/pull/205988

>From 21dd188c0d6b094375fb038f288c10947bdac949 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Fri, 26 Jun 2026 14:24:31 +0000
Subject: [PATCH 1/2] Restrict linalg.contract results

---
 .../mlir/Dialect/Linalg/IR/LinalgStructuredOps.td |  2 +-
 mlir/test/Dialect/Linalg/invalid.mlir             | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index 5998f736ced34..fc5c9770d969b 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -881,7 +881,7 @@ def ContractOp : LinalgStructuredBase_Op<"contract", [
     AffineMapArrayAttr:$indexing_maps,
     DefaultValuedOptionalAttr<TypeFnAttr, "TypeFn::cast_signed">:$cast
   );
-  let results = (outs Variadic<AnyShaped>:$result_tensors);
+  let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
   // NB: The only reason this op has a region - and it get populated at op build
   //     time - is that currently the LinalgOp interface exposes methods that
   //     assume a relevant region is available to be queried at any time.
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index a446cfcc4eec1..91b2a495f27bb 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -422,6 +422,21 @@ func.func @illegal_fill_tensor_with_memref_return
 
 // -----
 
+func.func @illegal_contract_memref_with_memref_return
+  (%arg0: memref<4x8xf32>, %arg1: memref<8x6xf32>, %arg2: memref<4x6xf32>) -> memref<4x6xf32>
+{
+  // expected-error @+1 {{result #0 must be variadic of ranked tensor of any non-token type values, but got 'memref<4x6xf32>'}}
+  %0 = linalg.contract
+      indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
+                       affine_map<(d0, d1, d2) -> (d2, d1)>,
+                       affine_map<(d0, d1, d2) -> (d0, d1)>]
+      ins(%arg0, %arg1 : memref<4x8xf32>, memref<8x6xf32>)
+      outs(%arg2 : memref<4x6xf32>) -> memref<4x6xf32>
+  return %0 : memref<4x6xf32>
+}
+
+// -----
+
 func.func @illegal_fill_element_type_truncation(%arg0 : tensor<2xf32>, %arg1 : f64) -> tensor<2xf32>
 {
   // expected-error @+1 {{'linalg.fill' op expected fill value type ('f64') to match output element type ('f32')}}

>From e00650040d409266c7acc532f2dd67c50a12b1fb Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Sat, 27 Jun 2026 10:41:21 +0000
Subject: [PATCH 2/2] Move contract invalid test

---
 mlir/test/Dialect/Linalg/invalid.mlir | 30 +++++++++++++--------------
 1 file changed, 15 insertions(+), 15 deletions(-)

diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index 91b2a495f27bb..c52163e244ea2 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -422,21 +422,6 @@ func.func @illegal_fill_tensor_with_memref_return
 
 // -----
 
-func.func @illegal_contract_memref_with_memref_return
-  (%arg0: memref<4x8xf32>, %arg1: memref<8x6xf32>, %arg2: memref<4x6xf32>) -> memref<4x6xf32>
-{
-  // expected-error @+1 {{result #0 must be variadic of ranked tensor of any non-token type values, but got 'memref<4x6xf32>'}}
-  %0 = linalg.contract
-      indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
-                       affine_map<(d0, d1, d2) -> (d2, d1)>,
-                       affine_map<(d0, d1, d2) -> (d0, d1)>]
-      ins(%arg0, %arg1 : memref<4x8xf32>, memref<8x6xf32>)
-      outs(%arg2 : memref<4x6xf32>) -> memref<4x6xf32>
-  return %0 : memref<4x6xf32>
-}
-
-// -----
-
 func.func @illegal_fill_element_type_truncation(%arg0 : tensor<2xf32>, %arg1 : f64) -> tensor<2xf32>
 {
   // expected-error @+1 {{'linalg.fill' op expected fill value type ('f64') to match output element type ('f32')}}
@@ -659,6 +644,21 @@ func.func @invalid_type_matmul(%arg0 : !x86.amx.tile<16x16xbf16>)
 
 // -----
 
+func.func @illegal_contract_memref_with_memref_return
+  (%arg0: memref<4x8xf32>, %arg1: memref<8x6xf32>, %arg2: memref<4x6xf32>) -> memref<4x6xf32>
+{
+  // expected-error @+1 {{result #0 must be variadic of ranked tensor of any non-token type values, but got 'memref<4x6xf32>'}}
+  %0 = linalg.contract
+      indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
+                       affine_map<(d0, d1, d2) -> (d2, d1)>,
+                       affine_map<(d0, d1, d2) -> (d0, d1)>]
+      ins(%arg0, %arg1 : memref<4x8xf32>, memref<8x6xf32>)
+      outs(%arg2 : memref<4x6xf32>) -> memref<4x6xf32>
+  return %0 : memref<4x6xf32>
+}
+
+// -----
+
 func.func @invalid_indexing_maps_placement_contraction(
     %lhs: tensor<4x1xf32>, %rhs: tensor<1x64xf32>, %init: tensor<4x64xf32>) {
   // expected-error @+3 {{custom op 'linalg.contract' expected 'indexing_maps' attribute}}