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

Code before:

                  T dx2 = static_cast<T>(x2) - x;
                  T dx1 = x - static_cast<T>(x1);
                  T dy2 = static_cast<T>(y2) - y;
                  T dy1 = y - static_cast<T>(y1);
                  *Y_gridpoint = dy2 * (dx2 * p11 + dx1 * p12) + dy1 * (dx2 * p21 + dx1 * p22);
                } else if (mode_ == Cubic) {
                  int64_t x0 = static_cast<int64_t>(std::floor(x)) - 1;  // top-left corner of the bbox
                  int64_t y0 = static_cast<int64_t>(std::floor(y)) - 1;

                  T p[4][4] = {};  // [H][W]
                  for (int64_t h = 0; h < 4; h++) {
                    for (int64_t w = 0; w < 4; w++) {
                      p[h][w] = PixelAtGrid(X_data, h + y0, w + x0, H_in, W_in, border);
                    }
                  }
                  T dx = static_cast<T>(x - x0 - 1);
                  T dy = static_cast<T>(y - y0 - 1);
                  *Y_gridpoint = GsBicubicInterpolate(p, dx, dy);
                }
              }
            }
          });
    };

    const bool can_use_fast_path = (mode_ == Linear && padding_mode_ == Zeros && !align_corners_);
    if (can_use_fast_path) {
      std::vector<BilinearSamplePlan2D<T>> sampling_plan;
      for (int64_t n = 0; n < N; n++) {
        // Choose fast path when all 4 neighbors are within the image and use zero for out-of-boundary neighbors.
        // This fast path can be 2-3x faster than the generic path with boundary check and supports Neon optimization.
        // sampling_plan helps precomputing a separate plan entry per output pixel.
        TryRunBilinearZerosFastPath2D(*input, *grid, Y, n, C, H_in, W_in, H_out, W_out, tp, sampling_plan);
      }
    } else {
      for (int64_t n = 0; n < N; n++) {
        run_generic_path_for_n(n);
      }
    }
  } else if (data_dims == 3) {
    // sample 3d;
    auto D_in = input_dims[2];
    auto H_in = input_dims[3];
    auto W_in = input_dims[4];
    auto D_out = grid_dims[1];
    auto H_out = grid_dims[2];
    auto W_out = grid_dims[3];
    TensorShape Y_shape = {N, C, D_out, H_out, W_out};
    auto& Y = *context->Output(0, Y_shape);
    // Return early if the output tensor is going to be of size 0
    if (Y.Shape().Size() == 0) {

Review comment (suggestion):
Not sure why we need the changes in the generic path ? Can it just be kept as is (i.e.) move the fast path logic above the generic path and bail out of Compute() if the fast path is taken. This means we can avoid the change to make the generic path a lambda and call it in another loop in the else part. 

This is just a nit. I am just trying to keep code change minimal...

Code after:
                    }
                    T dx = static_cast<T>(x - x0 - 1);
                    T dy = static_cast<T>(y - y0 - 1);
                    *Y_gridpoint = GsBicubicInterpolate(p, dx, dy);
                  }
                }
              }
            });
      }
    }
  } else if (data_dims == 3) {
    // sample 3d;
    auto D_in = input_dims[2];
    auto H_in = input_dims[3];
    auto W_in = input_dims[4];
    auto D_out = grid_dims[1];
    auto H_out = grid_dims[2];
    auto W_out = grid_dims[3];
    TensorShape Y_shape = {N, C, D_out, H_out, W_out};
    auto& Y = *context->Output(0, Y_shape);
    // Return early if the output tensor is going to be of size 0
    if (Y.Shape().Size() == 0) {
      return Status::OK();
    }

    T x_min = -0.5f;
    T x_max = W_in - 0.5f;
    T y_min = -0.5f;
    T y_max = H_in - 0.5f;
    T z_min = -0.5f;
    T z_max = D_in - 0.5f;

    if (align_corners_) {
      x_min = 0.f;
      x_max = W_in - 1.f;
      y_min = 0.f;
      y_max = H_in - 1.f;
      z_min = 0.f;
      z_max = D_in - 1.f;
    }
    T border[] = {x_min, y_min, z_min, x_max, y_max, z_max};

    concurrency::ThreadPool* tp = D_out * H_out * W_out > 64 ? context->GetOperatorThreadPool() : nullptr;
    for (int64_t n = 0; n < N; n++) {
      const T* grid_data = grid->Data<T>() + n * (D_out * H_out * W_out) * 3;
      concurrency::ThreadPool::TrySimpleParallelFor(
          tp, onnxruntime::narrow<std::ptrdiff_t>(C),
          [&](std::ptrdiff_t c) {
            const T* X_data = input->Data<T>() + (n * C + c) * (D_in * H_in * W_in);
            T* Y_data = Y.MutableData<T>() + (n * C + c) * (D_out * H_out * W_out);