[tosa][mlir] Add handling of more than one dynamic dimension for reshape lowering

[NatashaKnk]

Summary

Tosa-to-linalg lowering for the reshape function. Works for simple cases where the dynamic dimensions are not modified.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

static LogicalResult reduceMatchAndRewriteHelper(Operation *op, uint64_t axis,
   rewriter.replaceOpWithNewOp<tosa::ReshapeOp>(op, resultTy,
                                                linalgOp.getResults());
   return success();
 }
 
 static bool findIntermediateShape(ArrayRef<int64_t> lhsShape,
                                   ArrayRef<int64_t> rhsShape,
                                   SmallVector<int64_t> &intermediateShape,
-                                  bool isDynamic) {
-  if (isDynamic) {
+                                  int64_t dynamicDims) {
+  if (dynamicDims == 1) {
     // TODO (natashaknk): Make dynamic intermediate shape not always be rank-1
     intermediateShape = {-1};
     return true;
   }
 
   if (lhsShape.empty() || rhsShape.empty()) {
     intermediateShape = {};
     return true;
   }
 
   unsigned currLhsDim = 0, currRhsDim = 0;
   while (currLhsDim < lhsShape.size() && currRhsDim < rhsShape.size()) {
     int64_t rhsSize = rhsShape[currRhsDim];
     int64_t lhsSize = lhsShape[currLhsDim];
     while (lhsSize != rhsSize && currLhsDim < lhsShape.size() &&
            currRhsDim < rhsShape.size()) {
+      if ((lhsSize == -1) != (rhsSize == -1))
+        return false;
+      if (lhsSize == -1 && rhsSize == -1) {
+        intermediateShape.push_back(lhsSize);
+        currLhsDim++;
+        currRhsDim++;
+        continue;
+      }
       if (lhsSize < rhsSize) {
         currLhsDim++;
         lhsSize *= lhsShape[currLhsDim];
       } else {
         currRhsDim++;
         rhsSize *= rhsShape[currRhsDim];
       }
     }
static bool findIntermediateShape(ArrayRef<int64_t> lhsShape,
   }
 
   return true;
 }
 
 static bool createReassociationMapsForCollapse(
     PatternRewriter &rewriter, ArrayRef<int64_t> srcShape,
     ArrayRef<int64_t> dstShape,
-    SmallVector<ReassociationExprs, 4> &reassociationMap, bool isDynamic) {
+    SmallVector<ReassociationExprs, 4> &reassociationMap, int64_t dynamicDims) {
+
+  if (dstShape.empty()) {
+    reassociationMap = {};
+    return true;
+  }
 
-  // If the shape is dynamic, create a map for collapsing into one dimension.
-  if (isDynamic) {
+  // If the shape is dynamic with one dynamic dim, create a map for collapsing
+  // into one dimension.
+  if (dynamicDims == 1) {
     SmallVector<AffineExpr, 2> exprs;
     for (int i = 0, s = srcShape.size(); i < s; ++i)
       exprs.push_back(rewriter.getAffineDimExpr(i));
     reassociationMap = {exprs};
     return true;
   }
 
-  if (dstShape.empty()) {
-    reassociationMap = {};
-    return true;
-  }
-
+  // If there are multiple dynamic dims, only reshape if all the other dims
+  // match up
   reassociationMap.resize(dstShape.size());
+
   unsigned currSrcDim = 0, currDstDim = 0;
   while (currSrcDim < srcShape.size() && currDstDim < dstShape.size()) {
     int64_t dstSize = dstShape[currDstDim];
     int64_t srcSize = srcShape[currSrcDim];
+    if ((srcSize == -1) != (dstSize == -1))
+      return false;
+
+    if (srcSize == -1) {
+      reassociationMap[currDstDim++].push_back(
+          rewriter.getAffineDimExpr(currSrcDim++));
+      if (currSrcDim)
+        continue;
+    }
+
     while (srcSize < dstSize && currSrcDim < srcShape.size()) {
       reassociationMap[currDstDim].push_back(
           rewriter.getAffineDimExpr(currSrcDim++));
       srcSize *= srcShape[currSrcDim];
     }
-    if (srcSize == dstSize) {
+
+    if (srcSize == dstSize && srcSize != -1) {
       reassociationMap[currDstDim].push_back(
           rewriter.getAffineDimExpr(currSrcDim++));
-      // If the next dim in collapsedShape is not 1, treat subsequent dims in
-      // expandedShape which are 1 to be collapsed.
+      // If the next dim in collapsedShape is not 1, treat subsequent dims
+      // in expandedShape which are 1 to be collapsed.
       if (currDstDim == dstShape.size() - 1 || dstShape[currDstDim + 1] != 1) {
         while (currSrcDim < srcShape.size() && srcShape[currSrcDim] == 1) {
           reassociationMap[currDstDim].push_back(
               rewriter.getAffineDimExpr(currSrcDim++));
         }
       }
     }
     currDstDim++;

 public:
   using OpConversionPattern<tosa::ReshapeOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(tosa::ReshapeOp reshape, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
     ShapedType operandTy = adaptor.input1().getType().cast<ShapedType>();
     ShapedType resultTy = reshape.getType().template cast<ShapedType>();
-    bool isDynamic = !operandTy.hasStaticShape();
-
-    if (isDynamic && resultTy.getRank() != 1) {
-      return rewriter.notifyMatchFailure(
-          reshape, "Cannot collapse dynamic dims to more than one dimension");
-    }
+    int64_t dynamicDims = 0;
 
     if (operandTy == resultTy) {
       rewriter.replaceOp(reshape, adaptor.getOperands()[0]);
       return success();
     }
 
+    if (operandTy.getNumDynamicDims() == resultTy.getNumDynamicDims()) {
+      dynamicDims = operandTy.getNumDynamicDims();
+    } else {
+      return rewriter.notifyMatchFailure(
+          reshape, "tosa.reshape must have the same number of dynamic "
+                   "dimentions for input and output");
+    }
+
     SmallVector<ReassociationExprs, 4> reassociationMap;
     if (!createReassociationMapsForCollapse(rewriter, operandTy.getShape(),
                                             resultTy.getShape(),
-                                            reassociationMap, isDynamic)) {
+                                            reassociationMap, dynamicDims)) {
       return rewriter.notifyMatchFailure(
           reshape,
           "tosa.reshape Attempting to collapse into an incompatible shape");
     }
 
     SmallVector<int64_t> intermediateShape;
     if (!findIntermediateShape(operandTy.getShape(), resultTy.getShape(),
-                               intermediateShape, isDynamic)) {
+                               intermediateShape, dynamicDims)) {
       return rewriter.notifyMatchFailure(
           reshape, "tosa.reshape Cannot collapse into given shape");
     }
 
     rewriter.replaceOpWithNewOp<tensor::CollapseShapeOp>(
         reshape, resultTy, adaptor.getOperands()[0], reassociationMap);
     return success();
   }
 };
 
 class ReshapeConverterExpand : public OpConversionPattern<tosa::ReshapeOp> {
 public:
   using OpConversionPattern<tosa::ReshapeOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(tosa::ReshapeOp reshape, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
     ShapedType operandTy = adaptor.input1().getType().cast<ShapedType>();
     ShapedType resultTy = reshape.getType().template cast<ShapedType>();
-    bool isDynamic = !operandTy.hasStaticShape();
+    int64_t dynamicDims = 0;
 
     if (operandTy == resultTy) {
       rewriter.replaceOp(reshape, adaptor.getOperands()[0]);
       return success();
     }
 
-    if (isDynamic && operandTy.getRank() != 1) {
+    if (operandTy.getNumDynamicDims() == resultTy.getNumDynamicDims()) {
+      dynamicDims = operandTy.getNumDynamicDims();
+    } else {
       return rewriter.notifyMatchFailure(
-          reshape, "Cannot expand dynamic dims from more than one dimension");
+          reshape, "tosa.reshape must have the same number of dynamic "
+                   "dimentions for input and output");
     }
 
     SmallVector<ReassociationExprs, 4> reassociationMap;
     if (!createReassociationMapsForCollapse(rewriter, resultTy.getShape(),
                                             operandTy.getShape(),
-                                            reassociationMap, isDynamic)) {
+                                            reassociationMap, dynamicDims)) {
       return rewriter.notifyMatchFailure(
           reshape,
           "tosa.reshape Attempting to expand into an incompatible shape");
     }
 
     SmallVector<int64_t> intermediateShape;
     if (!findIntermediateShape(operandTy.getShape(), resultTy.getShape(),
-                               intermediateShape, isDynamic) ||
+                               intermediateShape, dynamicDims) ||
         intermediateShape != operandTy.getShape()) {
       return rewriter.notifyMatchFailure(
           reshape, "tosa.reshape Cannot expand into given shape");
     }
     rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
         reshape, resultTy, adaptor.getOperands()[0], reassociationMap);
     return success();
   }
 };
 
 class ReshapeConverterCollapseExpand
     : public OpConversionPattern<tosa::ReshapeOp> {
 public:
   using OpConversionPattern<tosa::ReshapeOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(tosa::ReshapeOp reshape, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
     ShapedType operandTy = adaptor.input1().getType().cast<ShapedType>();
     ShapedType resultTy = reshape.getType().template cast<ShapedType>();
-    bool isDynamic = !operandTy.hasStaticShape();
+    int64_t dynamicDims = 0;
 
     if (operandTy == resultTy) {
       rewriter.replaceOp(reshape, adaptor.getOperands()[0]);
       return success();
     }
 
+    if (operandTy.getNumDynamicDims() == resultTy.getNumDynamicDims()) {
+      dynamicDims = operandTy.getNumDynamicDims();
+    } else {
+      return rewriter.notifyMatchFailure(
+          reshape, "tosa.reshape must have the same number of dynamic "
+                   "dimentions for input and output");
+    }
+
     SmallVector<int64_t> intermediateShape;
     if (!findIntermediateShape(resultTy.getShape(), operandTy.getShape(),
-                               intermediateShape, isDynamic)) {
+                               intermediateShape, dynamicDims)) {
       return rewriter.notifyMatchFailure(
-          reshape, "tosa.reshape Cannot identify an intermediate shape between "
-                   "the given two shapes");
+          reshape, "tosa.reshape Cannot identify an intermediate shape"
+                   "between the given two shapes");
     }
 
     Value collapse = rewriter.create<tosa::ReshapeOp>(
         reshape.getLoc(),
         RankedTensorType::get(intermediateShape,
                               reshape.getType().getElementType()),
         adaptor.input1());
     Value expand =

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

