[mlir] tosa.reshape - Add InferTensorType interface

[AviadCo]

Summary

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

[mgehre-amd]

Why are we inferring the return type of reshape and thus require the new_shape attribute, instead of removing the new_shape attribute and just taking the return type as the type that we need to reshape to?

[jpienaar]

In D148498#4273023, @mgehre-amd wrote:

Why are we inferring the return type of reshape and thus require the new_shape attribute, instead of removing the new_shape attribute and just taking the return type as the type that we need to reshape to?

I believe as that matches the spec most directly. But it is a good question if required in dialect.

This change didn't change this, but if this was moved to being on the type, then inferReturnTypes couldn't be used anymore as there is not sufficient information in the op's operands or attributes to infer return type. Else I would say the above question doesn't affect this, but a change to the above would require a revert of this. So might be good to get verification about my above assumption/desirability of this here.

[AviadCo]

Well, removing new_shape is a big refactor which I think reuqires propare discussion if we would like to do so.
I also agree with @jpienaar's comment, but until we decide otherwise I think that infer shape can be useful and missing functionality to check element type of TOSA reshape operation.
We can raise a discussion to get ride of new_shape attribute and part of the work will be to revert this change.

[jpienaar]

Looks fine given this is existing matching.

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

689

Nit: you can also use getElementTypeOrSelf(operands.getType()[0])

(new style also recommends cast<Foo>(x) rather than x.cast<Foo>() since that was added)

[AviadCo]

Refactored cast.

[AviadCo]

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

689

Good point, refactored.

[eric-k256]

In D148498#4282904, @AviadCo wrote:

Well, removing new_shape is a big refactor which I think reuqires propare discussion if we would like to do so.
I also agree with @jpienaar's comment, but until we decide otherwise I think that infer shape can be useful and missing functionality to check element type of TOSA reshape operation.
We can raise a discussion to get ride of new_shape attribute and part of the work will be to revert this change.

@jpienaar was right that new_shape was the most direct implementation of reshape from the TOSA specification to the dialect.

As long as the dialect can describe the functionality of TOSA reshape, I think we can have flexibility on the implementation. @AviadCo's plan to get this change merged, and separate out the removal of new_shape to a new thread seems reasonable.

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

 // TOSA Spec Section 2.9
 // Operator Class: Data Layout / Memory Reinterpretation.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: concat
 //===----------------------------------------------------------------------===//
 def Tosa_ConcatOp : Tosa_Op<"concat", [
-                              InferTensorType,
-    Pure]> {
+    InferTensorType, Pure]> {
   let summary = "Concatenates tensors along one dimension.";
 
   let description = [{
     Concatenate a variadic amount of tensors along a given axis. No data
     conversion happens during a concat operation.
   }];
 
   let arguments = (ins
def Tosa_PadOp : Tosa_Op<"pad", [
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reshape
 //===----------------------------------------------------------------------===//
 def Tosa_ReshapeOp: Tosa_Op<"reshape", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+  InferTensorType, Pure]> {
   let summary = "Reshape operator";
 
   let description = [{
     Returns a tensor with the same type/values as the input, with a new shape
     specified by the shape argument. Reshape may operate on tensors of any rank.
     No data conversion happens during a reshape operation.
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     DenseI64ArrayAttr:$new_shape
   );
 
   let results = (outs
     Tosa_RankedTensor:$output
   );
+
+  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: reverse
 //===----------------------------------------------------------------------===//
 def Tosa_ReverseOp: Tosa_Op<"reverse", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>, Pure]> {

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

if (dim != ShapedType::kDynamic)
       dim *= multiples[i];
     outputShape.push_back(dim);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
+bool tosa::ReshapeOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
+  if (l.size() != r.size() || l.size() != 1)
+    return false;
+  return getElementTypeOrSelf(l[0]) == getElementTypeOrSelf(r[0]);
+}
+
 LogicalResult tosa::ReshapeOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   ReshapeOpAdaptor adaptor(operands, attributes);
   ShapeAdaptor inputShape = operands.getShape(0);
+  Type inputType = getElementTypeOrSelf(operands.getType()[0]);
   llvm::SmallVector<int64_t> newShapeValue =
       convertToMlirShape(adaptor.getNewShape());
 
   // We cannot infer from the total number of elements so we must take the
   // shape attribute as exact.
   if (!inputShape.hasRank() || !inputShape.hasStaticShape()) {
-    inferredReturnShapes.push_back(ShapedTypeComponents(newShapeValue));
+    inferredReturnShapes.push_back(
+        ShapedTypeComponents(newShapeValue, inputType));
     return success();
   }
 
   // Determine the number of elements covered by the slice of all static
   // dimensions. This allows us to infer the length of the remaining dynamic
   // dimension.
   int64_t numElements = inputShape.getNumElements();
   int64_t staticMul = 1;
   for (auto val : newShapeValue) {
     if (!ShapedType::isDynamic(val)) {
       staticMul *= val;
     }
   }
 
   // Determine the length of the dynamic dimension.
   for (auto &val : newShapeValue) {
     if (ShapedType::isDynamic(val))
       val = numElements / staticMul;
   }
 
-  inferredReturnShapes.push_back(ShapedTypeComponents(newShapeValue));
+  inferredReturnShapes.push_back(
+      ShapedTypeComponents(newShapeValue, inputType));
   return success();
 }
 
 mlir::LogicalResult tosa::ReshapeOp::verify() {
   ShapedType inputType = getInput1().getType().cast<ShapedType>();
   ShapedType outputType = getType().cast<ShapedType>();
 
   if (inputType.hasStaticShape() && outputType.hasStaticShape()) {

mlir/test/Dialect/Tosa/invalid.mlir

 
 // -----
 
 func.func @test_reduce_prod_type_mismatch(%arg0 : tensor<2x3x4x5xf32>) -> () {
   // expected-error@+1 {{'tosa.reduce_prod' op inferred type(s) 'tensor<2x1x4x5xf32>' are incompatible with return type(s) of operation 'tensor<2x3x4x5xf32>'}}
   %0 = "tosa.reduce_prod"(%arg0) {axis = 1 : i64} : (tensor<2x3x4x5xf32>) -> tensor<2x3x4x5xf32>
   return
 }
+
+// -----
+
+func.func @test_reshape_type_mismatch(%arg0 : tensor<13x21x3xf32>) -> () {
+  // expected-error@+1 {{'tosa.reshape' op inferred type(s) 'tensor<13x21x3x1xf32>' are incompatible with return type(s) of operation 'tensor<13x21x3x1xi32>'}}
+  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 13, 21, 3, 1>} : (tensor<13x21x3xf32>) -> tensor<13x21x3x1xi32>
+  return
+}