[mlir][tosa] Fix tosa.reshape failures due to implicit broadcasting

[rsuderman]

Summary

Make broadcastable needs the output shape to determine whether the operation
includes additional broadcasting. Include some canonicalizations for TOSA
to remove unneeded reshape.

[rsuderman]

Rebase on head to validate pre-merge is safe.

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

def Tosa_ReshapeOp: Tosa_Op<"reshape", [
   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 arguments = (ins
     Tosa_Tensor:$input1,
     I64ArrayAttr:$new_shape
   );
 
   let results = (outs
     Tosa_RankedTensor:$output
   );

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

for (auto it : llvm::enumerate(type.getShape())) {
         newShape.push_back(it.value());
         affineExprs.push_back(
             mlir::getAffineDimExpr(it.index(), rewriter.getContext()));
       }
     }
 
     if (newShape.size() != rank) {
       operand = rewriter.create<tosa::ReshapeOp>(
-          loc, RankedTensorType::get(newShape, type.getElementType()), operand);
+          loc, RankedTensorType::get(newShape, type.getElementType()), operand,
+          rewriter.getI64ArrayAttr(newShape));
     }
 
     operands.push_back(operand);
     indexingMaps.push_back(AffineMap::get(
         /*dimCount=*/type.getRank(), /*symbolCount=*/0, affineExprs,
         rewriter.getContext()));
   }
 

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

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "mlir/Transforms/RegionUtils.h"
 
 using namespace mlir;
 using namespace mlir::tosa;
 
 #include "mlir/Dialect/Tosa/IR/TosaOpsDialect.cpp.inc"
Operation *TosaDialect::materializeConstant(OpBuilder &builder, Attribute value,
   // Tosa dialect constants only support ElementsAttr unlike standard dialect
   // constant which supports all attributes.
   if (value.isa<ElementsAttr>())
     return builder.create<tosa::ConstOp>(loc, type, value.cast<ElementsAttr>());
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
+// Operator Canonicalizers.
+//===----------------------------------------------------------------------===//
+
+struct RemoveReshapeNoop : public OpRewritePattern<tosa::ReshapeOp> {
+  using OpRewritePattern<tosa::ReshapeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::ReshapeOp op,
+                                PatternRewriter &rewriter) const override {
+    if (op.input1().getType() != op.getType())
+      return failure();
+
+    rewriter.replaceOp(op, op.input1());
+    return success();
+  }
+};
+
+struct ReshapeReshapeOptimization : public OpRewritePattern<tosa::ReshapeOp> {
+  using OpRewritePattern<tosa::ReshapeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::ReshapeOp op,
+                                PatternRewriter &rewriter) const override {
+    Value input = op.input1();
+    Operation *definingOp = input.getDefiningOp();
+    if (!definingOp)
+      return failure();
+
+    if (tosa::ReshapeOp reshapeOp = dyn_cast<tosa::ReshapeOp>(definingOp)) {
+      rewriter.replaceOpWithNewOp<tosa::ReshapeOp>(
+          op, op.getType(), reshapeOp.input1(), op.new_shape());
+      return success();
+    }
+
+    return failure();
+  }
+};
+
+void ReshapeOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                            MLIRContext *context) {
+  results.insert<ReshapeReshapeOptimization, RemoveReshapeNoop>(context);
+}
+
+//===----------------------------------------------------------------------===//
 // Operator Folders.
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ConstOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.empty() && "constant has no operands");
   return valueAttr();
 }
 

mlir/lib/Dialect/Tosa/Transforms/TosaMakeBroadcastable.cpp

static LogicalResult reshapeLowerToHigher(PatternRewriter &rewriter,
   ArrayRef<int64_t> higherRankShape =
       higherTensorValue.getType().cast<RankedTensorType>().getShape();
   (void)higherRankShape;
   ArrayRef<int64_t> lowerRankShape =
       lowerTensorValue.getType().cast<RankedTensorType>().getShape();
 
   SmallVector<int64_t, 4> reshapeOutputShape;
 
-  computeReshapeOutput(higherRankShape, lowerRankShape, reshapeOutputShape);
+  computeReshapeOutput(outputType.getShape(), lowerRankShape,
+                       reshapeOutputShape);
 
   auto reshapeInputType = lowerTensorValue.getType().cast<RankedTensorType>();
   auto reshapeOutputType = RankedTensorType::get(
       ArrayRef<int64_t>(reshapeOutputShape), reshapeInputType.getElementType());
 
   // Verify the rank agrees with the output type if the output type is ranked.
   if (outputType) {
     if (outputType.getShape().size() != reshapeOutputShape.size() ||

mlir/test/Dialect/Tosa/broadcast.mlir

 // CHECK-LABEL: broadcast18
 func @test_broadcast18(%arg0: tensor<14x1xf32>, %arg1: tensor<1x15xf32>) -> tensor<14x15xf32> {
   //  CHECK: add
   %0 = "tosa.add"(%arg0, %arg1) : (tensor<14x1xf32>, tensor<1x15xf32>) -> tensor<14x15xf32>
   return %0 : tensor<14x15xf32>
 }
 
 // -----
+// CHECK-LABEL: broadcast19
+func @broadcast19(%arg0: tensor<64x64x1xf32>, %arg1: tensor<1x17xf32>) -> (tensor<64x64x17xf32> ) {
+  //  CHECK: reshape
+  //  CHECK: sub
+  %0 = "tosa.sub"(%arg0, %arg1) : (tensor<64x64x1xf32>, tensor<1x17xf32>) -> tensor<64x64x17xf32>
+  return %0 : tensor<64x64x17xf32>
+}
+
+// -----
 // CHECK-LABEL: broadcast_mul
 func @test_broadcast_mul(%arg0: tensor<15x14xi32>, %arg1: tensor<17x16x15x14xi32>) -> tensor<17x16x15x14xi32> {
   //  CHECK: reshape
   %0 = "tosa.mul"(%arg0, %arg1) {shift = 1 : i32 } : (tensor<15x14xi32>, tensor<17x16x15x14xi32>) -> tensor<17x16x15x14xi32>
   return %0 : tensor<17x16x15x14xi32>
 }
 
 // -----