func.func @test_reshape_downrank_6D_dyn(%arg0: tensor<1x2x?x5x7x11xf32>) -> tensor<?x5x77xf32> {
   // CHECK: tensor.collapse_shape %arg0 {{\[}}[0, 1, 2, 3, 4, 5]]
   // CHECK: tensor.expand_shape %0 {{\[}}[0, 1, 2]]
   %0 = "tosa.reshape"(%arg0) {new_shape = [-1, 5, 77]} : (tensor<1x2x?x5x7x11xf32>) -> tensor<?x5x77xf32>
   return %0 : tensor<?x5x77xf32>
 }
 
 // -----
 
+//CHECK-LABEL: @test_reshape_collapse_multiple_dyn_dims
+func.func @test_reshape_collapse_multiple_dyn_dims(%arg0: tensor<1x?x?x256x1xf32>) -> (tensor<1x?x?x256xf32>) {
+  // CHECK: tensor.collapse_shape %arg0 {{\[}}[0], [1], [2], [3, 4]]
+  %0 = "tosa.reshape"(%arg0) {new_shape = [1, -1, -1, 256]} : (tensor<1x?x?x256x1xf32>) -> tensor<1x?x?x256xf32>
+  return %0 : tensor<1x?x?x256xf32>
+}
+
+// -----
+
+//CHECK-LABEL: @test_reshape_expand_multiple_dyn_dims
+func.func @test_reshape_expand_multiple_dyn_dims(%arg0: tensor<?x?x?x6x8xf32>) -> (tensor<?x?x?x2x3x2x4xf32>) {
+  // CHECK: tensor.expand_shape %arg0 {{\[}}[0], [1], [2], [3, 4], [5, 6]]
+  %0 = "tosa.reshape"(%arg0) {new_shape = [-1, -1, -1, 2, 3, 2, 4]} : (tensor<?x?x?x6x8xf32>) -> tensor<?x?x?x2x3x2x4xf32>
+  return %0 : tensor<?x?x?x2x3x2x4xf32>
+}
+
+
+// -----
+
+//CHECK-LABEL: @test_reshape_collapse_expand_multiple_dyn_dims
+func.func @test_reshape_collapse_expand_multiple_dyn_dims(%arg0: tensor<2x3x8x?x?xf32>) -> (tensor<6x2x4x?x?xf32>) {
+  // CHECK: tensor.collapse_shape %arg0 {{\[}}[0, 1], [2], [3], [4]]
+  // CHECK: tensor.expand_shape %0 {{\[}}[0], [1, 2], [3], [4]]
+  %0 = "tosa.reshape"(%arg0) {new_shape = [6, 2, 4, -1, -1]} : (tensor<2x3x8x?x?xf32>) -> tensor<6x2x4x?x?xf32>
+  return %0 : tensor<6x2x4x?x?xf32>
+}
+
+// -----
+
 // CHECK-LABEL: @test_identity
 func.func @test_identity(%arg0: tensor<1xf32>, %arg1: tensor<1xi32>) -> (tensor<1xf32>, tensor<1xi32>) {
   %0 = "tosa.identity"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
   %1 = "tosa.identity"(%arg1) : (tensor<1xi32>) -> tensor<1xi32>
 
   // CHECK: return %arg0, %arg1
   return %0, %1 : tensor<1xf32>, tensor<1xi32>
 }
