[tosa][mlir] Add dynamic width/height support for depthwise convolution in tosa-to-linalg

[NatashaKnk]

Summary

In addition, fixed a small bug with padding incorrectly inferring output shape for dynaic inputs in convolution

[rsuderman]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

51–54

You should be able to use ShapedType::kDynamicSize instead of -1. Makes things clearer what -1 means.

110

Rather than putting width / height int he loop, just copy the array then special case the width / height after. It limits the loop body but still handles the special cases.

[jpienaar]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

51–54

Or ShapedType::isDynamic(inputShape[i])

106

And here we will always have a rank?

111

Consider continuing if not dynamic to reduce the nesting here.

113

stride_y ?

[NatashaKnk]

Add small clarity fixes

[NatashaKnk]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

51–54

Is there a preferred approach between ShapedType::isDynamic(inputShape[i]) and something like inputTy.isDynamicDim(i)?

106

I'm assuming that running convolution on unranked input should error, correct? If that's the case, should I check for it in the outside function as a sanity check?

111

I'm guessing it won't be necessary if I follow Robs advice above and take the special cases out of the loop?

113

I'm not sure I'm understanding the question, but stride_y is how the tosa spec refers to the height parameter for stride.

[jpienaar]

Nice, thanks

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

106

You are correct in the verifier of conv we check if ranked.

113

Oh I just meant there is a typo sride instead of stride

137

dynDims[weightDim]?

183

kwarg style comments for 1 and 2 ? (that way we can clang-tidy check too)

[NatashaKnk]

Small additional cleanup

[NatashaKnk]

Fix typos in documentation

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

static mlir::Value applyPad(Location loc, Value input, ArrayRef<int64_t> pad,
   assert((inputShape.size() * 2) == pad.size());
 
   SmallVector<int64_t, 4> paddedShape;
   SmallVector<OpFoldResult, 8> lowIndices;
   SmallVector<OpFoldResult, 8> highIndices;
   for (int i = 0, s = inputShape.size(); i < s; i++) {
     auto lowPad = pad[i * 2];
     auto highPad = pad[i * 2 + 1];
+    if (ShapedType::isDynamic(inputShape[i]))
+      paddedShape.push_back(inputShape[i]);
+    else
       paddedShape.push_back(inputShape[i] + highPad + lowPad);
     lowIndices.push_back(rewriter.getIndexAttr(lowPad));
     highIndices.push_back(rewriter.getIndexAttr(highPad));
   }
 
   Value padValue = rewriter.create<arith::ConstantOp>(loc, padAttr);
 
   return tensor::createPadScalarOp(RankedTensorType::get(paddedShape, inputETy),
                                    input, padValue, lowIndices, highIndices,
                                    /*nofold=*/false, loc, rewriter)
       .result();
 }
 
 static mlir::Value reifyConstantDim(Attribute attr,
                                     ImplicitLocOpBuilder &builder) {
   return builder.createOrFold<arith::IndexCastOp>(
       builder.getIndexType(), builder.create<arith::ConstantOp>(attr));
 }
 
 // Calculating the output width/height using the formula:
