[mlir][tosa] Add InferTensorType interface to tosa reduce operations

[AviadCo]

Summary

When this interface is used, a call to inferReturnTypeComponents()
is generated on creation and verification of the op.

[AviadCo]

Addeed mlir tests for tosa reduce operations infer shape.

[mamrami]

mlir/test/Dialect/Tosa/invalid.mlir

102 ↗

(On Diff #510396)

I would remove the function name check and the newlines, to align to the other tests.

[mamrami]

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

944

Can we have a static method with this implementation? Then it can be called here and in tosa::ConcatOp::isCompatibleReturnTypes to avoid code duplication.
Or maybe we can stay with this implementation only, and have another macro that will be defined for tosa::ConcatOp as well as for the reduce ops.

[AviadCo]

Fixed Maya's comments.

[AviadCo]

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

944

Good point, made new macro and included tosa::ConcatOp in it.

mlir/test/Dialect/Tosa/invalid.mlir

102 ↗

(On Diff #510396)

Ack.

[AviadCo]

Removed space and redundent code from mlir checks.

[eric-k256]

Looks good to me.

[jpienaar]

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

933

So rhs size being 1 would be an error but if lhs were 1 then it's fine? (For example, if r.size() was 0 then the below would assert)

944

I think the proposal to call the static method from the macro sounds good. Just reduces work for compiler.

[AviadCo]

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

933

The tosa operations that have this macro are not changing the tensor length/element type, and all of them working on tensors with length 1.

Actually lhs and rhs must be from the same length from the first condition (otherwise we will return false). From the second condition lhs mustn't be 1. From the combination of those conditions we can deduce that we will return false also if rhs is 1 also (either lhs is 1 also nor rhs != lhs).

944

Ack. thanks.

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

 // TOSA Spec Section 2.8
 // Operator Class: Reduction Ops.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: reduce_all
 //===----------------------------------------------------------------------===//
 def Tosa_ReduceAllOp : Tosa_Op<"reduce_all", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+    InferTensorType, Pure]> {
   let summary = "Reduce All operator";
 
   let description = [{
     Reduce a tensor along the given axis with a logical AND operation
   }];
 
   let arguments = (ins
     Tosa_Tensor1Dto4D:$input,
     I64Attr:$axis
   );
 
   let results = (outs
     Tosa_Tensor1Dto4D:$output
   );
 
   let hasFolder = 1;
+
+  let extraClassDeclaration = [{
+    /// Returns true when two result types are compatible for this op;
+    /// Method used by InferTypeOpInterface.
+    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+  }];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reduce_any
 //===----------------------------------------------------------------------===//
 def Tosa_ReduceAnyOp : Tosa_Op<"reduce_any", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+    InferTensorType, Pure]> {
   let summary = "Reduce Any operator";
 
   let description = [{
     Reduce a tensor along the given axis with a logical OR operation
   }];
 
   let arguments = (ins
     Tosa_Tensor1Dto4D:$input,
     I64Attr:$axis
   );
 
   let results = (outs
     Tosa_Tensor1Dto4D:$output
   );
 
   let hasFolder = 1;
+
+  let extraClassDeclaration = [{
+    /// Returns true when two result types are compatible for this op;
+    /// Method used by InferTypeOpInterface.
+    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+  }];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reduce_max
 //===----------------------------------------------------------------------===//
 def Tosa_ReduceMaxOp : Tosa_Op<"reduce_max", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+    InferTensorType, Pure]> {
   let summary = "Reduce Max operator";
 
   let description = [{
     Reduce a tensor along the given axis with a maximum operation
   }];
 
   let arguments = (ins
     Tosa_Tensor1Dto4D:$input,
     I64Attr:$axis
   );
 
   let results = (outs
     Tosa_Tensor1Dto4D:$output
   );
 
   let hasFolder = 1;
+
+  let extraClassDeclaration = [{
+    /// Returns true when two result types are compatible for this op;
+    /// Method used by InferTypeOpInterface.
+    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+  }];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reduce_min
 //===----------------------------------------------------------------------===//
 def Tosa_ReduceMinOp : Tosa_Op<"reduce_min", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+    InferTensorType, Pure]> {
   let summary = "Reduce Min operator";
 
   let description = [{
     Reduce a tensor along the given axis with a minimum operation
   }];
 
   let arguments = (ins
     Tosa_Tensor1Dto4D:$input,
     I64Attr:$axis
   );
 
   let results = (outs
     Tosa_Tensor1Dto4D:$output
   );
 
   let hasFolder = 1;
+
+  let extraClassDeclaration = [{
+    /// Returns true when two result types are compatible for this op;
+    /// Method used by InferTypeOpInterface.
+    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+  }];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reduce_prod
 //===----------------------------------------------------------------------===//
 def Tosa_ReduceProdOp : Tosa_Op<"reduce_prod", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+    InferTensorType, Pure]> {
   let summary = "Reduce Prod operator";
 
   let description = [{
     Reduce a tensor along the given axis by computing the product of the axis.
   }];
 
   let arguments = (ins
     Tosa_Tensor1Dto4D:$input,
     I64Attr:$axis
   );
 
   let results = (outs
     Tosa_Tensor1Dto4D:$output
   );
 
   let hasFolder = 1;
+
+  let extraClassDeclaration = [{
+    /// Returns true when two result types are compatible for this op;
+    /// Method used by InferTypeOpInterface.
+    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+  }];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reduce_sum
 //===----------------------------------------------------------------------===//
 def Tosa_ReduceSumOp : Tosa_Op<"reduce_sum", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+    InferTensorType, Pure]> {
   let summary = "Reduce Sum operator";
 
   let description = [{
     Reduce a tensor along the given axis by computing the sum of the axis.
   }];
 
   let arguments = (ins
     Tosa_Tensor1Dto4D:$input,
     I64Attr:$axis
   );
 
   let results = (outs
     Tosa_Tensor1Dto4D:$output
   );
 
   let hasFolder = 1;
+
+  let extraClassDeclaration = [{
+    /// Returns true when two result types are compatible for this op;
+    /// Method used by InferTypeOpInterface.
+    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+  }];
 }
 
 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.9
 // Operator Class: Data Layout / Memory Reinterpretation.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//

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

if (outputShape[2] == ShapedType::kDynamic)
       outputShape[2] = inputShape.getDimSize(2);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 static LogicalResult ReduceInferReturnTypes(
-    ShapeAdaptor operandShape, IntegerAttr axis,
+    ShapeAdaptor operandShape, Type inputType, IntegerAttr axis,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   if (!operandShape.hasRank()) {
-    inferredReturnShapes.push_back(ShapedTypeComponents());
+    inferredReturnShapes.push_back(ShapedTypeComponents(inputType));
     return success();
   }
 
   SmallVector<int64_t> outputShape;
   operandShape.getDims(outputShape);
   int64_t axisVal = axis.getValue().getSExtValue();
   outputShape[axisVal] = 1;
-  inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
+  inferredReturnShapes.push_back(ShapedTypeComponents(outputShape, inputType));
   return success();
 }
 
 #define REDUCE_SHAPE_INFER(OP)                                                 \
   LogicalResult OP::inferReturnTypeComponents(                                 \
       MLIRContext *context, ::std::optional<Location> location,                \
       ValueShapeRange operands, DictionaryAttr attributes,                     \
       RegionRange regions,                                                     \
       SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {           \
-    return ReduceInferReturnTypes(operands.getShape(0),                        \
+    Type inputType =                                                           \
+        operands.getType()[0].cast<TensorType>().getElementType();             \
+    return ReduceInferReturnTypes(operands.getShape(0), inputType,             \
                                   attributes.get("axis").cast<IntegerAttr>(),  \
                                   inferredReturnShapes);                       \
+  }                                                                            \
+                                                                               \
+  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]));                       \
   }
 
 REDUCE_SHAPE_INFER(tosa::ReduceAllOp)
 REDUCE_SHAPE_INFER(tosa::ReduceAnyOp)
 REDUCE_SHAPE_INFER(tosa::ReduceMaxOp)
 REDUCE_SHAPE_INFER(tosa::ReduceMinOp)
 REDUCE_SHAPE_INFER(tosa::ReduceProdOp)
 REDUCE_SHAPE_INFER(tosa::ReduceSumOp)