[mlir][tosa] Add support for dynamic width and height for pooling operations in the tosa to linalg lowering

[NatashaKnk]

Pull in most recent version

[rsuderman]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

100–120

You can check but I don't think you actually need a template, function overloading should be enough. You would just have two functions with slightly different args but the same name.

117

You should check that dim is withinin the right range and return nullptr if not.

711–719

Ditto on op->getAttr comment.

792–800

rather than op->getAttr you should be able to do op.kernel(), etc.

[NatashaKnk]

Remove template, small fixes

[NatashaKnk]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

100–120

Whoops, good point

[silvas]

This patch appears to be miscompiling. See discussion on https://github.com/llvm/torch-mlir/issues/1361

Can we revert this patch?

[NatashaKnk]

Yes, AFAIK there shouldn't be issues in reverting this patch.

[silvas]

Okay, let's revert it then.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

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:
 // 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));
+static mlir::Value getConvOrPoolOutputDim(Location loc, Value inputDim,
+                                          Attribute padBeforeAttr,
+                                          Attribute padAfterAttr,
+                                          Value kernelDim, Attribute strideAttr,
+                                          Attribute dilationAttr, Type inputETy,
+                                          ImplicitLocOpBuilder &builder) {
+  auto one = builder.create<arith::ConstantOp>(
+      loc, IntegerAttr::get(inputDim.getType(), 1));
   Value padBefore = reifyConstantDim(padBeforeAttr, builder);
-  Value paddedBefore = builder.create<arith::AddIOp>(initDim, padBefore);
+  Value paddedBefore = builder.create<arith::AddIOp>(inputDim, 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::AddIOp>(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) {
+// For convolution, the kernel is a value.
+Value getKernelDim(Location loc, Value kernel, uint64_t dim,
+                   ImplicitLocOpBuilder &builder) {
+  return builder.create<tensor::DimOp>(loc, kernel, dim).getResult();
+}
+
+// For pooling, the kernel is an attribute.
+Value getKernelDim(Location loc, ArrayAttr kernel, uint64_t dim,
+                   ImplicitLocOpBuilder &builder) {
+  auto kernelArr = kernel.getValue();
+  if (dim >= kernelArr.size()) return nullptr;
+  Attribute kernelDimAttr = kernelArr[dim];
+  return reifyConstantDim(kernelDimAttr, builder);
+}
+
+// Creates a vector of the dynamic output dims convolution and pooling ops.
+template <typename T>
+static SmallVector<Value> inferDynamicDimsForConvOrPool(
+    Location loc, Value input, T 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);
   }
 
+  ImplicitLocOpBuilder builder(loc, rewriter);
   // 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();
+    Value inputHDim =
+        builder.create<tensor::DimOp>(loc, input, heightDim).getResult();
+    Value kernelHDim = getKernelDim(loc, weight, weightHDim, builder);
     // 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);
+    dynDims[heightDim] = getConvOrPoolOutputDim(
+        loc, inputHDim, padAttr.getValue()[0], padAttr.getValue()[1],
+        kernelHDim, strideAttr.getValue()[0], dilationAttr.getValue()[0],
+        inputETy, builder);
   }
 
   // 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();
+    Value inputWDim =
+        builder.create<tensor::DimOp>(loc, input, weightDim).getResult();
+    Value kernelWDim = getKernelDim(loc, weight, weightWDim, builder);
     // 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);
+    dynDims[weightDim] = getConvOrPoolOutputDim(
+        loc, inputWDim, padAttr.getValue()[2], padAttr.getValue()[3],
+        kernelWDim, strideAttr.getValue()[1], dilationAttr.getValue()[1],
+        inputETy, builder);
   }
 
   SmallVector<Value> filteredDims = condenseValues(dynDims);
   return filteredDims;
 }
 
 // Creates a map to collapse the last dimension of the Depthwise convolution op
 // due to a shape mismatch
