From 6e895fe70a227472a51e1f88f8d795234bb183d3 Mon Sep 17 00:00:00 2001 From: Yi-Hong Lyu Date: Wed, 19 Jul 2023 15:23:48 -0700 Subject: [PATCH] Parallelize Max (#16745) It gives up to 7.5% improvement in LLaMA 7B case. --- onnxruntime/core/providers/cpu/math/element_wise_ops.cc | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/cpu/math/element_wise_ops.cc b/onnxruntime/core/providers/cpu/math/element_wise_ops.cc index a50c462703..3192c8573c 100644 --- a/onnxruntime/core/providers/cpu/math/element_wise_ops.cc +++ b/onnxruntime/core/providers/cpu/math/element_wise_ops.cc @@ -855,7 +855,12 @@ struct Max_8::ComputeImpl { }}; int input_count = inst.Node().InputArgCount().front(); - UntypedBroadcastVariadic(input_count, *context, typed_allocator, funcs); + // TODO: Parallelize across spans in UntypedBroadcastVariadic to avoid specific logic here + if (input_count == 2) { + UntypedBroadcastTwo(*context, funcs, 1.0); + } else { + UntypedBroadcastVariadic(input_count, *context, typed_allocator, funcs); + } return Status::OK(); }