[mlir][tosa] Fix crash when inferring shape of tosa.equal

[sabauma]

Summary

The tosa-infer-shapes pass crashes when trying to infer the output shape
of tosa.equal when the input shape is unranked.
This is due to tosa-infer-shapes requiring at least information on the
base type of the resulting operation from inferReturnTypeComponents.
This change enhances EqualOp::inferReturnTypeComponents to always supply
the inferred elementType.

[eric-k256]

LGTM.

[sabauma]

If there are no other comments, could someone with submission privileges please submit the patch?

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

LogicalResult tosa::ConcatOp::inferReturnTypeComponents(
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape, inputType));
   return success();
 }
 
 LogicalResult tosa::EqualOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+  auto elementType = IntegerType::get(context, /*width=*/1);
+
   llvm::SmallVector<int64_t> outShape;
   if (resolveBroadcastShape(operands, outShape).failed()) {
-    inferredReturnShapes.push_back(ShapedTypeComponents());
+    inferredReturnShapes.push_back(ShapedTypeComponents(elementType));
     return success();
   }
 
-  inferredReturnShapes.push_back(
-      ShapedTypeComponents(outShape, IntegerType::get(context, /*width=*/1)));
+  inferredReturnShapes.push_back(ShapedTypeComponents(outShape, elementType));
   return success();
 }
 
 bool tosa::EqualOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
   if (l.size() != r.size() || l.size() != 1)
     return false;
   return succeeded(verifyCompatibleShape(l[0], r[0]));
 }

mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir

 // -----
 
 // CHECK-LABEL: @test_dynamic_batch_fft2d
 func.func @test_dynamic_batch_fft2d(%arg0: tensor<?x4x8xf32>, %arg1: tensor<?x4x8xf32>) -> (tensor<?x4x8xf32>, tensor<?x4x8xf32>) {
   // CHECK: -> (tensor<?x4x8xf32>, tensor<?x4x8xf32>)
   %output_real, %output_imag = "tosa.fft2d"(%arg0, %arg1) {inverse = false} : (tensor<?x4x8xf32>, tensor<?x4x8xf32>) -> (tensor<?x4x8xf32>, tensor<?x4x8xf32>)
   return %output_real, %output_imag : tensor<?x4x8xf32>, tensor<?x4x8xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @test_unranked_equal
+func.func @test_unranked_equal(%arg0 : tensor<*xf32>, %arg1 : tensor<f32>) -> () {
+  // CHECK: "tosa.equal"(%arg0, %arg1) : (tensor<*xf32>, tensor<f32>) -> tensor<*xi1>
+  %0 = "tosa.equal"(%arg0, %arg1) : (tensor<*xf32>, tensor<f32>) -> tensor<*xi1>
+
+  return
+}