[tosa][mlir] Add support for dynamic width/height for Conv2D inputs in tosa-to-linalg

[NatashaKnk]

Summary

Infers output shape for dynamic width/height inputs.

[jpienaar]

Nice, thanks!

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

64

As is currently named the function may be confusing as the name is quite general (index refers to the type here, but many could think of index into tensor/index in axis sense). If we follow the other shape work, reifyConstantDim could work (else materialize would feel more natural locally, but the method to materialize the shape computations as actual ops in the module is using reify).

66

This is adding a cast unconditionally, one could use createOrFold (https://mlir.llvm.org/doxygen/classmlir_1_1ImplicitLocOpBuilder.html#aec0ceee6a6e834449c1f7c65d45b26e5) or just check the type before creating instead. (That also allows to assert that the type is compatible, so if you got a float attr instead of an integer one, you could flag it)

73

Could you rewrite the formula in terms of what the inputs are below? (And then like H=F(IH, ...)) If one just reads the comment one could think that this returns 2 values, one for the w and one for the h, but instead you have noticed it is one formula needed that enables calculating either by passing in different attributes needed.

134

Nit: prefer to not compute end in for (https://llvm.org/docs/CodingStandards.html#use-range-based-for-loops-wherever-possible)

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

444

So (OOC) the reification of shape computations are currently combined with lowering to linalg? (There isn't a an intermediate step)

[rsuderman]

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

444

Thats my thought. It would be nice to have a better computation mechanism for reify-ing shapes but we explicitly do it during the lowering. I was wondering whether a refactor in the future to push it to a reification would help simplify things.

[jpienaar]

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

444

I think we can get there with InferTensorTypeWithReify + perhaps some cleanups in its reify mechanism (ValueShapeRange is intended to allow injecting info without mutating). Such a refactor would be great!

[NatashaKnk]

Small clarity updates

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

static mlir::Value applyPad(Location loc, Value input, ArrayRef<int64_t> pad,
   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>(
+      loc, IntegerAttr::get(initDim.getType(), 1));
+  Value padBefore = reifyConstantDim(padBeforeAttr, builder);
+  Value paddedBefore = builder.create<arith::AddIOp>(initDim, padBefore);
+  Value padAfter = reifyConstantDim(padAfterAttr, builder);
+  Value paddedAfter = builder.create<arith::AddIOp>(paddedBefore, padAfter);
+
+  Value subOne = builder.create<arith::SubIOp>(kernelDim, one);
+  Value dilation = reifyConstantDim(dilationAttr, builder);
+  Value dilated = builder.create<arith::MulIOp>(dilation, subOne);
+  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);
+}
+
 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");
 
-    auto dynamicDimsOr =
-        checkHasDynamicBatchDims(rewriter, op, {input, op.output()});
-    if (!dynamicDimsOr.hasValue())
-      return failure();
-    SmallVector<Value> dynamicDims = dynamicDimsOr.getValue();
-
     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);
+
     auto weightShape = weightTy.getShape();
 
     // Apply padding as necessary.
     Attribute zeroAttr = rewriter.getZeroAttr(inputETy);
     if (isQuantized) {
       auto quantizationInfo =
           op->getAttr("quantization_info").cast<tosa::ConvOpQuantizationAttr>();
       auto iZp = quantizationInfo.input_zp().getValue().getSExtValue();
Value weightPermValue =
         rewriter.create<arith::ConstantOp>(loc, weightPermAttr);
     Type newWeightTy =
         RankedTensorType::get(newWeightShape, weightTy.getElementType());
     weight = rewriter.create<tosa::TransposeOp>(loc, newWeightTy, weight,
                                                 weightPermValue);
 
     Attribute resultZeroAttr = rewriter.getZeroAttr(resultETy);
     Value initTensor = rewriter.create<linalg::InitTensorOp>(
-        loc, dynamicDims, resultTy.getShape(), resultETy);
+        loc, filteredDims, resultTy.getShape(), resultETy);
     Value zero = rewriter.create<arith::ConstantOp>(loc, resultZeroAttr);
     Value zeroTensor =
         rewriter.create<linalg::FillOp>(loc, zero, initTensor).getResult(0);
 
     // Extract the attributes for convolution.
     llvm::SmallVector<int64_t> stride, dilation;
     getValuesFromIntArrayAttribute(strideTosaAttr, stride);
     getValuesFromIntArrayAttribute(dilationTosaAttr, dilation);
 
     // Create the convolution op.
     auto strideAttr = DenseIntElementsAttr::get(
         RankedTensorType::get({2}, rewriter.getI64Type()), stride);
     auto dilationAttr = DenseIntElementsAttr::get(
         RankedTensorType::get({2}, rewriter.getI64Type()), dilation);
 
     // Create maps for the bias broadcasting
     SmallVector<AffineMap, 4> indexingMaps;
     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()));
 
     Value biasInitTensor = rewriter.create<linalg::InitTensorOp>(
-        loc, dynamicDims, resultTy.getShape(), resultETy);
+        loc, filteredDims, resultTy.getShape(), resultETy);
 
     if (isQuantized) {
       auto quantizationInfo =
           op->getAttr("quantization_info").cast<tosa::ConvOpQuantizationAttr>();
       auto iZp = rewriter.getI32IntegerAttr(
           quantizationInfo.input_zp().getValue().getSExtValue());
       auto kZp = rewriter.getI32IntegerAttr(
           quantizationInfo.weight_zp().getValue().getSExtValue());

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

func @conv2d_dyn(%input: tensor<?x49x42x27xf32>, %weights: tensor<28x3x3x27xf32>, %bias: tensor<28xf32>) -> () {
   // CHECK:   %[[ADD:.+]] = arith.addf
   // CHECK:   linalg.yield %[[ADD]] : f32
   %0 = "tosa.conv2d"(%input, %weights, %bias) {pad = [0, 0, 0, 0], stride = [1, 1], dilation = [2, 1]} : (tensor<?x49x42x27xf32>, tensor<28x3x3x27xf32>, tensor<28xf32>)  -> (tensor<?x45x40x28xf32>)
   return
 }
 
 // -----
 
+// CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
+// CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+
+// CHECK-LABEL: @conv2d_dyn_w_h
+func @conv2d_dyn_w_h(%input: tensor<1x?x?x27xf32>, %weights: tensor<28x3x3x27xf32>, %bias: tensor<28xf32>) -> () {
+  // Computing output height
+  // CHECK: %[[C1:.+]] = arith.constant 1
+  // CHECK: %[[H:.+]] = tensor.dim %arg0, %[[C1]]
+  // CHECK: %[[C1_0:.+]] = arith.constant 1
+  // CHECK: %[[KH:.+]] = tensor.dim %arg1, %[[C1_0]]
+  // CHECK: %[[ONE:.+]] = arith.constant 1 : index
+  // CHECK: %[[PAD_0:.+]] = arith.constant 0 : index
+  // CHECK: %[[ADD_PAD_0:.+]] = arith.addi %[[H]], %[[PAD_0]] : index
+  // CHECK: %[[PAD_1:.+]] = arith.constant 0 : index
+  // CHECK: %[[ADD_PAD_1:.+]] = arith.addi %[[ADD_PAD_0]], %[[PAD_1]] : index
+  // CHECK: %[[SUB_ONE:.+]] = arith.subi %[[KH]], %[[ONE]] : index
+  // CHECK: %[[DIL_H:.+]] = arith.constant 2 : index
+  // CHECK: %[[DILATED:.+]] = arith.muli %[[DIL_H]], %[[SUB_ONE]] : index
+  // CHECK: %[[ADD_ONE:.+]] = arith.addi %[[DILATED]], %[[ONE]] : index
+  // CHECK: %[[SUBTRACTED:.+]] = arith.subi %[[ADD_PAD_1]], %[[ADD_ONE]] : index
+  // CHECK: %[[STRIDE_H:.+]] = arith.constant 1 : index
+  // CHECK: %[[DIVIDED:.+]] = arith.divui %[[SUBTRACTED]], %[[STRIDE_H]] : index
+  // CHECK: %[[H_OUT:.+]] = arith.subi %[[DIVIDED]], %[[ONE]] : index
+
+  // Computing output width
+  // CHECK: %[[C2:.+]] = arith.constant 2
+  // CHECK: %[[W:.+]] = tensor.dim %arg0, %[[C2]]
+  // CHECK: %[[C2_0:.+]] = arith.constant 2
+  // CHECK: %[[KW:.+]] = tensor.dim %arg1, %[[C2_0]]
+  // CHECK: %[[ONE_0:.+]] = arith.constant 1 : index
+  // CHECK: %[[PAD_2:.+]] = arith.constant 0 : index
+  // CHECK: %[[ADD_PAD_2:.+]] = arith.addi %[[W]], %[[PAD_2]] : index
+  // CHECK: %[[PAD_3:.+]] = arith.constant 0 : index
+  // CHECK: %[[ADD_PAD_3:.+]] = arith.addi %[[ADD_PAD_2]], %[[PAD_3]] : index
+  // CHECK: %[[SUB_ONE_0:.+]] = arith.subi %[[KW]], %[[ONE_0]] : index
+  // CHECK: %[[DIL_W:.+]] = arith.constant 1 : index
+  // CHECK: %[[DILATED_0:.+]] = arith.muli %[[DIL_W]], %[[SUB_ONE_0]] : index
+  // CHECK: %[[ADD_ONE_0:.+]] = arith.addi %[[DILATED_0]], %[[ONE_0]] : index
+  // CHECK: %[[SUBTRACTED_0:.+]] = arith.subi %[[ADD_PAD_3]], %[[ADD_ONE_0]] : index
+  // CHECK: %[[STRIDE_W:.+]] = arith.constant 1 : index
+  // CHECK: %[[DIVIDED_0:.+]] = arith.divui %[[SUBTRACTED_0]], %[[STRIDE_W]] : index
+  // CHECK: %[[W_OUT:.+]] = arith.subi %[[DIVIDED_0]], %[[ONE_0]] : index
+
+  // Running convolution
+  // CHECK: %[[PERM:.+]] = arith.constant dense<[1, 2, 3, 0]>
+  // CHECK: %[[WEIGHT:.+]] = "tosa.transpose"(%arg1, %[[PERM]])
+  // CHECK: %[[M_IN:.+]] = linalg.init_tensor [1, %[[H_OUT]], %[[W_OUT]], 28]
+  // CHECK: %[[CST:.+]] = arith.constant 0
+  // CHECK: %[[FILL:.+]] = linalg.fill
+  // CHECK: %[[B_IN:.+]] = linalg.init_tensor [1, %[[H_OUT]], %[[W_OUT]], 28]
+  // CHECK: %[[CONV:.+]] = linalg.conv_2d_nhwc_hwcf {dilations = dense<[2, 1]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %[[WEIGHT]] : tensor<1x?x?x27xf32>, tensor<3x3x27x28xf32>) outs(%[[FILL]] : tensor<1x?x?x28xf32>)
+  // CHECK: %[[B:.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, %[[CONV]] : tensor<28xf32>, tensor<1x?x?x28xf32>) outs(%[[B_IN]] : tensor<1x?x?x28xf32>)
+  // CHECK:   %[[ADD:.+]] = arith.addf
+  // CHECK:   linalg.yield %[[ADD]] : f32
+  %0 = "tosa.conv2d"(%input, %weights, %bias) {pad = [0, 0, 0, 0], stride = [1, 1], dilation = [2, 1]} : (tensor<1x?x?x27xf32>, tensor<28x3x3x27xf32>, tensor<28xf32>)  -> (tensor<1x?x?x28xf32>)
+  return
+}
+
+// -----
+
 // CHECK-LABEL: @conv2d_padded_f32
 func @conv2d_padded_f32(%input: tensor<1x47x40x28xf32>, %weights: tensor<28x3x3x28xf32>, %bias: tensor<28xf32>) -> () {
   // CHECK: %[[C0:.+]] = arith.constant 0
   // CHECK: tensor.pad %arg0 low[0, 1, 1, 0] high[0, 1, 1, 0]
   // CHECK:   tensor.yield %[[C0]]
   // CHECK: linalg.conv_2d_nhwc_hwcf
   %0 = "tosa.conv2d"(%input, %weights, %bias) {pad = [1, 1, 1, 1], stride = [1, 1], dilation = [2, 1]} : (tensor<1x47x40x28xf32>, tensor<28x3x3x28xf32>, tensor<28xf32>)  -> (tensor<1x45x40x28xf32>)
   return