[mlir][tosa] Fix constant folding of tosa.mul

[sabauma]

Summary

The constant folder for tosa.mul produces a tensor attribute whose type
may not match the result type of the operation when broadcasting is
needed. This results in a tosa.const op whose attribute's type does not
match the type of the const op.
This change explicitly expands the attribute to the expected result
type.

[sabauma]

Add verification check on the structure of tosa.const

[jpienaar]

We should not be folding if the type changes AFAIK (I think that may not be enforced ...). Is this change effectively ensuring it doesn't change types?

[sabauma]

In D150439#4342736, @jpienaar wrote:

We should not be folding if the type changes AFAIK (I think that may not be enforced ...). Is this change effectively ensuring it doesn't change types?

This change is ensuring that the attribute produced by folding tosa.mul has a type which matches the result type of the tosa.mul operation. The folder for tosa.mul was not broadcasting the resulting attribute when multiplying by zero. The resizeSplat is performing the required broadcast, in this case.
The IR operations are always well typed (input/output types), before and after folding, as far as I can tell. However, you end up with `tosa.const operations like this:

mlir
%0 = "tosa.const"() {value = dense<0.000000e+00> : tensor<1x1xf32>} : () -> tensor<100x100xf32>

I raised a question about this on this discourse thread: https://discourse.llvm.org/t/correctness-properties-of-fold-methods/70564

The change to the tosa.const definition is just adding a sanity check that would have caught this issue earlier.

[sabauma]

@jpienaar Would you mind submitting this change then, since I do not have commit privileges?

[eric-k256]

I thought I had accepted this already, but don't see it recorded. LGTM for the reasons described in the discourse thread.

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

OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
 
   auto resultETy = resultTy.getElementType();
   auto lhsAttr = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>();
   auto rhsAttr = adaptor.getInput2().dyn_cast_or_null<DenseElementsAttr>();
 
   const int64_t shift = resultETy.isa<IntegerType>() ? getShift() : 0;
   if (rhsTy == resultTy) {
     if (isSplatZero(resultETy, lhsAttr))
-      return lhsAttr;
+      return lhsAttr.resizeSplat(resultTy);
     if (isSplatOne(resultETy, lhsAttr, shift))
       return rhs;
   }
   if (lhsTy == resultTy) {
     if (isSplatZero(resultETy, rhsAttr))
-      return rhsAttr;
+      return rhsAttr.resizeSplat(resultTy);
     if (isSplatOne(resultETy, rhsAttr, shift))
       return lhs;
   }
 
   return mulBinaryFolder(lhsAttr, rhsAttr, resultTy, getShift());
 }
 
 OpFoldResult SubOp::fold(FoldAdaptor adaptor) {

mlir/test/Dialect/Tosa/canonicalize.mlir

 func.func @mul_one_int(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> {
   // CHECK: return %arg0
   // CHECK-NOT: tosa.mul
   %ones = "tosa.const"() {value = dense<1> : tensor<2x3xi32>} : () -> tensor<2x3xi32>
   %1 = "tosa.mul"(%arg0, %ones) {shift = 0 : i32} : (tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
   return %1 : tensor<2x3xi32>
 }
 
+// CHECK-LABEL: @mul_zero_broadcast
+func.func @mul_zero_broadcast(%arg0: tensor<2x3xf32>) -> (tensor<2x3xf32>, tensor<2x3xf32>) {
+  // CHECK: %[[ZERO:.*]] = "tosa.const"() <{value = dense<0.000000e+00> : tensor<2x3xf32>}> : () -> tensor<2x3xf32> 
+  // CHECK-NOT: tosa.mul
+  %zeros = "tosa.const"() {value = dense<0.0> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
+  %1 = "tosa.mul"(%arg0, %zeros) {shift = 0 : i32} : (tensor<2x3xf32>, tensor<1x1xf32>) -> tensor<2x3xf32>
+
+  // CHECK-NOT: tosa.mul
+  // CHECK: return %[[ZERO]], %[[ZERO]]
+  %2 = "tosa.mul"(%zeros, %arg0) {shift = 0 : i32} : (tensor<1x1xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
+  return %1, %2 : tensor<2x3xf32>, tensor<2x3xf32>
+}
+
 // CHECK-LABEL: @select_same_value
 func.func @select_same_value(%arg0: tensor<2x3xi1>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
   %0 = "tosa.select"(%arg0, %arg1, %arg1) : (tensor<2x3xi1>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
   // CHECK: return %arg1
   // CHECK-NOT: tosa.select
   return %0 : tensor<2x3xi32>
 }