[mlir][tosa] Only match rfft2d of floats in linalg conversion

[KB9]

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

Fixes #92064

[llvmbot]

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir-tosa

Author: Kavan Bickerstaff (KB9)

Changes

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

Fixes #92064

---

Full diff: https://github.com/llvm/llvm-project/pull/93432.diff

2 Files Affected:

• (modified) mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp (+4-1)
• (modified) mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir (+9)

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@+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>
+}

[KB9]

Ping

[sjarus]

Sorry for the delay!

[KB9]

Thanks for the review! Would you mind merging this? I don't have the permissions to do so.

[sjarus]

Done!