Repository: microsoft/onnxruntime
File: onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc (C++)

Code before:

        for (int i = 0; i < num_expert_threads; ++i) {
          const size_t thread_offset = static_cast<size_t>(i) * output_buffer_size;
          const float* src = thread_local_outputs + thread_offset + start_idx;
          float* dst = buffer + start_idx;

          size_t j = 0;
          const size_t chunk_size = end_idx - start_idx;
          const size_t unroll_factor = static_cast<size_t>(GetUnrollFactor(static_cast<int64_t>(chunk_size)));
          for (; j + unroll_factor <= chunk_size; j += unroll_factor) {
            for (size_t loop_k = 0; loop_k < unroll_factor; ++loop_k) {
              dst[j + loop_k] += src[j + loop_k];
            }
          }
          for (; j < chunk_size; ++j) {
            dst[j] += src[j];
          }
        }
      });
    } else {
      for (int i = 0; i < num_expert_threads; ++i) {
        const size_t thread_offset = static_cast<size_t>(i) * output_buffer_size;
        const float* src = thread_local_outputs + thread_offset;

        size_t j = 0;
        const size_t unroll_factor = static_cast<size_t>(GetUnrollFactor(static_cast<int64_t>(output_buffer_size)));
        for (; j + unroll_factor <= output_buffer_size; j += unroll_factor) {
          for (size_t loop_k = 0; loop_k < unroll_factor; ++loop_k) {
            buffer[j + loop_k] += src[j + loop_k];
          }
        }
        for (; j < output_buffer_size; ++j) {
          buffer[j] += src[j];
        }
      }
    }
  };

  if constexpr (std::is_same_v<T, MLFloat16>) {
    auto final_output_float_ptr = IAllocator::MakeUniquePtr<float>(allocator, output_buffer_size);
    float* final_output_float = final_output_float_ptr.get();
    accumulate(final_output_float);

    MlasConvertFloatToHalfBuffer(final_output_float,
                                 reinterpret_cast<MLFloat16*>(output->MutableData<T>()),
                                 static_cast<size_t>(output_buffer_size));
  } else {
    accumulate(output->MutableData<T>());
  }

  return Status::OK();

Review comment (suggestion):
>const size_t unroll_factor = static_cast<size_t>(GetUnrollFactor(static_cast<int64_t>(output_buffer_size))); [](http://example.com/codeflow?start=8&length=108)

Both casts should be narrow on this line #Resolved

Code after:

        for (int i = 0; i < num_expert_threads; ++i) {
          const size_t thread_offset = static_cast<size_t>(i) * output_buffer_size;
          const float* src = thread_local_outputs + thread_offset + start_idx;
          float* dst = buffer + start_idx;

          size_t j = 0;
          const size_t chunk_size = end_idx - start_idx;
          const size_t unroll_factor = static_cast<size_t>(GetUnrollFactor(static_cast<int64_t>(chunk_size)));
          for (; j + unroll_factor <= chunk_size; j += unroll_factor) {
            for (size_t loop_k = 0; loop_k < unroll_factor; ++loop_k) {
              dst[j + loop_k] += src[j + loop_k];
            }
          }
          for (; j < chunk_size; ++j) {
            dst[j] += src[j];
          }
        }
      });
    } else {
      for (int i = 0; i < num_expert_threads; ++i) {
        const size_t thread_offset = static_cast<size_t>(i) * output_buffer_size;
        const float* src = thread_local_outputs + thread_offset;

        size_t j = 0;
        const size_t unroll_factor = narrow<size_t>(GetUnrollFactor(narrow<int64_t>(output_buffer_size)));
        for (; j + unroll_factor <= output_buffer_size; j += unroll_factor) {
          for (size_t loop_k = 0; loop_k < unroll_factor; ++loop_k) {
            buffer[j + loop_k] += src[j + loop_k];
          }
        }
        for (; j < output_buffer_size; ++j) {
          buffer[j] += src[j];
        }
      }
    }
  };

  if constexpr (std::is_same_v<T, MLFloat16>) {
    auto final_output_float_ptr = IAllocator::MakeUniquePtr<float>(allocator, output_buffer_size);
    float* final_output_float = final_output_float_ptr.get();
    accumulate(final_output_float);

    MlasConvertFloatToHalfBuffer(final_output_float,
                                 reinterpret_cast<MLFloat16*>(output->MutableData<T>()),
                                 static_cast<size_t>(output_buffer_size));
  } else {
    accumulate(output->MutableData<T>());
  }

  return Status::OK();