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

Code before:
      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,
                                   int64_t W_in,
                                   int64_t H_out,
                                   int64_t W_out,
                                   concurrency::ThreadPool* tp,
                                   std::vector<BilinearSamplePlan2D<T>>& sampling_plan) {
  const int64_t plane_in = H_in * W_in;
  const size_t plane_out = H_out * W_out;
  sampling_plan.resize(plane_out);

  const T* grid_data = grid.Data<T>() + n * plane_out * 2;
  PrecomputeBilinearSamplePlan2D(grid_data, H_out, W_out, H_in, W_in, sampling_plan);

  const T* input_data = input.Data<T>();
  T* output_data = output.MutableData<T>();

  if (plane_out == 0) {
    return;
  }

  concurrency::ThreadPool::TrySimpleParallelFor(
      tp, onnxruntime::narrow<std::ptrdiff_t>(C),
      [&](std::ptrdiff_t c) {
        const T* X_data = input_data + (n * C + c) * plane_in;
        T* Y_data = output_data + (n * C + c) * plane_out;
        EvaluatePlanForChannel(X_data, Y_data, W_in, sampling_plan.data(), plane_out);
      });
}

}  // namespace

// When grid sampling, padding is applied before interpolation.
// For instance, in bilinear mode and zeros padding-mode, pixel p at actual

Review comment (suggestion):
```suggestion
  const size_t plane_out = static_cast<size_t>(H_out) * static_cast<size_t>(W_out);
```

Code after:
      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,
                                   int64_t W_in,
                                   int64_t H_out,
                                   int64_t W_out,
                                   concurrency::ThreadPool* tp,
                                   std::vector<BilinearSamplePlan2D<T>>& sampling_plan) {
  const size_t plane_in = static_cast<size_t>(H_in) * static_cast<size_t>(W_in);
  const size_t plane_out = static_cast<size_t>(H_out) * static_cast<size_t>(W_out);
  sampling_plan.resize(plane_out);

  const T* grid_data = grid.Data<T>() + n * plane_out * 2;
  PrecomputeBilinearSamplePlan2D(grid_data, H_out, W_out, H_in, W_in, sampling_plan);

  const T* input_data = input.Data<T>();
  T* output_data = output.MutableData<T>();

  if (plane_out == 0) {
    return;
  }

  concurrency::ThreadPool::TrySimpleParallelFor(
      tp, onnxruntime::narrow<std::ptrdiff_t>(C),
      [&](std::ptrdiff_t c) {
        const T* X_data = input_data + (n * C + c) * plane_in;
        T* Y_data = output_data + (n * C + c) * plane_out;
        EvaluatePlanForChannel(X_data, Y_data, W_in, sampling_plan.data(), plane_out);
      });
}

}  // namespace

// When grid sampling, padding is applied before interpolation.
// For instance, in bilinear mode and zeros padding-mode, pixel p at actual