[TOSA] Loosen folding restrictions for tosa.add,tosa.sub, tosa.mul

[sjw36]

Summary

Allow folding of different tensor types when the constant tensor is broadcast. Removed redundant and incorrect AddZero and MulOne canonical optimizations.

[sjw36]

update format

[sjw36]

@nicolasvasilache any feedback?

[eric-k256]

Adding @rsuderman, the author of the existing fold code as a more appropriate reviewer.

TOSA operators don't allow implicit rank changes, they need to be done with an explicit tosa.reshape. Implicit broadcast within a rank is allowed, so this could be changed to allow broadcast, but not dimension differences.

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

1

This looks like an accidental inclusion?

[rsuderman]

The change looks good, just make the fix the @eric-k256 pointed out.

[mgehre-amd]

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

624

I wonder if we can generalize the two branches into a single, non-templated isSplatZero(rhsAttr) && lhsTy == resultTy branch

[sjw36]

• simplified logic
• removed irrelevant change

[sjw36]

@mgehre-amd thanks, simplified the logic.
@rsuderman cleaned up PR.

[eric-k256]

I like the cleanup and the overall simplification of the code, but the broadcast issue still remains. Not specifically called out, but in addition to the marked issues, updating the commit message would also be good.

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

559

Same comment on the new patch, this should check against a broadcastable shape, ranks must match, and dimension sizes must be equal, or 1.
See resolveBroadcastShape in TosaOps.cpp for a reference.

mlir/test/Dialect/Tosa/canonicalize.mlir

23

This isn't a valid TOSA add, the rank of both tensors needs to match. So tensor<1x1x1xi32> would be valid (would broadcast along dimensions of size 1).

209

Same issue here.

[sjw36]

@eric-k256 thanks. I didn't see a hard requirement in the spec and was confused by the test above for add_zero_different_shape. Also turns out there is a AddZeroOptimization and MulOneOptimization that also allow mixed rank folding.. also kind of redundant. Shall I remove those as well?

[eric-k256]

In D145738#4228947, @sjw36 wrote:

@eric-k256 thanks. I didn't see a hard requirement in the spec and was confused by the test above for add_zero_different_shape. Also turns out there is a AddZeroOptimization and MulOneOptimization that also allow mixed rank folding.. also kind of redundant. Shall I remove those as well?

No problem. Section 1.9.4 covers broadcasting, it's possible if you were only looking at ADD/MUL it wouldn't be as obvious, but it says: TOSA broadcast requires the rank of both tensors to be the same. https://www.mlplatform.org/tosa/tosa_spec.html#_broadcasting

Yes, to reduce confusion, it would be good if you could also remove those. I must not have caught that when they went in.

[sjw36]

• added rank checks for the folders
• further simplified
• also removed redundant and incorrect pattern rewriters for AddZero and MulOne

[mgehre-amd]

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

617

Why do we need to check this here?

From what I understand, a sub is invalid when the rank of lhs and rhs differs. The verifier would catch this.
So passes/transforms can assume that they get valid IR as input, and thus don't need to each
verify the validity by themselves.

[sjw36]

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

617

Agreed. The canonicalize test below has these bad IR cases but the verifier doesn't seem to catch them in any case, @eric-k256 perhaps file a ticket?

[eric-k256]

Good catch @mgehre-amd. I had been thinking along a similar path to what you described, your comment helped solidify my thoughts.
If we track fixing the verifier, then agree that we can drop the check here. In terms of this change, we should change the tests so that if we fix the verifier, we don't then have to change these tests. (We're likely to find other tests that fail, but let's make our future work easier 😄 ). While doing that, changing the commit message to drop the mention of lower dimension helps avoid confusion.

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

617

Yes, a ticket sounds reasonable to me, with the caveat that I'm not sure how well those get monitored for MLIR overall, and less for TOSA. I expect that I can't be assigned to the issue. I don't know if you wanted to file it, or would like me to. If you file it, just let me know the number, and I'll subscribe. It seems like an important issue to fix, so I'll try to find someone here to give it a shot, unless you're already planning on working on it.