-// Out =((initDim+padBefore+padAttr-(dilation*(kernelDim-1)+1))/stride+1
 // H = ((IH+pad_top+pad_bottom-(dilation_y*(KH-1)+1))/stride_y)+1
 // W = ((IW+pad_left+pad_right-(dilation_x*(KW-1)+1))/stride_x)+1
 static mlir::Value
 getConvOutputDim(Location loc, Value initDim, Attribute padBeforeAttr,
                  Attribute padAfterAttr, Value kernelDim, Attribute strideAttr,
                  Attribute dilationAttr, Type inputETy, OpBuilder &rewriter) {
   ImplicitLocOpBuilder builder(loc, rewriter);
   auto one = rewriter.create<arith::ConstantOp>(
getConvOutputDim(Location loc, Value initDim, Attribute padBeforeAttr,
   Value addOne = builder.create<arith::AddIOp>(dilated, one);
 
   Value subtract = builder.create<arith::SubIOp>(paddedAfter, addOne);
   Value stride = reifyConstantDim(strideAttr, builder);
   Value divide = builder.create<arith::DivUIOp>(subtract, stride);
   return builder.create<arith::SubIOp>(divide, one);
 }
 
+// Creates a vector of the dynamic output dims for Conv2D and Depthwise_Conv2D
+static SmallVector<Value> inferDynamicDimsForConv(
+    Location loc, Value input, Value weight, ShapedType resultTy,
+    ArrayAttr padAttr, ArrayAttr strideAttr, ArrayAttr dilationAttr,
+    int64_t weightHDim, int64_t weightWDim, OpBuilder &rewriter) {
+  ShapedType inputTy = input.getType().cast<ShapedType>();
+  Type inputETy = inputTy.getElementType();
+  int64_t inputRank = inputTy.getRank();
+  int64_t heightDim = 1;
+  int64_t weightDim = 2;
+
+  SmallVector<Value> dynDims;
+  dynDims.resize(resultTy.getRank());
+  for (int i = 0; i < inputRank; i++) {
+    if (inputTy.isDynamicDim(i) && i != heightDim && i != weightDim)
+      dynDims[i] = rewriter.create<tensor::DimOp>(loc, input, i);
+  }
+
+  // Dynamic input height
+  if (inputTy.isDynamicDim(heightDim)) {
+    Value initHDim =
+        rewriter.create<tensor::DimOp>(loc, input, heightDim).getResult();
+    Value kernelHDim =
+        rewriter.create<tensor::DimOp>(loc, weight, weightHDim).getResult();
+    // H = F(IH, pad_top, pad_bottom, dilation_y, KH, stride_y)
+    dynDims[heightDim] = getConvOutputDim(
+        loc, initHDim, padAttr.getValue()[0], padAttr.getValue()[1], kernelHDim,
+        strideAttr.getValue()[0], dilationAttr.getValue()[0], inputETy,
+        rewriter);
+  }
+
+  // Dynamic input weight
+  if (inputTy.isDynamicDim(weightDim)) {
+    Value initWDim =
+        rewriter.create<tensor::DimOp>(loc, input, weightDim).getResult();
+    Value kernelWDim =
+        rewriter.create<tensor::DimOp>(loc, weight, weightWDim).getResult();
+    // W = F(IW, pad_left, pad_right, dilation_x, KW, stride_x)
+    dynDims[weightDim] = getConvOutputDim(
+        loc, initWDim, padAttr.getValue()[2], padAttr.getValue()[3], kernelWDim,
+        strideAttr.getValue()[1], dilationAttr.getValue()[1], inputETy,
+        rewriter);
+  }
+
+  SmallVector<Value> filteredDims = condenseValues(dynDims);
+  return filteredDims;
+}
+
 namespace {
 
 class ConvConverter : public OpConversionPattern<tosa::Conv2DOp> {
 public:
   using OpConversionPattern<tosa::Conv2DOp>::OpConversionPattern;
   LogicalResult
   matchAndRewrite(tosa::Conv2DOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
     Location loc = op->getLoc();
     Value input = op->getOperand(0);
     Value weight = op->getOperand(1);
     Value bias = op->getOperand(2);
 
     ShapedType inputTy = input.getType().cast<ShapedType>();
     ShapedType weightTy = weight.getType().cast<ShapedType>();
     ShapedType biasTy = bias.getType().cast<ShapedType>();
     ShapedType resultTy = op->getResult(0).getType().cast<ShapedType>();
-    int64_t inputRank = inputTy.getRank();
 
     Type inputETy = inputTy.getElementType();
     Type resultETy = resultTy.getElementType();
 
     auto padAttr = op->getAttr("pad").cast<ArrayAttr>();
     auto strideTosaAttr = op->getAttr("stride").cast<ArrayAttr>();
     auto dilationTosaAttr = op->getAttr("dilation").cast<ArrayAttr>();
     bool isQuantized = op->hasAttr("quantization_info");
 
     if (!weightTy.hasStaticShape() || !biasTy.hasStaticShape())
       return rewriter.notifyMatchFailure(
           op, "tosa.conv ops require static shapes for weight and bias");
 
     if (inputETy.isUnsignedInteger())
       return rewriter.notifyMatchFailure(
           op, "tosa.conv ops does not support unsigned integer input");
 
-    SmallVector<Value> dynDims;
-    dynDims.resize(resultTy.getRank());
-    for (int i = 0; i < inputRank; i++) {
-      if (inputTy.isDynamicDim(i)) {
-        // Dynamic input height
-        // H = F(IH, pad_top, pad_bottom, dilation_y, KH, sride_y)
-        if (i == 1) {
-          Value initHDim =
-              rewriter.create<tensor::DimOp>(loc, input, 1).getResult();
-          Value kernelHDim =
-              rewriter.create<tensor::DimOp>(loc, weight, 1).getResult();
-          dynDims[i] = getConvOutputDim(
-              loc, initHDim, padAttr.getValue()[0], padAttr.getValue()[1],
-              kernelHDim, strideTosaAttr.getValue()[0],
-              dilationTosaAttr.getValue()[0], inputETy, rewriter);
-
-          // Dynamic input weight
-          // W = F(IH, pad_left, pad_right, dilation_x, KW, sride_x)
-        } else if (i == 2) {
-          Value initWDim =
-              rewriter.create<tensor::DimOp>(loc, input, 2).getResult();
-          Value kernelWDim =
-              rewriter.create<tensor::DimOp>(loc, weight, 2).getResult();
-          dynDims[i] = getConvOutputDim(
-              loc, initWDim, padAttr.getValue()[2], padAttr.getValue()[3],
-              kernelWDim, strideTosaAttr.getValue()[1],
-              dilationTosaAttr.getValue()[1], inputETy, rewriter);
-
-        } else {
-          dynDims[i] = rewriter.create<tensor::DimOp>(loc, input, i);
-        }
-      }
-    }
-
-    SmallVector<Value> filteredDims = condenseValues(dynDims);
+    SmallVector<Value> filteredDims = inferDynamicDimsForConv(
+        loc, input, weight, resultTy, padAttr, strideTosaAttr, dilationTosaAttr,
+        /*weightHDim=*/1, /*weightWDim=*/2, rewriter);
 
     auto weightShape = weightTy.getShape();
 
     // Apply padding as necessary.
     Attribute zeroAttr = rewriter.getZeroAttr(inputETy);
     if (isQuantized) {
       auto quantizationInfo =
           op->getAttr("quantization_info").cast<tosa::ConvOpQuantizationAttr>();
matchAndRewrite(tosa::DepthwiseConv2DOp op, OpAdaptor adaptor,
 
     Type inputETy = inputTy.getElementType();
     Type resultETy = resultTy.getElementType();
 
     auto padAttr = op->getAttr("pad").cast<ArrayAttr>();
     auto strideTosaAttr = op->getAttr("stride").cast<ArrayAttr>();
     auto dilationTosaAttr = op->getAttr("dilation").cast<ArrayAttr>();
 
+    if (!weightTy.hasStaticShape() || !biasTy.hasStaticShape())
+      return rewriter.notifyMatchFailure(
+          op, "tosa.depthwise_conv ops require static shapes");
+
+    // Compute output dynamic dims
+    SmallVector<Value> filteredDims = inferDynamicDimsForConv(
+        loc, input, weight, resultTy, padAttr, strideTosaAttr, dilationTosaAttr,
+        0, 1, rewriter);
+
     bool isQuantized = op->hasAttr("quantization_info");
     IntegerAttr iZp;
     IntegerAttr kZp;
     if (isQuantized) {
       auto quantizationInfo =
           op->getAttr("quantization_info").cast<tosa::ConvOpQuantizationAttr>();
       iZp = rewriter.getI32IntegerAttr(
           quantizationInfo.input_zp().getValue().getSExtValue());
       kZp = rewriter.getI32IntegerAttr(
           quantizationInfo.weight_zp().getValue().getSExtValue());
     }
 
-    if (!weightTy.hasStaticShape() || !biasTy.hasStaticShape())
-      return rewriter.notifyMatchFailure(
-          op, "tosa.depthwise_conv ops require static shapes");
-
-    auto dynamicDimsOr =
-        checkHasDynamicBatchDims(rewriter, op, {input, op.output()});
-    if (!dynamicDimsOr.hasValue())
-      return failure();
-    SmallVector<Value> dynamicDims = dynamicDimsOr.getValue();
-
     auto weightShape = weightTy.getShape();
     auto resultShape = resultTy.getShape();
 
     // Apply padding as necessary.
     Attribute zeroAttr = rewriter.getZeroAttr(inputETy);
     if (isQuantized) {
       auto quantizationInfo =
           op->getAttr("quantization_info").cast<tosa::ConvOpQuantizationAttr>();
matchAndRewrite(tosa::DepthwiseConv2DOp op, OpAdaptor adaptor,
     indexingMaps.push_back(AffineMap::get(
         /*dimCount=*/resultTy.getRank(), /*symbolCount=*/0,
         {rewriter.getAffineDimExpr(3)}, rewriter.getContext()));
     indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultTy.getRank()));
     indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultTy.getRank()));
 
     Attribute resultZeroAttr = rewriter.getZeroAttr(resultETy);
     Value initTensor = rewriter.create<linalg::InitTensorOp>(
-        loc, dynamicDims, linalgConvTy.getShape(), resultETy);
+        loc, filteredDims, linalgConvTy.getShape(), resultETy);
     Value zero = rewriter.create<arith::ConstantOp>(loc, resultZeroAttr);
     Value zeroTensor = rewriter
                            .create<linalg::FillOp>(loc, ValueRange{zero},
                                                    ValueRange{initTensor})
                            .result();
 
     Value biasInitTensor = rewriter.create<linalg::InitTensorOp>(
-        loc, dynamicDims, resultTy.getShape(), resultETy);
+        loc, filteredDims, resultTy.getShape(), resultETy);
     if (!isQuantized) {
       Value conv = rewriter
                        .create<linalg::DepthwiseConv2DNhwcHwcmOp>(
                            loc, linalgConvTy, ValueRange{input, weight},
                            ValueRange{zeroTensor}, strideAttr, dilationAttr)
                        .getResult(0);
       Value convReshape = rewriter.create<tosa::ReshapeOp>(
           loc, resultTy, conv, rewriter.getI64ArrayAttr(resultTy.getShape()));

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

 func @depthwise_conv(%arg0 : tensor<1x7x5x3xf32>, %arg1 : tensor<3x1x3x11xf32>, %arg2 : tensor<33xf32>) -> () {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 5, 5, 3, 11]
   // CHECK: [[CST0:%.+]] = arith.constant 0
   // CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]]
   // CHECK: [[OUT:%.+]] = linalg.init_tensor [1, 5, 5, 33]
   // CHECK: [[DEPTH:%.+]] = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<1x7x5x3xf32>, tensor<3x1x3x11xf32>) outs([[FILL]] : tensor<1x5x5x3x11xf32>)
   // CHECK: [[COLLAPSED:%.+]] = "tosa.reshape"([[DEPTH]]) {new_shape = [1, 5, 5, 33]}
   // CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<33xf32>, tensor<1x5x5x33xf32>) outs([[OUT]] : tensor<1x5x5x33xf32>) {
-  // CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):  
+  // CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
   // CHECK:   [[ADD:%.+]] = arith.addf %arg3, %arg4 : f32
   // CHECK:   linalg.yield [[ADD]] : f32
   // CHECK: } -> tensor<1x5x5x33xf32>
   %2 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) { pad = [0, 0, 0, 0], stride = [1, 1], dilation = [1, 1] } : (tensor<1x7x5x3xf32>, tensor<3x1x3x11xf32>, tensor<33xf32>)  -> (tensor<1x5x5x33xf32>)
   return
 }
 
 // -----
 
 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d3)>
 // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 
 // CHECK-LABEL: @depthwise_conv_dyn
 func @depthwise_conv_dyn(%arg0 : tensor<?x7x5x3xf32>, %arg1 : tensor<3x1x3x11xf32>, %arg2 : tensor<33xf32>) -> () {
   // CHECK: %[[C0:.+]] = arith.constant 0
   // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]]
   // CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[BATCH]], 5, 5, 3, 11]
   // CHECK: %[[CST0:.+]] = arith.constant 0
   // CHECK: %[[FILL:.+]] = linalg.fill
   // CHECK: %[[OUT:.+]] = linalg.init_tensor [%[[BATCH]], 5, 5, 33]
   // CHECK: %[[DEPTH:.+]] = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<?x7x5x3xf32>, tensor<3x1x3x11xf32>) outs(%[[FILL]] : tensor<?x5x5x3x11xf32>)
   // CHECK: %[[COLLAPSED:.+]] = "tosa.reshape"(%[[DEPTH]]) {new_shape = [-1, 5, 5, 33]}
   // CHECK: %[[BIAS:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, %[[COLLAPSED]] : tensor<33xf32>, tensor<?x5x5x33xf32>) outs(%[[OUT]] : tensor<?x5x5x33xf32>) {
-  // CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):  
+  // CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
   // CHECK:   %[[ADD:.+]] = arith.addf %arg3, %arg4 : f32
   // CHECK:   linalg.yield %[[ADD]] : f32
   // CHECK: } -> tensor<?x5x5x33xf32>
   %2 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) { pad = [0, 0, 0, 0], stride = [1, 1], dilation = [1, 1] } : (tensor<?x7x5x3xf32>, tensor<3x1x3x11xf32>, tensor<33xf32>)  -> (tensor<?x5x5x33xf32>)
   return
 }
 
 // -----
