[mlir][tosa] Revert added support for dynamic height/weight for pooling in tosa-to-linalg

[NatashaKnk]

Summary

Partial rollback to D133389

[jpienaar]

Thanks Natasha, I'll help with submit

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 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>(
+static mlir::Value
+getConvOutputDim(Location loc, Value inputDim, 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(inputDim.getType(), 1));
   Value padBefore = reifyConstantDim(padBeforeAttr, builder);
   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);
 }
 
-// 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) {
+// 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);
   }
 
-  ImplicitLocOpBuilder builder(loc, rewriter);
   // Dynamic input height
   if (inputTy.isDynamicDim(heightDim)) {
-    Value inputHDim =
-        builder.create<tensor::DimOp>(loc, input, heightDim).getResult();
-    Value kernelHDim = getKernelDim(loc, weight, weightHDim, builder);
+    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] = getConvOrPoolOutputDim(
-        loc, inputHDim, padAttr.getValue()[0], padAttr.getValue()[1],
-        kernelHDim, strideAttr.getValue()[0], dilationAttr.getValue()[0],
-        inputETy, builder);
+    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 inputWDim =
-        builder.create<tensor::DimOp>(loc, input, weightDim).getResult();
-    Value kernelWDim = getKernelDim(loc, weight, weightWDim, builder);
+    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] = getConvOrPoolOutputDim(
-        loc, inputWDim, padAttr.getValue()[2], padAttr.getValue()[3],
-        kernelWDim, strideAttr.getValue()[1], dilationAttr.getValue()[1],
-        inputETy, builder);
+    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;
 }
 
 // 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 = inferDynamicDimsForConvOrPool(
+    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) {
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 = inferDynamicDimsForConvOrPool(
+    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 =
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 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);
+    auto dynamicDimsOr =
+        checkHasDynamicBatchDims(rewriter, op, {input, op.getOutput()});
+    if (!dynamicDimsOr.has_value())
+      return failure();
+    SmallVector<Value> dynamicDims = dynamicDimsOr.value();
 
     // 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, filteredDims, resultTy.getShape(), resultTy.getElementType());
+        loc, dynamicDims, 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 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);
+    auto dynamicDimsOr =
+        checkHasDynamicBatchDims(rewriter, op, {input, op.getOutput()});
+    if (!dynamicDimsOr.has_value())
+      return failure();
+    SmallVector<Value> dynamicDims = dynamicDimsOr.value();
 
     // 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, filteredDims, accTy.getShape(), accETy);
+        loc, dynamicDims, 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, filteredDims, resultTy.getShape(), resultETy);
+        loc, dynamicDims, 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<?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
+func.func @max_pool_dyn(%arg0: tensor<?x6x34x62xf32>) -> () {
+  // CHECK: %[[C0:.+]] = arith.constant 0
+  // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]]
   // CHECK: %[[CONST:.+]] = arith.constant -3.40282347E+38
-  // CHECK: %[[INIT:.+]] = linalg.init_tensor
-  // CHECK: %[[FILL:.+]] = linalg.fill ins(%cst_18 : f32) outs(%20 : tensor<?x?x?x64xf32>) -> tensor<?x?x?x64xf32>
+  // CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[BATCH]], 4, 32, 62]
+  // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CONST]]{{.*}}outs(%[[INIT]]
   // CHECK: %[[KERNEL:.+]] = linalg.init_tensor [3, 3]
-  // 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>)
+  // 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>)
   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