Repository: microsoft/onnxruntime
File: onnxruntime/core/providers/cpu/math/element_wise_ops.cc (C++)

Code before:
template <>
Status Erf<MLFloat16>::Compute(OpKernelContext* context) const {
  const auto* X = context->Input<Tensor>(0);
  const auto& x_shape = X->Shape();
  auto* Y = context->Output(0, x_shape);
  const auto* input_data = X->Data<MLFloat16>();
  auto* output_data = Y->MutableData<MLFloat16>();
  concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
  int64_t elem_count = X->Shape().Size();
  constexpr int64_t length_per_task = 4096;
  int64_t task_count = (elem_count + length_per_task - 1) / length_per_task;

  const auto narrow_task_count = onnxruntime::narrow<std::ptrdiff_t>(task_count);

  concurrency::ThreadPool::TryBatchParallelFor(
      tp, narrow_task_count,
      [&](ptrdiff_t task_idx) {
        const auto start = task_idx * length_per_task;
        const int64_t count = std::min(length_per_task, elem_count - start);
        const auto narrow_count = onnxruntime::narrow<std::ptrdiff_t>(count);

        const MLFloat16* p_input = input_data + start;
        MLFloat16* p_output = output_data + start;

        std::vector<float> input_fp32(narrow_count);
        std::vector<float> output_fp32(narrow_count);

        MlasConvertHalfToFloatBuffer(p_input, input_fp32.data(), narrow_count);

        MlasComputeErf(input_fp32.data(), output_fp32.data(), narrow_count);

        MlasConvertFloatToHalfBuffer(output_fp32.data(), p_output, narrow_count);
      },
      0);

  return Status::OK();
}

class Mod final : public OpKernel {
 public:
  Mod(const OpKernelInfo& info) : OpKernel(info) {
    int64_t fmod = 0;
    Status s = info.GetAttr<int64_t>("fmod", &fmod);
    if (s.IsOK()) {
      ORT_ENFORCE((fmod == 0) || (fmod == 1), "fmod must have value either 0 or 1");
      fmod_ = (fmod == 1);
    }
  }

  Status Compute(OpKernelContext* context) const override;


Review comment (performance):
`std::vector` will make dynamic allocations. does this affect the performance?

maybe we can avoid some dynamic allocations by using the arena if it is enabled.

see here for an example (`OpKernelContext::GetTempSpaceAllocator` and `IAllocator::MakeUniquePtr`):
https://github.com/microsoft/onnxruntime/blob/3ec8829d449efb6eb4cb08cdb17b5ffe70595ed1/onnxruntime/contrib_ops/cpu/skip_layer_norm.cc#L167-L189

Code after:

template <>
Status Erf<MLFloat16>::Compute(OpKernelContext* context) const {
  const auto* X = context->Input<Tensor>(0);
  const auto& x_shape = X->Shape();
  auto* Y = context->Output(0, x_shape);
  const auto* input_data = X->Data<MLFloat16>();
  auto* output_data = Y->MutableData<MLFloat16>();
  concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
  int64_t elem_count = X->Shape().Size();
  constexpr int64_t length_per_task = 4096;
  int64_t task_count = (elem_count + length_per_task - 1) / length_per_task;

  const auto narrow_task_count = onnxruntime::narrow<std::ptrdiff_t>(task_count);

  // get allocator for temporary buffers
  AllocatorPtr alloc;
  ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc));

  concurrency::ThreadPool::TryBatchParallelFor(
      tp, narrow_task_count,
      [&](ptrdiff_t task_idx) {
        const auto start = task_idx * length_per_task;
        const int64_t count = std::min(length_per_task, elem_count - start);
        const auto narrow_count = onnxruntime::narrow<std::ptrdiff_t>(count);

        const MLFloat16* p_input = input_data + start;
        MLFloat16* p_output = output_data + start;

        // allocate temp buffers using ORT allocator
        IAllocatorUniquePtr<float> input_fp32 = IAllocator::MakeUniquePtr<float>(alloc, narrow_count);
        IAllocatorUniquePtr<float> output_fp32 = IAllocator::MakeUniquePtr<float>(alloc, narrow_count);

        // convert, compute, convert back
        MlasConvertHalfToFloatBuffer(p_input, input_fp32.get(), narrow_count);
        MlasComputeErf(input_fp32.get(), output_fp32.get(), narrow_count);
        MlasConvertFloatToHalfBuffer(output_fp32.get(), p_output, narrow_count);
      },
      0);

  return Status::OK();
}

class Mod final : public OpKernel {
 public:
  Mod(const OpKernelInfo& info) : OpKernel(info) {
    int64_t fmod = 0;
    Status s = info.GetAttr<int64_t>("fmod", &fmod);
    if (s.IsOK()) {
      ORT_ENFORCE((fmod == 0) || (fmod == 1), "fmod must have value either 0 or 1");
      fmod_ = (fmod == 1);