func @depthwise_conv_quant_dilations(%arg0 : tensor<1x14x14x4xi8>, %arg1 : tensor<3x3x4x128xi8>, %arg2 : tensor<512xi32>) -> () {
   // CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<512xi32>, tensor<1x10x10x512xi32>) outs([[OUT]] : tensor<1x10x10x512xi32>) {
   // CHECK: ^bb0(%arg3: i32, %arg4: i32, %arg5: i32):  
   // CHECK:   [[ADD:%.+]] = arith.addi %arg3, %arg4 : i32
   // CHECK:   linalg.yield [[ADD]] : i32
   // CHECK: } -> tensor<1x10x10x512xi32>
   %0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [0, 0, 0, 0], quantization_info = {input_zp = -128 : i32, weight_zp = 42 : i32}, stride = [1, 1], dilation = [2, 2] } : (tensor<1x14x14x4xi8>, tensor<3x3x4x128xi8>, tensor<512xi32>)  -> tensor<1x10x10x512xi32>
   return
 }
+
+// CHECK-LABEL: @depthwise_conv2d_dyn_w_h
+func @depthwise_conv2d_dyn_w_h(%arg0: tensor<2x?x?x3xf32>, %arg1: tensor<3x6x3x5xf32>, %arg2: tensor<15xf32>) {
+  // CHECK: arith.addi
+  // CHECK: arith.subi
+  // CHECK: arith.muli
+  // CHECK: arith.divui
+  // CHECK: %[[PADDED:.+]] = tensor.pad %arg0 low[0, 1, 3, 0] high[0, 2, 4, 0] {
+  // CHECK: ^bb0(%arg3: index, %arg4: index, %arg5: index, %arg6: index):
+  // CHECK: tensor.yield %cst : f32
+  // CHECK:  } : tensor<2x?x?x3xf32> to tensor<2x?x?x3xf32>
+  // CHECK: %[[CONV:.+]] = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<[2, 1]> : tensor<2xi64>, strides = dense<[1, 2]> : tensor<2xi64>} ins(%[[PADDED]], %arg1 : tensor<2x?x?x3xf32>, tensor<3x6x3x5xf32>) outs(%22 : tensor<2x?x?x3x5xf32>) -> tensor<2x?x?x3x5xf32>
+  // CHECK: %[[RESHAPED:.+]] = "tosa.reshape"(%[[CONV]]) {new_shape = [2, -1, -1, 15]} : (tensor<2x?x?x3x5xf32>) -> tensor<2x?x?x15xf32>
+  %0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [1, 2, 3, 4], dilation = [2, 1], stride = [1, 2]} : (tensor<2x?x?x3xf32>, tensor<3x6x3x5xf32>, tensor<15xf32>) -> tensor<2x?x?x15xf32>
+  return
+}