Repository: microsoft/onnxruntime
File: onnxruntime/core/providers/cpu/tensor/grid_sample.cc (C++)

Code before:
    if (x1 >= 0 && x1 < W_in && y2 >= 0 && y2 < H_in) {
      mask |= kBottomLeftMask;
    }
    if (x2 >= 0 && x2 < W_in && y2 >= 0 && y2 < H_in) {
      mask |= kBottomRightMask;
    }

    plan.x1 = x1;
    plan.x2 = x2;
    plan.y1 = y1;
    plan.y2 = y2;
    plan.w11 = dy2 * dx2;
    plan.w12 = dy2 * dx1;
    plan.w21 = dy1 * dx2;
    plan.w22 = dy1 * dx1;
    plan.mask = mask;
  }
}

template <typename T>
void EvaluatePlanForChannel(const T* input_data,
                            T* output_data,
                            int64_t W_in,
                            const BilinearSamplePlan2D<T>* plan_data,
                            int64_t point_count) {
  ORT_ENFORCE(input_data != nullptr, "EvaluatePlanForChannel requires non-null input_data.");
  for (int64_t idx = 0; idx < point_count; ++idx) {
    const auto& plan = plan_data[idx];
    if (plan.mask == 0) {
      output_data[idx] = T{};
      continue;
    }

#if defined(__ARM_NEON) || defined(__ARM_NEON__) || defined(__aarch64__) || defined(_M_ARM64)
    if constexpr (std::is_same_v<T, float>) {
      if (plan.mask == kAllNeighborsMask) {
        const float* row1_ptr = input_data + plan.y1 * W_in + plan.x1;
        const float* row2_ptr = input_data + plan.y2 * W_in + plan.x1;

        float32x2_t row1 = vld1_f32(row1_ptr);  // [p11, p12]
        float32x2_t row2 = vld1_f32(row2_ptr);  // [p21, p22]
        float32x4_t neighbors = vcombine_f32(row1, row2);

        float32x2_t weights_row1 = vdup_n_f32(plan.w12);
        weights_row1 = vset_lane_f32(plan.w11, weights_row1, 0);
        float32x2_t weights_row2 = vdup_n_f32(plan.w22);
        weights_row2 = vset_lane_f32(plan.w21, weights_row2, 0);
        float32x4_t weights = vcombine_f32(weights_row1, weights_row2);

        float32x4_t products = vmulq_f32(neighbors, weights);
        float32x2_t sum_pairs = vadd_f32(vget_low_f32(products), vget_high_f32(products));

Review comment (refactor):
Can the input just be a reference ? Looking at the call site, we don't expect nullable data here - so using a ref might simplify this function body and look cleaner (avoid such enforce, etc.)

Code after:
    }

    plan.x1 = x1;
    plan.x2 = x2;
    plan.y1 = y1;
    plan.y2 = y2;
    plan.w11 = dy2 * dx2;
    plan.w12 = dy2 * dx1;
    plan.w21 = dy1 * dx2;
    plan.w22 = dy1 * dx1;
    plan.mask = mask;
  }
}

template <typename T>
void EvaluatePlanForChannel(const T* input_data,
                            T* output_data,
                            int64_t W_in,
                            const BilinearSamplePlan2D<T>* plan_data,
                            size_t point_count) {
  for (size_t idx = 0; idx < point_count; ++idx) {
    const auto& plan = plan_data[idx];
    if (plan.mask == 0) {
      output_data[idx] = T{};
      continue;
    }

    T p11 = T{};
    T p12 = T{};
    T p21 = T{};
    T p22 = T{};

    if (plan.mask == kAllNeighborsMask) {
      const int64_t row1 = plan.y1 * W_in;
      const int64_t row2 = plan.y2 * W_in;
      p11 = input_data[row1 + plan.x1];
      p12 = input_data[row1 + plan.x2];
      p21 = input_data[row2 + plan.x1];
      p22 = input_data[row2 + plan.x2];
    } else {
      if (plan.mask & kTopLeftMask) {
        p11 = input_data[plan.y1 * W_in + plan.x1];
      }
      if (plan.mask & kTopRightMask) {
        p12 = input_data[plan.y1 * W_in + plan.x2];
      }
      if (plan.mask & kBottomLeftMask) {
        p21 = input_data[plan.y2 * W_in + plan.x1];
      }
      if (plan.mask & kBottomRightMask) {
        p22 = input_data[plan.y2 * W_in + plan.x2];