[Mlir-commits] [mlir] [mlir][linalg] Reject linalg.map payloads that read the init argument (PR #235745)

[Chanj Wons]

https://github.com/wonschangge updated https://github.com/llvm/llvm-project/pull/235745

>From 4fc886a8b119c427178eb6af5d582ec46864a7b1 Mon Sep 17 00:00:00 2001
From: Chanj Wons <892108131 at qq.com>
Date: Mon, 26 Oct 2026 18:56:08 +0800
Subject: [PATCH] [mlir][linalg] Reject linalg.map payloads that read the init
 argument

payloadUsesValueFromOperand already treats the destination of linalg.map as
unused, but MapOp::verify did not enforce that. Tighten verification so reading
the outs block argument is rejected (use linalg.generic instead), and document
the destination-passing-style invariant.

Fixes #196959
---
 .../mlir/Dialect/Linalg/IR/LinalgStructuredOps.td |  6 +++++-
 mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp          |  9 +++++++++
 mlir/test/Dialect/Linalg/invalid.mlir             | 15 +++++++++++++++
 3 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index 2b6fc04ac0991f..c6936702bf23ba 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -236,7 +236,11 @@ def MapOp : LinalgStructuredBase_Op<"map", [
   let summary = "Elementwise operations";
   let description = [{
     Models elementwise operations on tensors in terms of arithmetic operations
-    on the corresponding elements.
+    on the corresponding elements. The mapper region receives one block
+    argument per `ins` operand plus a trailing argument for the `outs`
+    destination; that destination argument must remain unused (destination
+    passing style only). Use `linalg.generic` when the payload needs to read
+    the destination value.
 
     Example:
     ```mlir
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 3c9797cc6acb5c..5947e4c72c56ee 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -1757,6 +1757,15 @@ LogicalResult MapOp::verify() {
     }
   }
 
+  // The init block argument is only present for destination-passing style.
+  // Unlike linalg.generic, the payload must not read the destination value;
+  // payloadUsesValueFromOperand already assumes this. Enforce it in verify.
+  BlockArgument initArg = blockArgs.back();
+  if (!initArg.use_empty())
+    return emitOpError()
+           << "expects the mapper region to not use the init/destination "
+              "block argument (use linalg.generic to read from outs)";
+
   return success();
 }
 
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index 70b375ef24941b..1de6d1c4e2db1a 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -1094,6 +1094,21 @@ func.func @map_input_output_shape_mismatch(
 
 // -----
 
+func.func @map_init_block_arg_used(
+    %arg0: tensor<64xf32>, %arg1: tensor<64xf32>) -> tensor<64xf32> {
+  // expected-error at +1{{'linalg.map' op expects the mapper region to not use the init/destination block argument}}
+  %mapped = linalg.map
+      ins(%arg0 : tensor<64xf32>)
+      outs(%arg1 : tensor<64xf32>)
+      (%in: f32, %init: f32) {
+        %0 = arith.addf %in, %init : f32
+        linalg.yield %0 : f32
+      }
+  func.return %mapped : tensor<64xf32>
+}
+
+// -----
+
 func.func @map_no_operands1() {
   // expected-error @+1 {{'linalg.map' op expected 1 or more operands, but found 0}}
   linalg.map { arith.addf }