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

Code before:
      KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<TYPE>()), \
      Expand_8<TYPE>);

#define REG_EXPAND_KERNEL(TYPE) REG_EXPAND_KERNEL_WITH_TYPE_NAME(TYPE, TYPE)

REG_EXPAND_KERNEL_WITH_TYPE_NAME(std::string, string)

template <>
Status Erf<float>::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 float* input_data = X->Data<float>();
  float* output_data = Y->MutableData<float>();
  concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
  int64_t elem_count = X->Shape().Size();
  constexpr int64_t length_per_task = 4096;  // this number comes from FastGelu.
  int64_t task_count = (elem_count + length_per_task - 1) / length_per_task;
  concurrency::ThreadPool::TryBatchParallelFor(
      tp, narrow<std::ptrdiff_t>(task_count),
      [&](ptrdiff_t task_idx) {
        const auto start = task_idx * length_per_task;
        const float* p_input = input_data + start;
        float* p_output = output_data + start;
        int64_t count = std::min(length_per_task, elem_count - start);
        MlasComputeErf(p_input, p_output, narrow<std::ptrdiff_t>(count));
      },
      0);

  return Status::OK();
}

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(

Review comment (suggestion):
can the `narrow` move outside? e.g., `elem_count`, `task_count`, and `length_per_task` could be made `ptrdiff_t` instead

Code after:
      Expand_8<TYPE>);

#define REG_EXPAND_KERNEL(TYPE) REG_EXPAND_KERNEL_WITH_TYPE_NAME(TYPE, TYPE)

REG_EXPAND_KERNEL_WITH_TYPE_NAME(std::string, string)

template <>
Status Erf<float>::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 float* input_data = X->Data<float>();
  float* output_data = Y->MutableData<float>();
  concurrency::ThreadPool* tp = context->GetOperatorThreadPool();

  const std::ptrdiff_t elem_count = onnxruntime::narrow<std::ptrdiff_t>(X->Shape().Size());
  constexpr std::ptrdiff_t length_per_task = 4096;  // this number comes from FastGelu.
  const std::ptrdiff_t task_count = (elem_count + length_per_task - 1) / length_per_task;

  concurrency::ThreadPool::TryBatchParallelFor(
      tp, task_count,
      [&](ptrdiff_t task_idx) {
        const std::ptrdiff_t start = task_idx * length_per_task;
        const float* p_input = input_data + start;
        float* p_output = output_data + start;
        const std::ptrdiff_t count = std::min(length_per_task, elem_count - start);
        MlasComputeErf(p_input, p_output, count);
      },
      0);

  return Status::OK();
}

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));