[Mlir-commits] [mlir] [mlir][linalg] Fix crashes in parser on linalg ops without operands (PR #97944)

[Felix Schneider]

https://github.com/ubfx created https://github.com/llvm/llvm-project/pull/97944

`parseDstStyleOp` parses both `ins()` and `outs()` optionally. The parsers for `linalg.transpose`, `linalg.broadcast` and `linalg.map` however assume that at least one operand is present in the state, leading to crashes otherwise.

This patch adds checks to the parsers which stop them from crashing if no operands were parsed. After the Ops are parsed successfuly, the verifier takes it from there.

Fix https://github.com/llvm/llvm-project/issues/97857

>From bd56160def5f6ac4c2f25a63eb4bcee54f2630ef Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Sun, 7 Jul 2024 11:59:22 +0200
Subject: [PATCH] [mlir][linalg] Fix crashes in parser on linalg ops without
 operands

`parseDstStyleOp` parses both `ins()` and `outs()` optionally. The parsers
for `linalg.transpose`, `linalg.broadcast` and `linalg.map` however
assume that at least one operand is present in the state, leading to crashes
otherwise.

This patch adds checks to the parsers which stop them from crashing if
no operands were parsed. After the Ops are parsed successfuly, the verifier
takes it from there.

Fix https://github.com/llvm/llvm-project/issues/97857
---
 mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp |  9 ++++---
 mlir/test/Dialect/Linalg/invalid.mlir    | 31 ++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 57d126603ebd7..dec84ed579096 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -1356,8 +1356,10 @@ ParseResult MapOp::parse(OpAsmParser &parser, OperationState &result) {
     return failure();
 
   if (payloadOpName.has_value()) {
-    addBodyWithPayloadOp(parser, result, payloadOpName.value(), payloadOpAttrs,
-                         ArrayRef(result.operands).drop_back());
+    if (!result.operands.empty())
+      addBodyWithPayloadOp(parser, result, payloadOpName.value(),
+                           payloadOpAttrs,
+                           ArrayRef(result.operands).drop_back());
   } else {
     SmallVector<OpAsmParser::Argument> regionArgs;
     if (parser.parseArgumentList(regionArgs, OpAsmParser::Delimiter::Paren,
@@ -1739,7 +1741,8 @@ static void buildIdentityRegion(OpBuilder &builder, Location loc,
                                 ValueRange outputs) {
   buildGenericRegion(builder, loc, region, inputs, outputs,
                      [](OpBuilder &b, Location loc, ValueRange args) {
-                       b.create<linalg::YieldOp>(loc, args[0]);
+                       if (!args.empty())
+                         b.create<linalg::YieldOp>(loc, args[0]);
                      });
 }
 
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index 44c81c31ace0f..cfd269e3e6e3a 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -455,6 +455,18 @@ func.func @map_input_output_shape_mismatch(
 
 // -----
 
+func.func @map_no_operands(
+    %lhs: tensor<64xf32>, %rhs: tensor<64xf32>, %init: tensor<64xf32>)
+    -> tensor<64xf32> {
+  // This must not crash the parser.
+  linalg.map { arith.addf }
+  // expected-error @+1 {{cannot name an operation with no results}}
+  %add = linalg.map { arith.addf }
+  func.return %add : tensor<64xf32>
+}
+
+// -----
+
 func.func @reduce_input_vs_init_dimension_mismatch(
     %input: tensor<16x32x64xf32>,
     %init: tensor<16x64xf32>)  -> tensor<16x64xf32> {
@@ -676,6 +688,16 @@ func.func @transpose_input_init_rank_mismatch(%input: tensor<16x32xf32>,
 
 // -----
 
+func.func @transpose_no_operands() -> tensor<32x64x16xf32> {
+  // This must not crash the parser.
+  linalg.transpose permutation = [1, 0, 2]
+  // expected-error @+1 {{cannot name an operation with no results}}
+  %transpose = linalg.transpose permutation = [1, 0, 2]
+  func.return %transpose : tensor<32x64x16xf32>
+}
+
+// -----
+
 func.func @broadcast_input_dims_rank_mismatch(
     %input: tensor<4x16xf32>, %init: tensor<4x8x16xf32>)
     -> tensor<4x8x16xf32> {
@@ -725,6 +747,15 @@ func.func @broadcast_size_1_extension_not_supported(
       dimensions = [1]
   func.return %bcast : tensor<4x?x16xf32>
 }
+// -----
+
+func.func @broadcast_no_operands()
+    -> tensor<4x?x16xf32> {
+  linalg.broadcast dimensions = [1]
+  // expected-error @+1 {{cannot name an operation with no results}}
+  %broadcast = linalg.broadcast dimensions = [1]
+  func.return %broadcast : tensor<32x64x16xf32>
+}
 
 // -----