Revert "[mlir][tosa] Change 'shape' of RESHAPE from attribute to input shape … (llvm#125789)"

This reverts commit 571a987.

Author: bjacob

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

@@ -1625,7 +1625,7 @@ def Tosa_ReshapeOp : Tosa_InferTensorTypeOp<"reshape"> {
 
   let arguments = (ins
     Tosa_Tensor:$input1,
-    Tosa_Shape:$shape
+    DenseI64ArrayAttr:$new_shape
   );
 
   let results = (outs

mlir/include/mlir/Dialect/Tosa/Utils/ConversionUtils.h

@@ -230,11 +230,8 @@ SmallVector<T> applyTOSAPermutation(ArrayRef<T> input,
 }
 
 // Computes shape value using tosa const_shape op.
-Value getTosaConstShape(ImplicitLocOpBuilder &builder,
-                        llvm::ArrayRef<int64_t> shape);
 Value getTosaConstShape(PatternRewriter &rewriter, Location loc,
                         llvm::ArrayRef<int64_t> shape);
-
 SmallVector<int64_t> convertFromMlirShape(ArrayRef<int64_t> shape);
 
 bool getConstShapeValue(Operation *op,

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

@@ -1956,10 +1956,9 @@ struct TileConverter : public OpConversionPattern<tosa::TileOp> {
           nestedBuilder.create<linalg::YieldOp>(op.getLoc(), *args.begin());
         });
 
-    auto shapeValue = getTosaConstShape(
-        rewriter, loc, mlir::tosa::convertFromMlirShape(resultTy.getShape()));
     rewriter.replaceOpWithNewOp<tosa::ReshapeOp>(
-        op, resultTy, genericOp.getResult(0), shapeValue);
+        op, resultTy, genericOp.getResult(0),
+        rewriter.getDenseI64ArrayAttr(resultTy.getShape()));
     return success();
   }
 };

mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp

@@ -16,7 +16,6 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
-#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/DialectConversion.h"
 
@@ -236,12 +235,7 @@ class ReshapeConverter : public OpConversionPattern<tosa::ReshapeOp> {
       return rewriter.notifyMatchFailure(reshape.getLoc(),
                                          "expected input type to be tensor");
     }
-
-    llvm::SmallVector<int64_t> newShape;
-    if (!tosa::getConstShapeValue(reshape.getShape().getDefiningOp(),
-                                  newShape)) {
-      return failure();
-    }
+    auto newShape = reshape.getNewShape();
 
     // Infer all intermediate types
     auto inputType = inferReshapeInputType(input, newShape);

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

@@ -180,7 +180,7 @@ struct TransposeIsReshape : public OpRewritePattern<tosa::TransposeOp> {
 
     rewriter.replaceOpWithNewOp<tosa::ReshapeOp>(
         op, op.getType(), op.getInput1(),
-        getTosaConstShape(rewriter, op.getLoc(), newShape));
+        rewriter.getDenseI64ArrayAttr(newShape));
     return success();
   }
 };
@@ -948,12 +948,8 @@ OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
     if (!getInput1().hasOneUse())
       return {};
 
-    llvm::SmallVector<int64_t> shapeVec;
-    if (!tosa::getConstShapeValue(getShape().getDefiningOp(), shapeVec))
-      return {};
-
     return operand.reshape(
-        llvm::cast<ShapedType>(operand.getType()).clone(shapeVec));
+        llvm::cast<ShapedType>(operand.getType()).clone(getNewShape()));
   }
 
   return {};

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