func.func @pad_float(%arg0 : tensor<1x2xf32>) -> (tensor<4x9xf32>) {
   %0 = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
   // TODO: Output contains multiple "arith.constant 1 : index".
   // CHECK-DAG: [[INDEX1:%.+]] = arith.constant 1 : index
   // CHECK-DAG: [[INDEX2:%.+]] = arith.constant 2 : index
   // CHECK-DAG: [[INDEX3:%.+]] = arith.constant 3 : index
   // CHECK-DAG: [[INDEX4:%.+]] = arith.constant 4 : index
   // CHECK-DAG: [[CST:%.+]] = arith.constant 0.000000e+00 : f32
   // CHECK: tensor.pad %arg0 low{{\[}}%{{.*}}, [[INDEX3]]] high{{\[}}[[INDEX2]], [[INDEX4]]]  {
-  // CHECK: ^bb0(%arg1: index, %arg2: index):  
+  // CHECK: ^bb0(%arg1: index, %arg2: index):
   // CHECK:   tensor.yield [[CST]]
   // CHECK: } : tensor<1x2xf32> to tensor<4x9xf32>
   %1 = "tosa.pad"(%arg0, %0)  : (tensor<1x2xf32>, tensor<2x2xi32>)  -> (tensor<4x9xf32>)
   return %1 : tensor<4x9xf32>
 }
 
 func.func @pad_int(%arg0 : tensor<1x2xi32>) -> (tensor<4x9xi32>) {
   %0 = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
func.func @pad_float_explicit(%arg0 : tensor<1x2xf32>) -> (tensor<4x9xf32>) {
   %0 = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
   // TODO: Output contains multiple "arith.constant 1 : index".
   // CHECK-DAG: [[INDEX1:%.+]] = arith.constant 1 : index
   // CHECK-DAG: [[INDEX2:%.+]] = arith.constant 2 : index
   // CHECK-DAG: [[INDEX3:%.+]] = arith.constant 3 : index
   // CHECK-DAG: [[INDEX4:%.+]] = arith.constant 4 : index
   // CHECK-DAG: [[CST:%.+]] = arith.constant 4.200000e+01 : f32
   // CHECK: tensor.pad %arg0 low{{\[}}%{{.*}}, [[INDEX3]]] high{{\[}}[[INDEX2]], [[INDEX4]]]  {
-  // CHECK: ^bb0(%arg1: index, %arg2: index):  
+  // CHECK: ^bb0(%arg1: index, %arg2: index):
   // CHECK:   tensor.yield [[CST]]
   // CHECK: } : tensor<1x2xf32> to tensor<4x9xf32>
   %1 = arith.constant dense<42.0> : tensor<f32>
   %2 = "tosa.pad"(%arg0, %0, %1)  : (tensor<1x2xf32>, tensor<2x2xi32>, tensor<f32>)  -> (tensor<4x9xf32>)
   return %2 : tensor<4x9xf32>
 }
 
 // -----