[mlir][tosa][fix] Add proper type checking trait for tosa mul

[tatwaichong]

Summary

when operating integer type tensors, tosa elementwise multiplication
requires the element type of result to be a 32-bit integer rather
than the same type as inputs.

Change-Id: Ifd3d7ebd879be5c6b2c8e23aa6d7ef41f39c6d41

[jpienaar]

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

776

So this is a Tensor of I32? E.g., could this just be Tosa_Int32Tensor and would this cover that?

[tatwaichong]

This is a fix to response the mistake introduced in D150472.
Quote ""tosa.mul"(%0, %1) {shift = 0 : i32} : (tensor<4xi8>, tensor<4xi8>) -> tensor<4xi32> is marked as illegal because result and operand types don't match, even though i8xi8 = i32 is according to the spec".

I propose to create a custom for tosa mul to check if its operands and result meet the spec requirement. Do you think if this a good manner to cope with this demand?

[tatwaichong]

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

776

Smart idea, but the input could be floating type.

[tatwaichong]

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

776

Or Tosa_Int32OrFpTensor?

[tatwaichong]

missing < in the expected output string of the test.

[tatwaichong]

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

776

but this way doesn't cover the type checking between operands and result of the operation. e.g. won't check if the operands and result are the same type when operating floating type. Is it right?

[mgehre-amd]

Thanks for working on this!
Some clarification from my side:

• We should make sure that tosa.mul"(%0, %1) {shift = 0 : i32} : (tensor<4xi8>, tensor<4xi8>) -> tensor<4xi32> compiles (because its legal according to the standard).
• We don't need to disallow tosa.mul"(%0, %1) {shift = 0 : i32} : (tensor<4xi16>, tensor<4xi16>) -> tensor<4xi16>; the LLVM implementation of TOSA is more relaxed than the standard in many ways,

and I think that freedom has been useful.

What do you think?

[tatwaichong]

update the trait in a relexed way to allow the result type can be other size but not i32 as long as this is a valid size.

[tatwaichong]

Yes. This demand is reasonable. We loosen the trait to accept those use cases.

[mgehre-amd]

Thanks!

[tatwaichong]

Is there any remaining issues? could you land this for me?

[eric-k256]

Can you rebase this? There is a conflict and probably easiest for you to fix. Then I can land it.

[tatwaichong]

rebase and reslove conflict

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

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_TOSA_IR_TOSAOPS_H
 #define MLIR_DIALECT_TOSA_IR_TOSAOPS_H
 
 #include "mlir/Bytecode/BytecodeOpInterface.h"
 #include "mlir/Dialect/Traits.h"
+#include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // TOSA dialect and structs includes.
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tosa/IR/TosaOpsDialect.h.inc"
 
 namespace mlir {
 class PatternRewriter;
 
 namespace tosa {
 
 #include "mlir/Dialect/Tosa/IR/TosaInterfaces.h.inc"
 
 } // namespace tosa
+
+namespace OpTrait {
+namespace tosa {
+
+// This trait verifies if the element type amoung operands and result
+// of multiplication match tosa specification.
+template <typename ConcreteType>
+class MulOperandsAndResultElementType
+    : public TraitBase<ConcreteType, MulOperandsAndResultElementType> {
+public:
+  static LogicalResult verifyTrait(Operation *op) {
+    auto resElemType = getElementTypeOrSelf(op->getResult(0));
+
+    // In cases of floating point type, op requires the same element
+    // type for all operands and result.
+    if (llvm::isa<FloatType>(resElemType))
+      return impl::verifySameOperandsAndResultElementType(op);
+
+    if (auto resIntType = resElemType.dyn_cast<IntegerType>()) {
+      IntegerType lhsIntType =
+          getElementTypeOrSelf(op->getOperand(0)).cast<IntegerType>();
+      IntegerType rhsIntType =
+          getElementTypeOrSelf(op->getOperand(1)).cast<IntegerType>();
+      if (lhsIntType != rhsIntType)
+        return op->emitOpError(
+            "requires the same element type for all operands");
+
+      // Though the spec requires the element type of result to be i32, a more
+      // relaxed way is provided at dialect level for easier cooperating with
+      // other dialects.
+      if (lhsIntType.getWidth() > resIntType.getWidth())
+        return op->emitOpError("invalid data type size for operands or result");
+
+      return success();
+    }
+
+    return failure();
+  }
+};
+
+} // namespace tosa
+} // namespace OpTrait
+
 } // namespace mlir
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/Tosa/IR/TosaAttributes.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tosa/IR/TosaOps.h.inc"
 
 #endif // MLIR_DIALECT_TOSA_IR_TOSAOPS_H

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