@@ -1335,16 +1335,8 @@ LogicalResult tosa::ReshapeOp::inferReturnTypeComponents(
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   ShapeAdaptor inputShape(adaptor.getInput1().getType());
   Type inputType = getElementTypeOrSelf(adaptor.getInput1().getType());
-  llvm::SmallVector<int64_t> newShapeValue;
-  if (!tosa::getConstShapeValue(adaptor.getShape().getDefiningOp(),
-                                newShapeValue)) {
-    auto rank = cast<tosa::shapeType>(adaptor.getShape().getType()).getRank();
-    SmallVector<int64_t> fallback(rank, ShapedType::kDynamic);
-    inferredReturnShapes.push_back(ShapedTypeComponents(fallback, inputType));
-    return success();
-  } else {
-    newShapeValue = convertToMlirShape(newShapeValue);
-  }
+  llvm::SmallVector<int64_t> newShapeValue =
+      convertToMlirShape(adaptor.getNewShape());
 
   // We cannot infer from the total number of elements so we must take the
   // shape attribute as exact.
@@ -1380,19 +1372,13 @@ llvm::LogicalResult tosa::ReshapeOp::verify() {
   TensorType inputType = getInput1().getType();
   RankedTensorType outputType = getType();
 
-  SmallVector<int64_t> shapeValues;
-  if (!tosa::getConstShapeValue(getShape().getDefiningOp(), shapeValues)) {
-    // skip following checks if shape is not constant
-    return mlir::success();
-  }
-
-  if ((int64_t)shapeValues.size() != outputType.getRank())
+  if ((int64_t)getNewShape().size() != outputType.getRank())
     return emitOpError() << "new shape does not match result rank";
 
   for (auto [newShapeDim, outputShapeDim] :
-       zip(shapeValues, outputType.getShape())) {
-    if (newShapeDim != -1 && newShapeDim != ShapedType::kDynamic &&
-        outputShapeDim != ShapedType::kDynamic && newShapeDim != outputShapeDim)
+       zip(getNewShape(), outputType.getShape())) {
+    if (newShapeDim != -1 && outputShapeDim != ShapedType::kDynamic &&
+        newShapeDim != outputShapeDim)
       return emitOpError() << "new shape is inconsistent with result shape";
 
     if (newShapeDim != ShapedType::kDynamic && newShapeDim < -1)
@@ -1411,18 +1397,18 @@ llvm::LogicalResult tosa::ReshapeOp::verify() {
     }
 
     int64_t newShapeElementsNum = std::accumulate(
-        shapeValues.begin(), shapeValues.end(), 1LL,
+        getNewShape().begin(), getNewShape().end(), 1LL,
         [](int64_t acc, int64_t dim) { return (dim > 0) ? acc * dim : acc; });
     bool isStaticNewShape =
-        llvm::all_of(shapeValues, [](int64_t s) { return s > 0; });
+        llvm::all_of(getNewShape(), [](int64_t s) { return s > 0; });
     if ((isStaticNewShape && inputElementsNum != newShapeElementsNum) ||
         (!isStaticNewShape && newShapeElementsNum > inputElementsNum)) {
       return emitOpError() << "cannot reshape " << inputElementsNum
                            << " elements into " << newShapeElementsNum;
     }
   }
 
-  int missingDims = llvm::count(shapeValues, -1);
+  int missingDims = llvm::count(getNewShape(), -1);
   if (missingDims > 1)
     return emitOpError() << "expected at most one target dimension to be -1";
 

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

@@ -20,6 +20,12 @@ using namespace mlir::tosa;
 
 namespace {
 
+SmallVector<int64_t> convertFromMlirShape(ArrayRef<int64_t> shape) {
+  return to_vector(llvm::map_range(shape, [](int64_t dim) {
+    return ShapedType::isDynamic(dim) ? -1 : dim;
+  }));
+}
+
 struct Conv2DIsFullyConnected : public OpRewritePattern<tosa::Conv2DOp> {
   explicit Conv2DIsFullyConnected(MLIRContext *context)
       : OpRewritePattern(context) {}
@@ -92,27 +98,25 @@ struct Conv2DIsFullyConnected : public OpRewritePattern<tosa::Conv2DOp> {
     llvm::SmallVector<int64_t, 2> revisedInputShape{combined, inputShape[3]};
     auto revisedInputShapeType =
         RankedTensorType::get(revisedInputShape, inputType.getElementType());
-    auto revisedInputShapeValue = getTosaConstShape(
-        rewriter, op.getLoc(), convertFromMlirShape(revisedInputShape));
-    auto reshapedInput =
-        rewriter
-            .create<tosa::ReshapeOp>(op.getLoc(), revisedInputShapeType, input,
-                                     revisedInputShapeValue)
-            .getResult();
+    auto reshapedInput = rewriter
+                             .create<tosa::ReshapeOp>(
+                                 op.getLoc(), revisedInputShapeType, input,
+                                 rewriter.getDenseI64ArrayAttr(
+                                     convertFromMlirShape(revisedInputShape)))
+                             .getResult();
 
     // Reshape kernel to [OC,KH,KW,IC] -> [OC, IC].
     llvm::SmallVector<int64_t, 2> revisedWeightShape{weightShape[0],
                                                      weightShape[3]};
     auto revisedWeightShapeType = RankedTensorType::get(
         revisedWeightShape,
         dyn_cast<RankedTensorType>(weight.getType()).getElementType());
-    auto revisedWeightShapeValue = getTosaConstShape(
-        rewriter, op.getLoc(), convertFromMlirShape(revisedWeightShape));
-    auto reshapedWeight =
-        rewriter
-            .create<tosa::ReshapeOp>(op.getLoc(), revisedWeightShapeType,
-                                     weight, revisedWeightShapeValue)
-            .getResult();
+    auto reshapedWeight = rewriter
+                              .create<tosa::ReshapeOp>(
+                                  op.getLoc(), revisedWeightShapeType, weight,
+                                  rewriter.getDenseI64ArrayAttr(
+                                      convertFromMlirShape(revisedWeightShape)))
+                              .getResult();
 
     // Perform a fully connected network over the reshaped input and weight.
     llvm::SmallVector<int64_t, 2> fullyConnectedShape{combined, weightShape[0]};
@@ -145,10 +149,9 @@ struct Conv2DIsFullyConnected : public OpRewritePattern<tosa::Conv2DOp> {
     // Reshape output to [N, IH, IW, OC].
     llvm::SmallVector<int64_t, 4> outputShape{inputShape[0], inputShape[1],
                                               inputShape[2], weightShape[0]};
-    auto outputShapeValue = getTosaConstShape(
-        rewriter, op.getLoc(), convertFromMlirShape(outputShape));
     rewriter.replaceOpWithNewOp<tosa::ReshapeOp>(
-        op, resultType, fullyConnectedValue, outputShapeValue);
+        op, resultType, fullyConnectedValue,
+        rewriter.getDenseI64ArrayAttr(convertFromMlirShape(outputShape)));
     return success();
   }
 };

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

@@ -55,11 +55,10 @@ struct DepthwiseConv2DIsMul : public OpRewritePattern<tosa::DepthwiseConv2DOp> {
     inputType = RankedTensorType::get(
         revisedInputShape,
         dyn_cast<RankedTensorType>(input.getType()).getElementType());
-    auto revisedInputShapeValue =
-        getTosaConstShape(rewriter, op.getLoc(), revisedInputShape);
     input = rewriter
-                .create<tosa::ReshapeOp>(op.getLoc(), inputType, input,
-                                         revisedInputShapeValue)
+                .create<tosa::ReshapeOp>(
+                    op.getLoc(), inputType, input,
+                    rewriter.getDenseI64ArrayAttr(revisedInputShape))
                 .getResult();
 
     Type inputETy = inputType.getElementType();
@@ -154,10 +153,9 @@ struct DepthwiseConv2DIsMul : public OpRewritePattern<tosa::DepthwiseConv2DOp> {
     auto outputShapeType = RankedTensorType::get(
         outputShape,
         dyn_cast<RankedTensorType>(input.getType()).getElementType());
-    auto outputShapeValue =
-        getTosaConstShape(rewriter, op->getLoc(), outputShape);
     Value outputValue = rewriter.create<tosa::ReshapeOp>(
-        op.getLoc(), outputShapeType, mulValue, outputShapeValue);
+        op.getLoc(), outputShapeType, mulValue,
+        rewriter.getDenseI64ArrayAttr(outputShape));
 
     Value bias = op.getBias();
     if (EqualizeRanks(rewriter, op.getLoc(), outputValue, bias).failed()) {

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

@@ -159,11 +159,9 @@ class TransposeConvStridedConverter
         outputChannels, weightHeight / stride[0],
         stride[0],      weightWidth / stride[1],
         stride[1],      inputChannels};
-
-    ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
     weight = CreateOpAndInferShape<tosa::ReshapeOp>(
-        builder, UnrankedTensorType::get(weightETy), weight,
-        getTosaConstShape(rewriter, loc, weightReshapeDims0));
+        rewriter, loc, UnrankedTensorType::get(weightETy), weight,
+        rewriter.getDenseI64ArrayAttr(weightReshapeDims0));
 
     // Transpose the factored-out stride to the output channels.
     Value transposeWeightVal = rewriter.create<tosa::ConstOp>(
@@ -175,13 +173,12 @@ class TransposeConvStridedConverter
         transposeWeightVal);
 
     // Collapse the strides and output channels into a single dimension.
-    llvm::SmallVector<int64_t, 4> weightReshapeDims1 = {
+    llvm::SmallVector<int64_t, 6> weightReshapeDims1 = {
         outputChannels * stride[0] * stride[1], weightHeight / stride[0],
         weightWidth / stride[1], inputChannels};
-
     weight = CreateOpAndInferShape<tosa::ReshapeOp>(
         rewriter, loc, UnrankedTensorType::get(weightETy), weight,
-        getTosaConstShape(rewriter, loc, weightReshapeDims1));
+        rewriter.getDenseI64ArrayAttr(weightReshapeDims1));
     ShapedType restridedWeightTy = cast<ShapedType>(weight.getType());
 
     weight = CreateOpAndInferShape<tosa::ReverseOp>(
@@ -260,13 +257,9 @@ class TransposeConvStridedConverter
     // Factor striding out of the convolution result.
     llvm::SmallVector<int64_t, 6> convReshapeDims0 = {
         batch, convHeight, convWidth, stride[0], stride[1], outputChannels};
-
-    auto convReshapeDims0Value =
-        getTosaConstShape(rewriter, loc, convReshapeDims0);
-
     conv2d = CreateOpAndInferShape<tosa::ReshapeOp>(
         rewriter, loc, UnrankedTensorType::get(resultETy), conv2d,
-        convReshapeDims0Value);
+        rewriter.getDenseI64ArrayAttr(convReshapeDims0));
 
     // Transpose the factored-out stride to the output channels.
     Value transposeConvVal = rewriter.create<tosa::ConstOp>(
@@ -280,13 +273,9 @@ class TransposeConvStridedConverter
     // Fuse striding behavior back into width / height.
     llvm::SmallVector<int64_t, 6> convReshapeDims1 = {
         batch, convHeight * stride[0], convWidth * stride[1], outputChannels};
-
-    auto convReshapeDims1Value =
-        getTosaConstShape(rewriter, loc, convReshapeDims1);
-
     conv2d = CreateOpAndInferShape<tosa::ReshapeOp>(
         rewriter, loc, UnrankedTensorType::get(resultETy), conv2d,
-        convReshapeDims1Value);
+        rewriter.getDenseI64ArrayAttr(convReshapeDims1));
 
     // Determine the amount to slice / pad from the result start.
     int64_t resultSliceTop = std::max<int64_t>(0, -pad[0]);

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

@@ -402,20 +402,13 @@ std::optional<Value> TosaReduceTransposes::buildMappedToValue(
     return std::nullopt;
 
   // Do not insert a TransposeOp, instead we fold the reshape and its attribute.
-  llvm::SmallVector<int64_t> newShape;
-  if (!tosa::getConstShapeValue(reshapeOp.getShape().getDefiningOp(),
-                                newShape)) {
-    // this mean shape is not constant
-    return std::nullopt;
-  }
-  ImplicitLocOpBuilder builder(reshapeOp.getLoc(), rewriter);
   auto foldedReshape = rewriter.create<ReshapeOp>(
       reshapeOp.getLoc(),
       RankedTensorType::get(applyTOSAPermutation(shape, hoistedPerms),
                             reshapeOutputType.getElementType()),
       reshapeOp.getInput1(),
-      getTosaConstShape(builder, applyTOSAPermutation(llvm::ArrayRef(newShape),
-                                                      hoistedPerms)));
+      rewriter.getDenseI64ArrayAttr(
+          applyTOSAPermutation(reshapeOp.getNewShape(), hoistedPerms)));
   return foldedReshape->getResult(0);
 }