From c2255c36ec121fdb998ce3db8deb7508c814b567 Mon Sep 17 00:00:00 2001 From: Richard Barnes Date: Mon, 21 Feb 2022 13:49:19 -0800 Subject: [PATCH] Fix binary search in bisect_percentile_op (#73146) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/73146 Binary search can overflow; this fixes it. Test Plan: Sandcastle Reviewed By: meyering Differential Revision: D34365186 fbshipit-source-id: f92a810b49ef5ce345d0b019b584fe3c1f5ae017 (cherry picked from commit 9c2133ec6f3ae9874fd6ae2ffb2a98bb7b68a6b0) --- caffe2/operators/bisect_percentile_op.h | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/caffe2/operators/bisect_percentile_op.h b/caffe2/operators/bisect_percentile_op.h index 0b5567a776c..2c2122c884c 100644 --- a/caffe2/operators/bisect_percentile_op.h +++ b/caffe2/operators/bisect_percentile_op.h @@ -115,13 +115,10 @@ class BisectPercentileOp final : public Operator { int lo, int hi, float val) { - int mid; - bool low_cond, high_cond; - while (lo < hi) { - mid = (lo + hi) >> 1; - low_cond = (data[mid] <= val); - high_cond = (val < data[mid + 1]); + const auto mid = lo + (hi - lo) / 2; + const bool low_cond = (data[mid] <= val); + const bool high_cond = (val < data[mid + 1]); if (low_cond && high_cond) { return mid; } else if (!low_cond) {