[mlir] Draft: Fix shape inference of tosa elementwise binary ops

[mamrami]

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

992

this function does not cover all changes needed. It is an expanded version of isCompatibleReturnTypes used in this file, with this additional check.

[mgehre-amd]

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

992

Can you explain the rational for the checks in this function? What is the goal here?

mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td

 class Tosa_ElemWiseUnaryOp<string mnemonic, list<Trait> traits = []> :
     Tosa_Op<mnemonic, !listconcat(traits, [
               DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                                         ["inferReturnTypeComponents"]>,
               Pure, SameOperandsAndResultElementType])> {
 }
 
 class Tosa_ElemWiseBinaryOp<string mnemonic, list<Trait> traits = []> :
-    Tosa_Op<mnemonic, !listconcat(traits, [
-              DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                                        ["inferReturnTypeComponents"]>,
+    Tosa_Op<mnemonic, !listconcat(traits, [InferTensorType,
               ResultsBroadcastableShape, Pure, SameOperandsAndResultElementType])> {
+  let extraClassDeclaration = [{
+    /// Returns true when two result types are compatible for this op;
+    /// Method used by InferTypeOpInterface.
+    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+  }];
 }
 
 #endif // TOSA_OP_BASE

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

static LogicalResult ReduceInferReturnTypes(
   SmallVector<int64_t> outputShape;
   operandShape.getDims(outputShape);
   int64_t axisVal = axis.getValue().getSExtValue();
   outputShape[axisVal] = 1;
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape, inputType));
   return success();
 }
 
+static bool binaryIsCompatibleReturnTypes(TypeRange inferred,
+                                          TypeRange returned) {
+  if (inferred.size() != returned.size() || inferred.size() != 1)
+    return false;
+  // return succeeded(verifyCompatibleShape(l[0], r[0]));
+  auto inferredType = llvm::dyn_cast<ShapedType>(inferred[0]);
+  auto returnedType = llvm::dyn_cast<ShapedType>(returned[0]);
+
+  // Either both or neither type should be shaped.
+  if (!inferredType)
+    return !inferredType;
+  if (!inferredType)
+    return false;
+
+  if (!inferredType.hasRank() || !inferredType.hasRank())
+    return true;
+
+  ArrayRef<int64_t> shape1 = inferredType.getShape();
+  ArrayRef<int64_t> shape2 = inferredType.getShape();
+  if (shape1.size() != shape2.size())
+    return false;
+  if (returnedType.hasStaticShape() && !inferredType.hasStaticShape())
+    return false;
+  for (auto dims : llvm::zip(shape1, shape2)) {
+    int64_t dim1 = std::get<0>(dims);
+    int64_t dim2 = std::get<1>(dims);
+    if (!ShapedType::isDynamic(dim1) && !ShapedType::isDynamic(dim2) &&
+        dim1 != dim2)
+      return false;
+  }
+  return true;
+}
+
 #define COMPATIBLE_RETURN_TYPES(OP)                                            \
   bool OP::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]));                       \
   }

 #undef REDUCE_SHAPE_INFER
 COMPATIBLE_RETURN_TYPES(tosa::ConcatOp)
 #undef COMPATIBLE_RETURN_TYPES
 
 static LogicalResult NAryInferReturnTypes(
     const ValueShapeRange &operands,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<int64_t> outShape;
+  Type inputType = operands.getType()[0].cast<TensorType>().getElementType();
   if (resolveBroadcastShape(operands, outShape).failed()) {
-    inferredReturnShapes.push_back(ShapedTypeComponents());
+    inferredReturnShapes.push_back(ShapedTypeComponents(inputType));
   } else {
-    inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
+    inferredReturnShapes.push_back(ShapedTypeComponents(outShape, inputType));
   }
   return success();
 }
 
 #define NARY_SHAPE_INFER(OP)                                                   \
   LogicalResult OP::inferReturnTypeComponents(                                 \
       MLIRContext *context, ::std::optional<Location> location,                \
       ValueShapeRange operands, DictionaryAttr attributes,                     \
       OpaqueProperties properties, RegionRange regions,                        \
       SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {           \
     return NAryInferReturnTypes(operands, inferredReturnShapes);               \
   }
 
+#define BINARY_COMPATIBLE_TYPES(OP)                                            \
+  bool OP::isCompatibleReturnTypes(TypeRange l, TypeRange r) {                 \
+    return binaryIsCompatibleReturnTypes(l, r);                                \
+  }
+
+BINARY_COMPATIBLE_TYPES(tosa::AddOp)
+BINARY_COMPATIBLE_TYPES(tosa::ArithmeticRightShiftOp)
+BINARY_COMPATIBLE_TYPES(tosa::BitwiseAndOp)
+BINARY_COMPATIBLE_TYPES(tosa::BitwiseOrOp)
+BINARY_COMPATIBLE_TYPES(tosa::BitwiseXorOp)
+BINARY_COMPATIBLE_TYPES(tosa::DivOp)
+BINARY_COMPATIBLE_TYPES(tosa::LogicalAndOp)
+BINARY_COMPATIBLE_TYPES(tosa::LogicalLeftShiftOp)
+BINARY_COMPATIBLE_TYPES(tosa::LogicalOrOp)
+BINARY_COMPATIBLE_TYPES(tosa::LogicalRightShiftOp)
+BINARY_COMPATIBLE_TYPES(tosa::LogicalXorOp)
+BINARY_COMPATIBLE_TYPES(tosa::MaximumOp)
+BINARY_COMPATIBLE_TYPES(tosa::MinimumOp)
+BINARY_COMPATIBLE_TYPES(tosa::MulOp)
+BINARY_COMPATIBLE_TYPES(tosa::PowOp)
+BINARY_COMPATIBLE_TYPES(tosa::SubOp)
+#undef BINARY_COMPATIBLE_TYPES
 NARY_SHAPE_INFER(tosa::AbsOp)
 NARY_SHAPE_INFER(tosa::AddOp)
 NARY_SHAPE_INFER(tosa::ArithmeticRightShiftOp)
 NARY_SHAPE_INFER(tosa::BitwiseAndOp)
 NARY_SHAPE_INFER(tosa::BitwiseOrOp)
 NARY_SHAPE_INFER(tosa::BitwiseXorOp)
 NARY_SHAPE_INFER(tosa::BitwiseNotOp)
 NARY_SHAPE_INFER(tosa::CastOp)

 NARY_SHAPE_INFER(tosa::RescaleOp)
 NARY_SHAPE_INFER(tosa::ReverseOp)
 NARY_SHAPE_INFER(tosa::RsqrtOp)
 NARY_SHAPE_INFER(tosa::SelectOp)
 NARY_SHAPE_INFER(tosa::SubOp)
 NARY_SHAPE_INFER(tosa::TanhOp)
 NARY_SHAPE_INFER(tosa::ErfOp)
 NARY_SHAPE_INFER(tosa::SigmoidOp)
-#undef PRED_SHAPE_INFER
+#undef NARY_SHAPE_INFER
 
 static LogicalResult poolingInferReturnTypes(
     const ValueShapeRange &operands, DictionaryAttr attributes,
     ArrayRef<int64_t> kernel, ArrayRef<int64_t> stride, ArrayRef<int64_t> pad,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   ShapeAdaptor inputShape = operands.getShape(0);
   llvm::SmallVector<int64_t> outputShape;
   outputShape.resize(4, ShapedType::kDynamic);