I just scanned through the code, and I think I understand the issue. The TOSA ops have the ResultsBroadcastableShape trait on them, which does do checking of the operands, however it uses a broader compatibility check than what TOSA allows, so the mismatch of ranks passes verification. It appears that we need to switch to a TOSA specific verification of broadcasting to catch this case.

[sjw36]

@eric-k256 I guess I will just remove the _different_shape tests and keep the new _bcast_ tests since that should be sufficient coverage. Is that acceptable?

[sjw36]

• removed rank checks
• removed invalid canonicalize tests (_different_shape)

[sjw36]

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

617

I unfortunately won't have time to work on that in the near future. Perhaps if you file the ticket you can assign it? I don't know the system here either.

[eric-k256]

Thanks for making all of the changes. Issue if you want to follow along: https://github.com/llvm/llvm-project/issues/61822. Not assigned to me, but if you're not working on it I'm less worried that someone else will pick it up without me knowing.

I may have overstepped on removing the existing optimization/canonicalization passes in favor of only the folder. @rsuderman, you'll probably want to make sure that's okay.
But from my perspective it looks ready to go.

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

let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 
-  let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: arithmetic_right_shift
 //===----------------------------------------------------------------------===//
 def Tosa_ArithmeticRightShiftOp : Tosa_Op<"arithmetic_right_shift", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
let arguments = (ins
     Tosa_Tensor:$input2,
     I32Attr:$shift
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 
-  let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: pow
 //===----------------------------------------------------------------------===//
 def Tosa_PowOp : Tosa_Op<"pow", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,

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

 //===- TosaCanonicalizations.cpp - Canonicalization patterns & folders ----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // \file
struct TransposeIsReshape : public OpRewritePattern<tosa::TransposeOp> {
   }
 };
 
 void TransposeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
   results.add<ConsolidateTransposeOptimization, TransposeIsReshape>(context);
 }
 
-struct AddZeroOptimization : public OpRewritePattern<tosa::AddOp> {
-  using OpRewritePattern::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(tosa::AddOp op,
-                                PatternRewriter &rewriter) const override {
-    auto input1 = op.getInput1();
-    auto input2 = op.getInput2();
-
-    DenseElementsAttr input1Attr;
-    if (matchPattern(input1, m_Constant(&input1Attr)) && input1Attr.isSplat() &&
-        input2.getType() == op.getType()) {
-      if (input1Attr.getType().getElementType().isa<IntegerType>() &&
-          input1Attr.getSplatValue<APInt>().isZero()) {
-        rewriter.replaceOp(op, op.getInput2());
-        return success();
-      }
-    }
-
-    DenseElementsAttr input2Attr;
-    if (matchPattern(input2, m_Constant(&input2Attr)) && input2Attr.isSplat() &&
-        input1.getType() == op.getType()) {
-      if (input2Attr.getType().getElementType().isa<IntegerType>() &&
-          input2Attr.getSplatValue<APInt>().isZero()) {
-        rewriter.replaceOp(op, op.getInput1());
-        return success();
-      }
-    }
-
-    return failure();
-  }
-};
-
-void AddOp::getCanonicalizationPatterns(RewritePatternSet &results,
-                                        MLIRContext *context) {
-  results.add<AddZeroOptimization>(context);
-}
-
-struct MulOneOptimization : public OpRewritePattern<tosa::MulOp> {
-  using OpRewritePattern::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(tosa::MulOp op,
-                                PatternRewriter &rewriter) const override {
-    auto input1 = op.getInput1();
-    auto input2 = op.getInput2();
-
-    DenseElementsAttr input1Attr;
-    if (matchPattern(input1, m_Constant(&input1Attr)) && input1Attr.isSplat() &&
-        input2.getType() == op.getType()) {
-      if (input1Attr.getType().getElementType().isa<FloatType>() &&
-          input1Attr.getSplatValue<APFloat>().isExactlyValue(1)) {
-        rewriter.replaceOp(op, op.getInput2());
-        return success();
-      }
-
-      if (input1Attr.getType().getElementType().isa<IntegerType>() &&
-          matchPattern(input1, m_One())) {
-        rewriter.replaceOp(op, op.getInput2());
-        return success();
-      }
-    }
-
-    DenseElementsAttr input2Attr;
-    if (matchPattern(input2, m_Constant(&input2Attr)) && input2Attr.isSplat() &&
-        input1.getType() == op.getType()) {
-      if (input2Attr.getType().getElementType().isa<FloatType>() &&
-          input2Attr.getSplatValue<APFloat>().isExactlyValue(1)) {
-        rewriter.replaceOp(op, op.getInput1());
-        return success();
-      }
-
-      if (input2Attr.getType().getElementType().isa<IntegerType>() &&
-          matchPattern(input2, m_One())) {
-        rewriter.replaceOp(op, op.getInput1());
-        return success();
-      }
-    }
-
-    return failure();
-  }
-};
-
-void MulOp::getCanonicalizationPatterns(RewritePatternSet &results,
-                                        MLIRContext *context) {
-  results.add<MulOneOptimization>(context);
-}
-
 struct MaterializePadValue : public OpRewritePattern<tosa::PadOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tosa::PadOp op,
                                 PatternRewriter &rewriter) const override {
     if (op.getPadConst())
       return failure();
 
if (lETy.isa<FloatType>()) {
       auto result = FloatFolder()(l, r);
       return DenseElementsAttr::get(returnTy, result);
     }
   }
 
   return {};
 }
 
