Convert StableHLO reshape to TOSA reshape (#2861)

[tosa.reshape](https://mlir.llvm.org/docs/Dialects/TOSA/#tosareshape-mlirtosareshapeop)
accepts shape as a second argument.
Conversion of the StableHLO reshape requires insertion of the
[tosa.const_shape](https://mlir.llvm.org/docs/Dialects/TOSA/#tosaconst_shape-mlirtosaconstshapeop)
op as an argument for the tosa.reshape.

Author: vvish

stablehlo/conversions/tosa/tests/unary.mlir

@@ -155,3 +155,11 @@ func.func @while(%arg0: tensor<i32>) -> tensor<i32> {
   }) : (tensor<i32>) -> (tensor<i32>)
   return %0 : tensor<i32>
 }
+
+// CHECK-LABEL: @reshape
+func.func @reshape(%arg0 : tensor<2x3xf32>) -> tensor<6xf32> {
+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<6> : tensor<1xindex>} : () -> !tosa.shape<1>
+  // CHECK: tosa.reshape %arg0, %[[VAR0]]
+  %0 = "stablehlo.reshape"(%arg0) : (tensor<2x3xf32>) -> tensor<6xf32>
+  return %0 : tensor<6xf32>
+}

stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.cpp

@@ -491,6 +491,38 @@ struct ConvertStablehloWhileOp : public OpRewritePattern<stablehlo::WhileOp> {
   }
 };
 
+struct ConvertStablehloReshapeOp
+    : public OpRewritePattern<mlir::stablehlo::ReshapeOp> {
+  using OpRewritePattern<mlir::stablehlo::ReshapeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(stablehlo::ReshapeOp op,
+                                PatternRewriter& rewriter) const override {
+    auto resultType = op.getResult().getType();
+    if (!resultType.hasStaticShape()) {
+      return rewriter.notifyMatchFailure(op, "result tensor must be static");
+    }
+
+    auto resultShape = resultType.getShape();
+    SmallVector<int64_t, 8> dimensions(resultShape.begin(), resultShape.end());
+
+    RankedTensorType shapeTensorType = RankedTensorType::get(
+        {static_cast<int64_t>(dimensions.size())}, rewriter.getIndexType());
+
+    auto denseAttr = DenseIntElementsAttr::get(shapeTensorType, dimensions);
+    auto shapeType =
+        tosa::shapeType::get(rewriter.getContext(), dimensions.size());
+
+    auto constShapeOp =
+        rewriter.create<tosa::ConstShapeOp>(op.getLoc(), shapeType, denseAttr);
+
+    auto reshapeOp = rewriter.create<tosa::ReshapeOp>(
+        op.getLoc(), resultType, op.getOperand(), constShapeOp);
+
+    rewriter.replaceOp(op, reshapeOp.getResult());
+    return success();
+  }
+};
+
 LogicalResult StablehloLegalizeToTosaPass::initialize(MLIRContext* ctx) {
   RewritePatternSet patternList(ctx);
   populateGeneratedPDLLPatterns(patternList);
@@ -509,6 +541,9 @@ LogicalResult StablehloLegalizeToTosaPass::initialize(MLIRContext* ctx) {
   patternList.addWithLabel<ConvertStablehloTransposeOp>({"StablehloTranspose"},
                                                         ctx);
   patternList.addWithLabel<ConvertStablehloWhileOp>({"StablehloWhile"}, ctx);
+  patternList.addWithLabel<ConvertStablehloReshapeOp>({"StablehloReshape"},
+                                                      ctx);
+
   patterns = std::move(patternList);
   return success();
 }