let arguments = (ins
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
+def MulOperandsAndResultElementType :
+  NativeOpTrait<"MulOperandsAndResultElementType"> {
+  let cppNamespace = "mlir::OpTrait::tosa";
+}
+
 //===----------------------------------------------------------------------===//
 // Operator: mul
 //===----------------------------------------------------------------------===//
 def Tosa_MulOp : Tosa_ElementwiseOp<"mul", [
     Commutative,
-    SameOperandsAndResultElementType]> {
+    MulOperandsAndResultElementType]> {
   let summary = "Multiplication operator";
 
   let description = [{
     Elementwise multiplication (Hadamard product) of input1 and input2.
     Axis of size 1 will be broadcast, as necessary.
     i8/i16 input type can be promoted to i32 result type.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2,
     I32Attr:$shift
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: pow
 //===----------------------------------------------------------------------===//

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

func.func @test_simple_f16(%arg0: tensor<1xf16>) -> () {
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_simple_i16
 func.func @test_simple_i16(%arg0: tensor<1xi16>) -> () {
   // CHECK: linalg.generic
+  // CHECK: arith.extsi
+  // CHECK: arith.extsi
   // CHECK: arith.muli
-  %0 = "tosa.mul"(%arg0, %arg0) {shift = 0 : i32} : (tensor<1xi16>, tensor<1xi16>) -> tensor<1xi16>
+  %0 = "tosa.mul"(%arg0, %arg0) {shift = 0 : i32} : (tensor<1xi16>, tensor<1xi16>) -> tensor<1xi32>
 
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_simple_ui8
 func.func @test_simple_ui8(%arg0: tensor<1xui8>) -> () {

mlir/test/Dialect/Tosa/constant-op-fold.mlir

func.func @fold_mul_one_lhs_i32(%arg0: tensor<i32>) -> tensor<i32> {
   %mul = "tosa.mul"(%one, %arg0) {shift = 6 : i32} : (tensor<i32>, tensor<i32>) -> tensor<i32>
   // CHECK: return %arg0
   return %mul : tensor<i32>
 }
 
 // -----
 
 // CHECK-LABEL: @fold_mul_splat_i8
-func.func @fold_mul_splat_i8() -> tensor<10xi8> {
+func.func @fold_mul_splat_i8() -> tensor<10xi32> {
   %one = "tosa.const"() {value = dense<17> : tensor<10xi8>} : () -> tensor<10xi8>
   %two = "tosa.const"() {value = dense<32> : tensor<10xi8>} : () -> tensor<10xi8>
-  %mul = "tosa.mul"(%one, %two) {shift = 3 : i32} : (tensor<10xi8>, tensor<10xi8>) -> tensor<10xi8>
-  // CHECK: %[[THREE:.+]] = "tosa.const"() <{value = dense<68> : tensor<10xi8>}
+  %mul = "tosa.mul"(%one, %two) {shift = 3 : i32} : (tensor<10xi8>, tensor<10xi8>) -> tensor<10xi32>
+  // CHECK: %[[THREE:.+]] = "tosa.const"() <{value = dense<68> : tensor<10xi32>}
   // CHECK: return %[[THREE]]
-  return %mul : tensor<10xi8>
+  return %mul : tensor<10xi32>
 }
 
 // -----
 
 // CHECK-LABEL: @fold_mul_splat_f32
 func.func @fold_mul_splat_f32() -> tensor<10xf32> {
   %one = "tosa.const"() {value = dense<3.0> : tensor<10xf32>} : () -> tensor<10xf32>
   %two = "tosa.const"() {value = dense<2.0> : tensor<10xf32>} : () -> tensor<10xf32>

mlir/test/Dialect/Tosa/ops.mlir

 // -----
 // CHECK-LABEL: mul
 func.func @test_mul(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x1x3xf32>) -> tensor<13x21x3xf32> {
   %0 = "tosa.mul"(%arg0, %arg1)  { shift = 1 : i32 } : (tensor<13x21x3xf32>, tensor<13x1x3xf32>) -> tensor<13x21x3xf32>
   return %0 : tensor<13x21x3xf32>
 }
 
 // -----
+// CHECK-LABEL: mul
+func.func @test_mul_relaxed_result_type(%arg0: tensor<13x21x3xi16>, %arg1: tensor<13x1x3xi16>) -> tensor<13x21x3xi16> {
+  %0 = "tosa.mul"(%arg0, %arg1)  { shift = 1 : i32 } : (tensor<13x21x3xi16>, tensor<13x1x3xi16>) -> tensor<13x21x3xi16>
+  return %0 : tensor<13x21x3xi16>
+}
+
+// -----
 // CHECK-LABEL: pow
 func.func @test_pow(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x21x1xf32>) -> tensor<13x21x3xf32> {
   %0 = "tosa.pow"(%arg0, %arg1) : (tensor<13x21x3xf32>, tensor<13x21x1xf32>) -> tensor<13x21x3xf32>
   return %0 : tensor<13x21x3xf32>
 }
 
 // -----
 // CHECK-LABEL: sub