[Mlir-commits] [mlir] [mlir][tosa] Only match rfft2d of floats in linalg conversion (PR #93432)

[Kavan Bickerstaff]

https://github.com/KB9 created https://github.com/llvm/llvm-project/pull/93432

This prevents an assertion being triggered by the cast to FloatType.

>From cc853d6063d7b6f7fffe93ed2b25774799db991a Mon Sep 17 00:00:00 2001
From: KB9 <kavanbickerstaff at googlemail.com>
Date: Mon, 27 May 2024 01:00:38 +0100
Subject: [PATCH] [mlir][tosa] Only match rfft2d of floats in linalg conversion

This prevents an assertion being triggered by the cast to FloatType.
---
 mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp        | 5 ++++-
 .../Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir  | 9 +++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index e6ba6e6bc602d..8ad8e41414656 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -2324,7 +2324,10 @@ struct RFFT2dConverter final : public OpRewritePattern<RFFT2dOp> {
     auto loc = rfft2d.getLoc();
     auto input = rfft2d.getInput();
     auto elementType =
-        cast<FloatType>(cast<ShapedType>(input.getType()).getElementType());
+        dyn_cast<FloatType>(cast<ShapedType>(input.getType()).getElementType());
+    if (!elementType)
+      return rewriter.notifyMatchFailure(rfft2d,
+                                         "only supports float element types");
 
     // Compute the output type and set of dynamic sizes
     llvm::SmallVector<Value> dynamicSizes;
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
index ad65410e635e9..b78577275a52a 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
@@ -27,3 +27,12 @@ func.func @unranked_add(%arg0 : tensor<10x10xf32> , %arg1 : tensor<10x10xf32>, %
   %2 = tosa.reshape %0 {new_shape = array<i64: 10, 10>} : (tensor<*xf32>) -> tensor<10x10xf32>
   return %2 : tensor<10x10xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @rfft2d_with_non_float_type
+func.func @rfft2d_with_non_float_type(%arg0 : tensor<1x1x1xi32>) -> (tensor<1x1x1xi32>, tensor<1x1x1xi32>) {
+  // expected-error at +1 {{failed to legalize operation 'tosa.rfft2d'}}
+  %real, %imag = tosa.rfft2d %arg0 : (tensor<1x1x1xi32>) -> (tensor<1x1x1xi32>, tensor<1x1x1xi32>)
+  return %real, %imag : tensor<1x1x1xi32>, tensor<1x1x1xi32>
+}