[mlir] tosa.concat - fix isCompatibleReturnTypes()

Reviewed By: eric-k256

Differential Revision: https://reviews.llvm.org/D146901

Author: amrami

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

@@ -425,6 +425,8 @@ LogicalResult tosa::FFT2dOp::inferReturnTypeComponents(
 bool tosa::ConcatOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
   if (l.size() != r.size() || l.size() != 1)
     return false;
+  if (getElementTypeOrSelf(l[0]) != getElementTypeOrSelf(r[0]))
+    return false;
   return succeeded(verifyCompatibleShape(l[0], r[0]));
 }
 

mlir/test/Dialect/Tosa/invalid.mlir

@@ -46,6 +46,14 @@ func.func @test_concat(%arg0 : tensor<2x1xf32>, %arg1 : tensor<2x2xf32>) -> tens
 
 // -----
 
+func.func @test_concat_element_type_mismatch(%arg0 : tensor<1x2xf32>, %arg1 : tensor<2x2xf32>) -> tensor<?x?xi8> {
+  // expected-error@+1 {{'tosa.concat' op inferred type(s) 'tensor<3x2xf32>' are incompatible with return type(s) of operation 'tensor<?x?xi8>}}
+  %0 = "tosa.concat"(%arg0, %arg1) {axis = 0 : i64} : (tensor<1x2xf32>, tensor<2x2xf32>) -> tensor<?x?xi8>
+  return %0 : tensor<?x?xi8>
+}
+
+// -----
+
 func.func @test_pad_non_const(%arg0: tensor<13x21x3xf32>, %arg1: tensor<3x2xi32>) -> tensor<13x21x3xf32> {
   // expected-error@+1 {{'tosa.pad' op padding of pad is not constant}}
   %0 = "tosa.pad"(%arg0, %arg1) : (tensor<13x21x3xf32>, tensor<3x2xi32>) -> tensor<13x21x3xf32>