Repository: microsoft/onnxruntime
File: onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc (C++)

Code before:
  // Compute the offset of gamma and beta to support broadcasting.
  int64_t i = LAYER_NORM_SCALE_BIAS_OFFSET(broadcast_param, task_idx, norm_size);

  for (int64_t h = 0; h < norm_size; h++, i++) {
    if (simplified) {
      p_output[h] = p_output[h] / mean_square * scale_data[i];
    } else if (nullptr == bias_data) {
      p_output[h] = (p_output[h] - mean) / mean_square * scale_data[i];
    } else {
      p_output[h] = (p_output[h] - mean) / mean_square * scale_data[i] + bias_data[i];
    }
  }

  if (mean_data != nullptr) {
    // ONNX spec doesn't support 'double' for 'U' so when 'T' == double, 'U' == float and we need to narrow
    mean_data[task_idx] = gsl::narrow_cast<float>(mean);
  }

  if (inv_std_dev_data != nullptr) {
    inv_std_dev_data[task_idx] = gsl::narrow_cast<float>(1 / mean_square);
  }
}

// Helper to convert int64_t -> Eigen::Index safely
inline Eigen::Index ToEigenIndex(int64_t v) {
  return static_cast<Eigen::Index>(v);
}

template <typename U>
void ComputeJob(
    const MLFloat16* X_data,
    const MLFloat16* scale_data,
    const MLFloat16* bias_data,
    const ptrdiff_t task_idx,
    const int64_t norm_size,
    const int64_t broadcast_param,
    const float* scale_float_ptr,
    const float* bias_float_ptr,
    float epsilon,
    bool simplified,
    MLFloat16* Y_data,
    U* mean_data,
    U* inv_std_dev_data,
    AllocatorPtr alloc) {
  ORT_UNUSED_PARAMETER(scale_data);
  ORT_UNUSED_PARAMETER(bias_data);
  ORT_UNUSED_PARAMETER(alloc);

  // reinterpret input/output MLFloat16* as Eigen::half*
  const Eigen::half* p_input = reinterpret_cast<const Eigen::half*>(X_data + task_idx * norm_size);
  Eigen::half* p_output = reinterpret_cast<Eigen::half*>(Y_data + task_idx * norm_size);

Review comment (suggestion):
I wouldn't consider `static_cast` a safe conversion. a checked conversion like `narrow` would be safer.

Code after:
  int64_t i = LAYER_NORM_SCALE_BIAS_OFFSET(broadcast_param, task_idx, norm_size);

  for (int64_t h = 0; h < norm_size; h++, i++) {
    if (simplified) {
      p_output[h] = p_output[h] / mean_square * scale_data[i];
    } else if (nullptr == bias_data) {
      p_output[h] = (p_output[h] - mean) / mean_square * scale_data[i];
    } else {
      p_output[h] = (p_output[h] - mean) / mean_square * scale_data[i] + bias_data[i];
    }
  }

  if (mean_data != nullptr) {
    // ONNX spec doesn't support 'double' for 'U' so when 'T' == double, 'U' == float and we need to narrow
    mean_data[task_idx] = gsl::narrow_cast<float>(mean);
  }

  if (inv_std_dev_data != nullptr) {
    inv_std_dev_data[task_idx] = gsl::narrow_cast<float>(1 / mean_square);
  }
}

// Helper to convert int64_t -> Eigen::Index safely
inline Eigen::Index ToEigenIndex(int64_t v) {
  return narrow<Eigen::Index>(v);
}

template <typename U>
void ComputeJob(
    const MLFloat16* X_data,
    const MLFloat16* scale_data,
    const MLFloat16* bias_data,
    const ptrdiff_t task_idx,
    const int64_t norm_size,
    const int64_t broadcast_param,
    const float* scale_float_ptr,
    const float* bias_float_ptr,
    float epsilon,
    bool simplified,
    MLFloat16* Y_data,
    U* mean_data,
    U* inv_std_dev_data,
    AllocatorPtr alloc) {
  ORT_UNUSED_PARAMETER(scale_data);  // only used in float/double overload
  ORT_UNUSED_PARAMETER(bias_data);   // only used in float/double overload
  ORT_UNUSED_PARAMETER(alloc);       // only required to create temporary float buffers

  // reinterpret input/output MLFloat16* as Eigen::half*
  const Eigen::half* p_input = reinterpret_cast<const Eigen::half*>(X_data + task_idx * norm_size);
  Eigen::half* p_output = reinterpret_cast<Eigen::half*>(Y_data + task_idx * norm_size);