+static bool isSplatZero(Type elemType, DenseElementsAttr val) {
+  if (elemType.isa<FloatType>())
+    return val && val.isSplat() && val.getSplatValue<APFloat>().isZero();
+  if (elemType.isa<IntegerType>())
+    return val && val.isSplat() && val.getSplatValue<APInt>().isZero();
+  return false;
+}
+
+static bool isSplatOne(Type elemType, DenseElementsAttr val, int64_t shift) {
+  if (elemType.isa<FloatType>())
+    return val && val.isSplat() &&
+           val.getSplatValue<APFloat>().isExactlyValue(1.0);
+  if (elemType.isa<IntegerType>()) {
+    const int64_t shifted = 1LL << shift;
+    return val && val.isSplat() &&
+           val.getSplatValue<APInt>().getSExtValue() == shifted;
+  }
+  return false;
+}
+
 OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
   auto lhsTy = getInput1().getType().dyn_cast<RankedTensorType>();
   auto rhsTy = getInput2().getType().dyn_cast<RankedTensorType>();
   auto resultTy = getType().dyn_cast<RankedTensorType>();
   if (!lhsTy || !rhsTy || !resultTy)
     return {};
-  if (lhsTy != rhsTy)
+  if (lhsTy.getRank() != rhsTy.getRank())
     return {};
 
   auto resultETy = resultTy.getElementType();
   auto lhsAttr = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>();
   auto rhsAttr = adaptor.getInput2().dyn_cast_or_null<DenseElementsAttr>();
 
