[Bug] AveragePool operator produces wrong shape when ceil mode is set to 1

[coffezhou]

Expected behavior

The AveragePool operator in TVM should produce right shape.

Actual behavior

For the following model,

when ceil_mode = 1, the shapes of maxpool_output and avgpool_output are as follows:

maxpool_output (1, 3, 17, 17)
avgpool_output (1, 3, 9, 9)
maxpool_output (1, 3, 17, 17)
avgpool_output (1, 3, 10, 10)

For the documents of AveragePool, if ceil_mode = 1, the shape can be calculated by the following formula:

output_spatial_shape[i] = ceil((input_spatial_shape[i] + pad_shape[i] - dilation[i] * (kernel_shape[i] - 1) - 1) / strides_spatial_shape[i] + 1)

that is, ceil((17+1-1*(2-1)-1)/2 + 1) = 9, which indicates that TVM produces wrong shape of average_output.

I also try to set ceil_mode = 0, the results are as follows:

maxpool_output (1, 3, 17, 17)
avgpool_output (1, 3, 9, 9)
maxpool_output (1, 3, 17, 17)
avgpool_output (1, 3, 9, 9)

In this case, the results are right.

Environment

OS: Ubuntu 20.04
TVM: 0.23.dev0 (f4e28d3)

onnxruntime: 1.23.2

Steps to reproduce

This bug can be reproduced by the following code with the model in the attachment.

from typing import Dict, List, Literal, Optional
import sys
import os

import numpy as np
import onnx
import onnxruntime
from onnx import ModelProto, TensorProto, helper

import tvm
import tvm.testing
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx

import argparse
import pickle

def test(
    model: ModelProto,
    inputs: Optional[Dict[str, np.ndarray]] = None,
    ir_version: int = 8,
    opset: int = 14,
) -> None:
    # Configure model format.
    if ir_version is not None:
        model.ir_version = ir_version
    if opset is not None:
        model.opset_import[0].version = opset
    
    with open("inputs.pkl", 'rb') as fp:
        inputs = pickle.load(fp)
    # Run the model through onnx to get the expected result.
    try:
        ort_session = onnxruntime.InferenceSession(
            model.SerializeToString(), providers=["CPUExecutionProvider"]
        )
        ort_output = ort_session.run([], inputs)
    except Exception as e:
        print(e)
        print("This model cannot be executed by onnxruntime!")
        sys.exit(1)

    print('maxpool_output', ort_output[0].shape)
    print('avgpool_output', ort_output[1].shape)

    tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True)
    tvm_model = relax.transform.DecomposeOpsForInference()(tvm_model)
    tvm_model = relax.transform.LegalizeOps()(tvm_model)

    tvm_model, params = relax.frontend.detach_params(tvm_model)
    with tvm.transform.PassContext(opt_level=3):
        ex = tvm.compile(tvm_model, target="llvm")
        vm = relax.VirtualMachine(ex, tvm.cpu())

    input_list = [
        inputs[key.name_hint] for key in tvm_model["main"].params if key.name_hint in inputs
    ]
    if params:
        input_list += params["main"]

    # Run model and check outputs.
    vm.set_input("main", *input_list)
    vm.invoke_stateful("main")
    tvm_output = vm.get_outputs("main")

    print('maxpool_output', tvm_output[0].shape)
    print('avgpool_output', tvm_output[1].shape)
    
    
if __name__ == "__main__":
    
    onnx_model = onnx.load("22.onnx")
    test(onnx_model)

testcase.zip

In the attachment, 11.onnx is the model with ceil_mode=1, 22.onnx is ceil_mode=0.

Triage

Please refer to the list of label tags here to find the relevant tags and add them below in a bullet format (example below).

• needs-triage