Optimize ReduceSum, ReduceMean, ReduceMin, ReduceMax

[xadupre]

Description:

Improves ReduceSum, ReduceMax, ReduceMean, ReduceMin on CPU for cases RKR (the tensor is equivalent to a 3D tensor with one reduced axis, one kept axis, one reduced axis. It reduces the memory consumption) (30% faster).

Motivation and Context

• performance
• Raised by issue ReduceMean consumes an unreasonable amount of VRAM #10099.

See the following script to measure the improvment:

import time
import numpy as np
from tqdm import tqdm
import onnx
import onnxruntime as ort


def create_onnx_model(node):
    # Size is (batch_size / 2) GB
    input_proto = onnx.helper.make_tensor_value_info('x', onnx.TensorProto.FLOAT, [None, *FEATURE_MAP_SHAPE])
    output_proto = onnx.helper.make_tensor_value_info('y', onnx.TensorProto.FLOAT, [1, FEATURE_MAP_SHAPE[0], 1, 1])

    node_def = onnx.helper.make_node(
        node,
        inputs=[input_proto.name],
        outputs=[output_proto.name],
        name='op',
        axes=(0, 2, 3),
    )

    graph_def = onnx.helper.make_graph(
        nodes=[node_def],
        name='test-model',
        inputs=[input_proto],
        outputs=[output_proto],
    )

    model_def = onnx.helper.make_model(graph_def, producer_name='onnx-example')
    model_def.ir_version = 4
    model_def.opset_import[0].version = 11

    onnx.checker.check_model(model_def, full_check=True)

    onnx.save_model(model_def, 'test.onnx')

    return ort.InferenceSession('test.onnx', providers=['CPUExecutionProvider'])


batch_size = 16
FEATURE_MAP_SHAPE = (128, 512, 512)
shape = (batch_size, *FEATURE_MAP_SHAPE)

N = 5

x = np.random.randn(*shape).astype(np.float32)

results = []
nodes = ['ReduceSum', 'ReduceMean', 'ReduceMin', 'ReduceMax', 'ReduceSum']
with tqdm(total=N * len(nodes)) as pbar:
    for node in nodes:
        model = create_onnx_model(node)
        durs = []
        for i in range(N):
            x *= -1
            begin = time.perf_counter()
            model.run(None, {'x': x})[0].shape
            duration = time.perf_counter() - begin
            durs.append(duration)
            pbar.update(1)
        durs.sort()
        results.append((node, durs))

for node, durs in results:
    print(node, sum(durs[1:-1]) / (len(durs) - 2), durs)


"""
PR:

ReduceSum 0.12530109999999914 [0.11439150000000353, 0.12402879999999783, 0.1247271000000012, 0.12714739999999836, 0.13132369999999582]
ReduceMean 0.12379500000000121 [0.12040319999999838, 0.12092700000000178, 0.12471620000000172, 0.12574180000000013, 0.13234720000000522]
ReduceMin 0.12942566666666502 [0.12426919999999342, 0.12519710000000117, 0.12768109999999666, 0.13539879999999727, 0.13887529999999515]
ReduceMax 0.1296077999999999 [0.1240763000000058, 0.128204199999999, 0.13001030000000213, 0.13060889999999858, 0.1819714000000019]

No PR:

ReduceSum 0.16718310000000022 [0.1642044000000027, 0.16425759999999912, 0.1684125000000023, 0.16887919999999923, 0.21092490000000197]
ReduceMean 0.17254450000000313 [0.16611610000000354, 0.16648840000000575, 0.1679203000000058, 0.18322479999999786, 0.23252000000000095]
ReduceMin 0.1797097666666687 [0.14939999999999998, 0.17422279999999546, 0.17787720000000462, 0.18702930000000606, 0.2097418000000033]
ReduceMax 0.17173149999999993 [0.1606559999999959, 0.16596030000000184, 0.17150569999999732, 0.17772850000000062, 0.18525440000000515]

"""

[edgchen1]

/azp run onnxruntime-binary-size-checks-ci-pipeline

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[xadupre]

/azp run "Android CI Pipeline"

[azure-pipelines]

No pipelines are associated with this pull request.

[xadupre]

/azp run Android_CI

[azure-pipelines]

No pipelines are associated with this pull request.

[xadupre]

Hi @skottmckay, would it be possible to continue the review?

[skottmckay]