Bug in Average Pool2D topi implementation for Integer inputs

[anijain2305]

def manual_test():
    data_shape=(1, 1, 2, 2)
    data_dtype='int32'
    data = relay.var("data", shape=data_shape,
            dtype=data_dtype)
    pool = relay.op.nn.avg_pool2d(data, pool_size=(2, 2))
    func = pool
    func = run_infer_type(func)
    func = relay.Function(relay.analysis.free_vars(func),
                func)
    print(func)
    with relay.build_config(opt_level=0):
        graph, lib, params = relay.build(func, "llvm", params=None)
        mod = graph_runtime.create(graph, lib, ctx=tvm.cpu(0))
        golden_data = np.array([5, 5, 5, 5]).reshape(data_shape).astype(data_dtype)
        mod.set_input("data",golden_data)
        mod.run()
        res = mod.get_output(0).asnumpy()
        print(res)

The above gives output 4 instead of 5.

Setting the golden_data to (1, 1, 1, 1) gives output 0.

I think we are dividing before accumulating. Since, we have only tested on Float32, this never showed up before. But, in int32, it performs int division.

[tqchen]

We want to update this the compute to move the division after the accumulation. @anijain2305 can you look into that?

[FrozenGene]

QAveragePool should be rewritten. And if we want to keep compatibility with TFLite, we should make the accumulation type be "Uint16", not "int32" here. I know it is very strange, but we meet the correctness problem when we parse TFLite model. We find TFLite's opt implementation is "UINT16", but the non-opt (reference) implementation is "int32". However, the default behavior of TFLite is opt. I wish we could consider this scenario. Otherwise, we can not keep the compatibility with TFLite.

[tqchen]

@FrozenGene for avg_pool, we can just use composed op to solve the problem.

avg_pool2d(int_8_value.astype("uint16")).astype("int8")

[anijain2305]

Yes, I will take a look next week.
I quickly tried to call topi::divide after tvm::sum. Threw me an error that the Reduction axis can only at level 0 and suggested to create a new Tensor.
@yzhliu suggested to call nested tvm::compute. I am not sure if having nested compute affects existing schedules. So, it might be little longer task, given what happens. Will try nested compute first.

[FrozenGene]

@anijain2305 Write two isolated tvm::compute could solve this issue. I also met this problem. I write one topi::nn::q_pool implementation

[anijain2305]

@tqchen @FrozenGene

I made following changes - basically creating an intermediate tensor for just sum.

    // TVM compute for summing the pooling windows.
    auto pool_sum = tvm::compute(out_shape,
    [&](const Array<Var>& output) {
      Array<Expr> indices;
      for (const Var& var : output) indices.push_back(var);
      indices.Set(height_axis, output[height_axis] * stride_height + dheight);
      indices.Set(width_axis, output[width_axis] * stride_width + dwidth);
      return tvm::sum(temp(indices), { dheight, dwidth });
    }, "tensor", "pool_sum");

    // TVM compute for dividing the reduced window sum by kernel size.
    return tvm::compute(out_shape,
    [&](const Array<Var>& output) {
      Array<Expr> indices;
      for (const Var& var : output) indices.push_back(var);

      if (count_include_pad) {
        return pool_sum(indices) / (kernel_height * kernel_width);
      } else {
        Expr h_start = output[height_axis] * stride_height - pad_top;
        Expr w_start = output[width_axis] * stride_width - pad_left;
        Expr h_end = ir::Min::make(h_start + kernel_height, height);
        Expr w_end = ir::Min::make(w_start + kernel_width, width);
        h_start = ir::Max::make(h_start, make_const(Int(32), 0));
        w_start = ir::Max::make(w_start, make_const(Int(32), 0));
        Expr divide_factor = ir::Max::make((h_end - h_start) * (w_end - w_start),
                                           make_const(Int(32), 1));
        return pool_sum(indices) / divide_factor;
      }
    }, "tensor", "pool_avg");

I am still getting the following error. Any direction?

TVMError: Check failed: 0 == level_: Reductions are only allowed at the top level of compute.
Please create another tensor for further composition.

[anijain2305]

Fixed. Closing.