[Mlir-commits] [mlir] [mlir][TosaToLinalg] Exit after `notifyMatchFailure` (PR #132012)

[Longsheng Mou]

https://github.com/CoTinker created https://github.com/llvm/llvm-project/pull/132012

This PR adds `return nullptr` when the shift value of `tosa.mul` is not constant to prevent a crash. Fixes #131766.

>From 2b2357cceda6a10bf706172353c3c3a41f77e0f3 Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Wed, 19 Mar 2025 20:02:07 +0800
Subject: [PATCH] [mlir][TosaToLinalg] Exit after `notifyMatchFailure`

This PR adds `return nullptr` when the shift value of
`tosa.mul` is not constant to prevent a crash.
---
 mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp      | 10 +++++-----
 .../TosaToLinalg/tosa-to-linalg-invalid.mlir           |  8 ++++++++
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index c0a25a56dbe2a..6e1e3343ac169 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -136,14 +136,14 @@ static Value createLinalgBodyCalculationForElementwiseOp(
 
   // tosa::MulOp
   if (isa<tosa::MulOp>(op)) {
-    auto shift_val = cast<tosa::MulOp>(op).getShift();
-    DenseElementsAttr shift_elem;
-    if (!shift_val.getImpl() ||
-        !matchPattern(shift_val, m_Constant(&shift_elem))) {
+    auto shiftVal = cast<tosa::MulOp>(op).getShift();
+    DenseElementsAttr shiftElem;
+    if (!matchPattern(shiftVal, m_Constant(&shiftElem))) {
       (void)rewriter.notifyMatchFailure(op, "shift value of mul not found");
+      return nullptr;
     }
 
-    int32_t shift = shift_elem.getValues<IntegerAttr>()[0].getInt();
+    int32_t shift = shiftElem.getValues<IntegerAttr>()[0].getInt();
 
     if (isa<FloatType>(elementTy)) {
       if (shift != 0) {
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
index d00846a4c3e02..69d8471df8032 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
@@ -73,3 +73,11 @@ func.func @unranked_gather(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>)
   %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x3xf32>, tensor<13x26xi32>) -> tensor<*xf32>
   return %0 : tensor<*xf32>
 }
+
+// -----
+
+func.func @mul_no_const_shift(%arg0: tensor<2x3xi32>, %arg1: tensor<2x3xi32>, %arg2: tensor<1xi8>) -> tensor<2x3xi32> {
+  // expected-error at +1 {{failed to legalize operation 'tosa.mul'}}
+  %0 = tosa.mul %arg0, %arg1, %arg2 : (tensor<2x3xi32>, tensor<2x3xi32>, tensor<1xi8>) -> tensor<2x3xi32>
+  return %0 : tensor<2x3xi32>
+}