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

Code before:
    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));
        float32x2_t accum = vpadd_f32(sum_pairs, sum_pairs);
        output_data[idx] = vget_lane_f32(accum, 0);
        continue;
      }
    }
#endif

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

    if (plan.mask == kAllNeighborsMask) {
      const int64_t row1 = plan.y1 * W_in;

Review comment (refactor):
Spoke offline with @edgchen1 - We both feel that the usage of intrinsics and platform specific code is better off in MLAS.

We have 2 possible approaches forward:

1) Add an MLAS API to perform the "specialized" dot product operation performed here (i.e.) read 2 values from one array + 2 values from another array + 4 weights and perform the dot product and move all this code to its implementation. We could do this in the same PR.

2) Remove this code for this PR and add the MLAS change in a separate PR later.

Which way forward depends on the perf gains from this vectorization logic. Can you please share that for your target workload ? if it isn't much gains, can we take approach 2 and think about adding this as a future enhancement when there is an appropriate workload to take advantage of this ? 

 

Code after:
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];
      }
    }

    output_data[idx] = plan.w11 * p11 + plan.w12 * p12 + plan.w21 * p21 + plan.w22 * p22;
  }
}

template <typename T>
void TryRunBilinearZerosFastPath2D(const Tensor& input,
                                   const Tensor& grid,
                                   Tensor& output,
                                   int64_t n,
                                   int64_t C,
                                   int64_t H_in,