matchAndRewrite(tosa::Conv2DOp op, OpAdaptor adaptor,
     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> filteredDims = inferDynamicDimsForConv(
+    SmallVector<Value> filteredDims = inferDynamicDimsForConvOrPool(
         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) {
matchAndRewrite(tosa::DepthwiseConv2DOp op, OpAdaptor adaptor,
     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(
+    SmallVector<Value> filteredDims = inferDynamicDimsForConvOrPool(
         loc, input, weight, resultTy, padAttr, strideTosaAttr, dilationTosaAttr,
         0, 1, rewriter);
 
     bool isQuantized = op->hasAttr("quantization_info");
     IntegerAttr iZp;
     IntegerAttr kZp;
     if (isQuantized) {
       auto quantizationInfo =
LogicalResult matchAndRewrite(tosa::MaxPool2dOp op,
                                 PatternRewriter &rewriter) const final {
     Location loc = op.getLoc();
     Value input = op.getInput();
     ShapedType inputTy = input.getType().cast<ShapedType>();
 
     ShapedType resultTy = op.getType().template cast<ShapedType>();
     Type resultETy = inputTy.getElementType();
 
-    auto dynamicDimsOr =
-        checkHasDynamicBatchDims(rewriter, op, {input, op.getOutput()});
-    if (!dynamicDimsOr.has_value())
-      return failure();
-    SmallVector<Value> dynamicDims = dynamicDimsOr.value();
+    auto kernelAttr = op.getKernel().cast<ArrayAttr>();
+    auto padAttr = op.getPad().cast<ArrayAttr>();
+    auto strideTosaAttr = op.getStride().cast<ArrayAttr>();
+    ArrayAttr dilationTosaAttr = rewriter.getI64ArrayAttr({1, 1});
+
+    SmallVector<Value> filteredDims = inferDynamicDimsForConvOrPool(
+        loc, input, kernelAttr, resultTy, padAttr, strideTosaAttr,
+        dilationTosaAttr,
+        /*weightHDim=*/0, /*weightWDim=*/1, rewriter);
 
     // Determine what the initial value needs to be for the max pool op.
     Attribute initialAttr;
     if (resultETy.isF32())
       initialAttr = rewriter.getFloatAttr(
           resultETy,
           APFloat::getLargest(resultETy.cast<FloatType>().getFloatSemantics(),
                               true));
LogicalResult matchAndRewrite(tosa::MaxPool2dOp op,
     getValuesFromIntArrayAttribute(op.getKernel(), kernel);
     getValuesFromIntArrayAttribute(op.getStride(), stride);
 
     Attribute strideAttr = rewriter.getI64VectorAttr(stride);
     Attribute dilationAttr = rewriter.getI64VectorAttr({1, 1});
 
     // Create the linalg op that performs pooling.
     Value initTensor = rewriter.create<linalg::InitTensorOp>(
-        loc, dynamicDims, resultTy.getShape(), resultTy.getElementType());
+        loc, filteredDims, resultTy.getShape(), resultTy.getElementType());
 
     Value filledInitTensor =
         rewriter
             .create<linalg::FillOp>(loc, ValueRange{initialValue},
                                     ValueRange{initTensor})
             .result();
 
     Value fakeWindowDims =
LogicalResult matchAndRewrite(tosa::AvgPool2dOp op,
 
     ShapedType resultTy = op.getType().template cast<ShapedType>();
     Type resultETy = op.getType().cast<ShapedType>().getElementType();
 
     Type accETy =
         inElementTy.isa<IntegerType>() ? rewriter.getI32Type() : inElementTy;
     ShapedType accTy = resultTy.clone(accETy);
 
-    auto dynamicDimsOr =
-        checkHasDynamicBatchDims(rewriter, op, {input, op.getOutput()});
-    if (!dynamicDimsOr.has_value())
-      return failure();
-    SmallVector<Value> dynamicDims = dynamicDimsOr.value();
+    auto kernelAttr = op.getKernel().cast<ArrayAttr>();
+    auto padArrayAttr = op.getPad().cast<ArrayAttr>();
+    auto strideTosaAttr = op.getStride().cast<ArrayAttr>();
+    ArrayAttr dilationTosaAttr = rewriter.getI64ArrayAttr({1, 1});
+
+    SmallVector<Value> filteredDims = inferDynamicDimsForConvOrPool(
+        loc, input, kernelAttr, resultTy, padArrayAttr, strideTosaAttr,
+        dilationTosaAttr,
+        /*weightHDim=*/0, /*weightWDim=*/1, rewriter);
 
     // Apply padding as necessary.
     llvm::SmallVector<int64_t> pad;
     pad.resize(2, 0);
     getValuesFromIntArrayAttribute(op.getPad(), pad);
     pad.resize(pad.size() + 2, 0);
     Attribute padAttr = rewriter.getZeroAttr(inElementTy);
     Value paddedInput = applyPad(loc, input, pad, padAttr, rewriter);
 
     Attribute initialAttr = rewriter.getZeroAttr(accETy);
     Value initialValue = rewriter.create<arith::ConstantOp>(loc, initialAttr);
 
     SmallVector<int64_t> kernel, stride;
     getValuesFromIntArrayAttribute(op.getKernel(), kernel);
     getValuesFromIntArrayAttribute(op.getStride(), stride);
 
     Attribute strideAttr = rewriter.getI64VectorAttr(stride);
     Attribute dilationAttr = rewriter.getI64VectorAttr({1, 1});
 
     // Create the linalg op that performs pooling.
     Value poolInitTensor = rewriter.create<linalg::InitTensorOp>(
-        loc, dynamicDims, accTy.getShape(), accETy);
+        loc, filteredDims, accTy.getShape(), accETy);
 
     Value filledInitTensor =
         rewriter
             .create<linalg::FillOp>(loc, ValueRange{initialValue},
                                     ValueRange{poolInitTensor})
             .result();
 
     Value fakeWindowDims =
         rewriter.create<linalg::InitTensorOp>(loc, kernel, accETy);
 
     // Sum across the pooled region.
     Value poolingOp = rewriter
                           .create<linalg::PoolingNhwcSumOp>(
                               loc, ArrayRef<Type>{accTy},
                               ValueRange{paddedInput, fakeWindowDims},
                               filledInitTensor, strideAttr, dilationAttr)
                           .getResult(0);
 
     // Normalize the summed value by the number of elements grouped in each
     // pool.
     auto poolingOpTy = poolingOp.getType().cast<ShapedType>();
     auto affineMap = rewriter.getMultiDimIdentityMap(resultTy.getRank());
 
     Value genericInitTensor = rewriter.create<linalg::InitTensorOp>(
-        loc, dynamicDims, resultTy.getShape(), resultETy);
+        loc, filteredDims, resultTy.getShape(), resultETy);
 
     auto genericOp = rewriter.create<linalg::GenericOp>(
         loc, ArrayRef<Type>({resultTy}), ValueRange{poolingOp},
         ValueRange{genericInitTensor},
         ArrayRef<AffineMap>({affineMap, affineMap}),
         getNParallelLoopsAttrs(resultTy.getRank()),
         [&](OpBuilder &b, Location loc, ValueRange args) {
           auto zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);

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

func.func @max_pool_padded(%arg0: tensor<1x6x34x62xf32>) -> () {
   // CHECK-DAG: [[FILL:%.+]] = linalg.fill ins([[INITVAL]]{{.*}}outs([[INIT]]
   // CHECK-DAG: [[KERNEL:%.+]] = linalg.init_tensor [3, 3]
   // CHECK: linalg.pooling_nhwc_max {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins([[PAD]], [[KERNEL]] : tensor<1x6x35x62xf32>, tensor<3x3xf32>) outs([[FILL]] : tensor<1x4x33x62xf32>)
   %0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 1], kernel = [3, 3], stride = [1, 1]} : (tensor<1x6x34x62xf32>)  -> (tensor<1x4x33x62xf32>)
   return
 }
 
 // CHECK-LABEL: @max_pool_dyn
-func.func @max_pool_dyn(%arg0: tensor<?x6x34x62xf32>) -> () {
-  // CHECK: %[[C0:.+]] = arith.constant 0
-  // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]]
+func.func @max_pool_dyn(%arg0: tensor<?x?x?x64xf32>) -> () {
+  // CHECK: %[[C0:.+]]  = arith.constant 0 : index
+  // CHECK: %[[DIM0:.+]]  = tensor.dim %arg0, %[[C0]] : tensor<?x?x?x64xf32>
+  // CHECK: %[[C1:.+]]  = arith.constant 1 : index
+  // CHECK: %[[DIM1:.+]]  = tensor.dim %arg0, %[[C1]] : tensor<?x?x?x64xf32>
+  // CHECK: arith.constant 2 : index
+  // CHECK: %[[C2:.+]]  = arith.constant 2 : index
+  // CHECK: %[[DIM2:.+]]  = tensor.dim %arg0, %[[C2]] : tensor<?x?x?x64xf32>
+  // CHECK: %[[PAD:.+]] = tensor.pad %arg0
   // CHECK: %[[CONST:.+]] = arith.constant -3.40282347E+38
-  // CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[BATCH]], 4, 32, 62]
-  // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CONST]]{{.*}}outs(%[[INIT]]
+  // CHECK: %[[INIT:.+]] = linalg.init_tensor
+  // CHECK: %[[FILL:.+]] = linalg.fill ins(%cst_18 : f32) outs(%20 : tensor<?x?x?x64xf32>) -> tensor<?x?x?x64xf32>
   // CHECK: %[[KERNEL:.+]] = linalg.init_tensor [3, 3]
-  // CHECK: linalg.pooling_nhwc_max {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%arg0, %[[KERNEL]] : tensor<?x6x34x62xf32>, tensor<3x3xf32>) outs(%[[FILL]] : tensor<?x4x32x62xf32>)
-  %0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 0], kernel = [3, 3], stride = [1, 1]} : (tensor<?x6x34x62xf32>)  -> (tensor<?x4x32x62xf32>)
+  // CHECK: linalg.pooling_nhwc_max {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>} ins(%[[PAD]], %[[KERNEL]] : tensor<?x?x?x64xf32>, tensor<3x3xf32>) outs(%[[FILL]] : tensor<?x?x?x64xf32>) -> tensor<?x?x?x64xf32>
+  %0 = "tosa.max_pool2d"(%arg0) {kernel = [3, 3], pad = [1, 1, 1, 1], stride = [2, 2]} : (tensor<?x?x?x64xf32>) -> (tensor<?x?x?x64xf32>)
   return
 }
 
 // CHECK-LABEL: @max_pool_i8
 func.func @max_pool_i8(%arg0: tensor<1x6x34x62xi8>) -> () {
   // CHECK: arith.constant -128
   // CHECK: linalg.pooling_nhwc_max
   %0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 0], kernel = [3, 3], stride = [1, 1]} : (tensor<1x6x34x62xi8>)  -> (tensor<1x4x32x62xi8>)
func.func @avg_pool_dyn(%arg0: tensor<?x6x34x62xf32>) -> (tensor<?x5x33x62xf32>) {
   // CHECK: %[[KERNEL:.+]] = linalg.init_tensor [4, 4]
   // CHECK: %[[POOL:.+]] = linalg.pooling_nhwc_sum {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%[[PAD]], %[[KERNEL]] : tensor<?x8x36x62xf32>, tensor<4x4xf32>) outs(%[[FILL]] : tensor<?x5x33x62xf32>)
   // CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[BATCH]], 5, 33, 62]
   // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%[[POOL]] : tensor<?x5x33x62xf32>) outs(%[[INIT]] : tensor<?x5x33x62xf32>)
   %0 = "tosa.avg_pool2d"(%arg0) {pad = [1, 1, 1, 1], kernel = [4, 4], stride = [1, 1]} : (tensor<?x6x34x62xf32>)  -> (tensor<?x5x33x62xf32>)
   return %0 : tensor<?x5x33x62xf32>
 }
 
+// CHECK-LABEL: @avg_pool_dyn_h
+func.func @avg_pool_dyn_h(%arg0: tensor<2x?x34x62xf32>) -> (tensor<2x?x33x62xf32>) {
+  // CHECK: %[[C1:.+]] = arith.constant 1
+  // CHECK: %[[DIM1:.+]] = tensor.dim %arg0, %[[C1]]
+  // CHECK: arith.addi
+  // CHECK: arith.addi
+  // CHECK: arith.addi
+  // CHECK: %[[RESULT:.+]] = arith.addi
+  // CHECK: %[[PAD:.+]] = tensor.pad %arg0 low[0, 1, 1, 0] high[0, 1, 1, 0]
+  // CHECK: %[[POOLINIT:.+]] = linalg.init_tensor [2, %[[RESULT]], 33, 62]
+  // CHECK: %[[FILL:.+]] = linalg.fill
+  // CHECK: %[[KERNEL:.+]] = linalg.init_tensor [4, 4]
+  // CHECK: %[[POOL:.+]] = linalg.pooling_nhwc_sum {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%[[PAD]], %[[KERNEL]] : tensor<2x?x36x62xf32>, tensor<4x4xf32>) outs(%[[FILL]] : tensor<2x?x33x62xf32>)
+  // CHECK: %[[INIT:.+]] = linalg.init_tensor [2, %[[RESULT]], 33, 62]
+  // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%[[POOL]] : tensor<2x?x33x62xf32>) outs(%[[INIT]] : tensor<2x?x33x62xf32>)
+  %0 = "tosa.avg_pool2d"(%arg0) {pad = [1, 1, 1, 1], kernel = [4, 4], stride = [1, 1]} : (tensor<2x?x34x62xf32>)  -> (tensor<2x?x33x62xf32>)
+  return %0 : tensor<2x?x33x62xf32>
+}
+
 // -----
 
 // CHECK-LABEL: @avg_pool_i8
 func.func @avg_pool_i8(%arg0 : tensor<1x128x128x2xi8>) -> () {
 
   // CHECK: linalg.pooling_nhwc_sum
   // CHECK: linalg.generic