-  if (lhsAttr && lhsAttr.isSplat() && resultETy.isa<FloatType>()) {
-    if (lhsAttr.getSplatValue<APFloat>().isZero())
-      return getInput2();
-  }
-
-  if (rhsAttr && rhsAttr.isSplat() && resultETy.isa<FloatType>()) {
-    if (rhsAttr.getSplatValue<APFloat>().isZero())
-      return getInput1();
-  }
-
-  if (lhsAttr && lhsAttr.isSplat() && resultETy.isa<IntegerType>()) {
-    if (lhsAttr.getSplatValue<APInt>().isZero())
-      return getInput2();
-  }
-
-  if (rhsAttr && rhsAttr.isSplat() && resultETy.isa<IntegerType>()) {
-    if (rhsAttr.getSplatValue<APInt>().isZero())
-      return getInput1();
-  }
+  if (lhsTy == resultTy && isSplatZero(resultETy, rhsAttr))
+  if (rhsTy == resultTy && isSplatZero(resultETy, lhsAttr))
 
   if (!lhsAttr || !rhsAttr)
     return {};
 
   return binaryFolder<std::plus<APInt>, std::plus<APFloat>>(lhsAttr, rhsAttr,
-                                                            lhsTy);
+                                                            resultTy);
 }
 
 OpFoldResult DivOp::fold(FoldAdaptor adaptor) {
   auto lhsTy = getInput1().getType().dyn_cast<RankedTensorType>();
   auto rhsTy = getInput2().getType().dyn_cast<RankedTensorType>();
   auto resultTy = getType().dyn_cast<RankedTensorType>();
   if (!lhsTy || !rhsTy || !resultTy)
     return {};

 OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
   auto lhs = getInput1();
   auto rhs = getInput2();
   auto lhsTy = lhs.getType().dyn_cast<RankedTensorType>();
   auto rhsTy = rhs.getType().dyn_cast<RankedTensorType>();
   auto resultTy = getType().dyn_cast<RankedTensorType>();
   if (!lhsTy || !rhsTy || !resultTy)
     return {};
-  if (lhsTy != rhsTy)
+  if (lhsTy.getRank() != rhsTy.getRank())
     return {};
 
   auto resultETy = resultTy.getElementType();
   auto lhsAttr = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>();
   auto rhsAttr = adaptor.getInput2().dyn_cast_or_null<DenseElementsAttr>();
 
-  if (lhsAttr && lhsAttr.isSplat() && resultETy.isa<FloatType>()) {
-    auto val = lhsAttr.getSplatValue<APFloat>();
-    if (val.isZero())
+  const int64_t shift = resultETy.isa<IntegerType>() ? getShift() : 0;
+  if (rhsTy == resultTy) {
+    if (isSplatZero(resultETy, lhsAttr))
       return lhsAttr;
-    if (val.isExactlyValue(1.0))
+    if (isSplatOne(resultETy, lhsAttr, shift))
       return rhs;
   }
-
-  if (rhsAttr && rhsAttr.isSplat() && resultETy.isa<FloatType>()) {
-    auto val = rhsAttr.getSplatValue<APFloat>();
-    if (val.isZero())
+  if (lhsTy == resultTy) {
+    if (isSplatZero(resultETy, rhsAttr))
       return rhsAttr;
-    if (val.isExactlyValue(1.0))
+    if (isSplatOne(resultETy, rhsAttr, shift))
       return lhs;
   }
 
-  if (lhsAttr && lhsAttr.isSplat() && resultETy.isa<IntegerType>()) {
-    auto val = lhsAttr.getSplatValue<APInt>();
-    if (val.isZero())
-      return lhsAttr;
-    const int64_t shift = getShift();
-    const int64_t shifted = 1LL << shift;
-    if (val.getSExtValue() == shifted)
-      return rhs;
-  }
-
-  if (rhsAttr && rhsAttr.isSplat() && resultETy.isa<IntegerType>()) {
-    auto val = rhsAttr.getSplatValue<APInt>();
-    const int64_t shift = getShift();
-    const int64_t shifted = 1LL << shift;
-    if (val.isZero())
-      return rhsAttr;
-    if (val.getSExtValue() == shifted)
-      return lhs;
-  }
-
-  return mulBinaryFolder(lhsAttr, rhsAttr, lhsTy, getShift());
+  return mulBinaryFolder(lhsAttr, rhsAttr, resultTy, getShift());
 }
 
 OpFoldResult SubOp::fold(FoldAdaptor adaptor) {
   auto lhsTy = getInput1().getType().dyn_cast<RankedTensorType>();
   auto rhsTy = getInput2().getType().dyn_cast<RankedTensorType>();
   auto resultTy = getType().dyn_cast<RankedTensorType>();
   if (!lhsTy || !rhsTy || !resultTy)
     return {};
-  if (lhsTy != rhsTy)
+  if (lhsTy.getRank() != rhsTy.getRank())
     return {};
 
   auto resultETy = resultTy.getElementType();
   auto lhsAttr = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>();
   auto rhsAttr = adaptor.getInput2().dyn_cast_or_null<DenseElementsAttr>();
 
-  if (rhsAttr && rhsAttr.isSplat() && resultETy.isa<FloatType>()) {
-    if (rhsAttr.getSplatValue<APFloat>().isZero())
-      return getInput1();
-  }
-
-  if (rhsAttr && rhsAttr.isSplat() && resultETy.isa<IntegerType>()) {
-    if (rhsAttr.getSplatValue<APInt>().isZero())
-      return getInput1();
-  }
+  if (lhsTy == resultTy && isSplatZero(resultETy, rhsAttr))
 
   if (!lhsAttr || !rhsAttr)
     return {};
 
   return binaryFolder<std::minus<APInt>, std::minus<APFloat>>(lhsAttr, rhsAttr,
-                                                              lhsTy);
+                                                              resultTy);
 }
 
 namespace {
 template <typename Cmp> struct ComparisonFold {
   ComparisonFold() = default;
   APInt operator()(const APInt &l, const APInt &r) {
     return APInt(1, Cmp()(l, r));
   }

mlir/test/Dialect/Tosa/canonicalize.mlir

 // CHECK-LABEL: @add_zero_different_shape
 func.func @add_zero_different_shape(%arg0: tensor<2x3xi32>) -> tensor<4x2x3xi32> {
   // CHECK: tosa.add
   %zeros = "tosa.const"() {value = dense<0> : tensor<4x2x3xi32>} : () -> tensor<4x2x3xi32>
   %1 = "tosa.add"(%arg0, %zeros) : (tensor<2x3xi32>, tensor<4x2x3xi32>) -> tensor<4x2x3xi32>
   return %1 : tensor<4x2x3xi32>
 }
 
+// CHECK-LABEL: @add_bcast_zero_int
+func.func @add_bcast_zero_int(%arg0: tensor<4x2x3xi32>) -> tensor<4x2x3xi32> {
+  // CHECK-NOT: tosa.add
+  // CHECK: return %arg0
+  %zeros = "tosa.const"() {value = dense<0> : tensor<1x1x1xi32>} : () -> tensor<1x1x1xi32>
+  %1 = "tosa.add"(%arg0, %zeros) : (tensor<4x2x3xi32>, tensor<1x1x1xi32>) -> tensor<4x2x3xi32>
+  return %1 : tensor<4x2x3xi32>
+}
 
 // CHECK-LABEL: @add_zero_int
 func.func @add_zero_int(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> {
   // CHECK: return %arg0
   // CHECK-NOT: tosa.add
   %zeros = "tosa.const"() {value = dense<0> : tensor<2x3xi32>} : () -> tensor<2x3xi32>
   %1 = "tosa.add"(%arg0, %zeros) : (tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
   return %1 : tensor<2x3xi32>

 func.func @mul_one_float(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
   // CHECK: return %arg0
   // CHECK-NOT: tosa.mul
   %ones = "tosa.const"() {value = dense<1.0> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
   %1 = "tosa.mul"(%arg0, %ones) {shift = 0 : i32} : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
   return %1 : tensor<2x3xf32>
 }
 
+// CHECK-LABEL: @mul_bcast_one_float
+func.func @mul_bcast_one_float(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
+  // CHECK: return %arg0
+  // CHECK-NOT: tosa.mul
+  %ones = "tosa.const"() {value = dense<1.0> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
+  %1 = "tosa.mul"(%ones, %arg0) {shift = 0 : i32} : (tensor<1x1xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
+  return %1 : tensor<2x3xf32>
+}
+
 // CHECK-LABEL: @mul_one_int
 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>
 }