Use iterators and fix basic_iota_iterator

src/include/migraphx/iota_iterator.hpp

@@ -81,8 +81,9 @@ struct basic_iota_iterator
         index--;
         return it;
     }
-    // TODO: operator->
     reference operator*() const { return f(index); }
+    pointer operator->() const { return &f(index); }
+    reference operator[](int n) const { return f(index + n); }
 };
 
 template <class T, class F>

src/include/migraphx/op/nonmaxsuppression.hpp

@@ -81,8 +81,8 @@ struct nonmaxsuppression
 
     struct box
     {
-        std::array<float, 2> x;
-        std::array<float, 2> y;
+        std::array<double, 2> x;
+        std::array<double, 2> y;
 
         void sort()
         {
@@ -90,9 +90,9 @@ struct nonmaxsuppression
             std::sort(y.begin(), y.end());
         }
 
-        std::array<float, 2>& operator[](std::size_t i) { return i == 0 ? x : y; }
+        std::array<double, 2>& operator[](std::size_t i) { return i == 0 ? x : y; }
 
-        float area() const
+        double area() const
         {
             assert(std::is_sorted(x.begin(), x.end()));
             assert(std::is_sorted(y.begin(), y.end()));
@@ -101,29 +101,29 @@ struct nonmaxsuppression
     };
 
     template <class T>
-    box batch_box(const T& boxes, std::size_t box_ind, std::size_t box_idx) const
+    box batch_box(T boxes, std::size_t box_idx) const
     {
         box result{};
-        auto start = box_ind + 4 * box_idx;
+        auto start = boxes + 4 * box_idx;
         if(center_point_box)
         {
-            float half_width  = boxes[start + 2] / 2.0;
-            float half_height = boxes[start + 3] / 2.0;
-            float x_center    = boxes[start + 0];
-            float y_center    = boxes[start + 1];
+            double half_width  = start[2] / 2.0;
+            double half_height = start[3] / 2.0;
+            double x_center    = start[0];
+            double y_center    = start[1];
             result.x          = {x_center - half_width, x_center + half_width};
             result.y          = {y_center - half_height, y_center + half_height};
         }
         else
         {
-            result.x = {static_cast<float>(boxes[start + 1]), static_cast<float>(boxes[start + 3])};
-            result.y = {static_cast<float>(boxes[start + 0]), static_cast<float>(boxes[start + 2])};
+            result.x = {static_cast<double>(start[1]), static_cast<double>(start[3])};
+            result.y = {static_cast<double>(start[0]), static_cast<double>(start[2])};
         }
 
         return result;
     }
 
-    inline bool suppress_by_iou(box b1, box b2, float iou_threshold) const
+    inline bool suppress_by_iou(box b1, box b2, double iou_threshold) const
     {
         b1.sort();
         b2.sort();
@@ -135,7 +135,7 @@ struct nonmaxsuppression
             intersection[i][1] = std::min(b1[i][1], b2[i][1]);
         }
 
-        std::vector<std::array<float, 2>> bbox = {intersection.x, intersection.y};
+        std::vector<std::array<double, 2>> bbox = {intersection.x, intersection.y};
         if(std::any_of(bbox.begin(), bbox.end(), [](auto bx) {
                return not std::is_sorted(bx.begin(), bx.end());
            }))
@@ -143,33 +143,33 @@ struct nonmaxsuppression
             return false;
         }
 
-        const float area1             = b1.area();
-        const float area2             = b2.area();
-        const float intersection_area = intersection.area();
-        const float union_area        = area1 + area2 - intersection_area;
+        const double area1             = b1.area();
+        const double area2             = b2.area();
+        const double intersection_area = intersection.area();
+        const double union_area        = area1 + area2 - intersection_area;
 
         if(area1 <= .0f or area2 <= .0f or union_area <= .0f)
         {
             return false;
         }
 
-        const float intersection_over_union = intersection_area / union_area;
+        const double intersection_over_union = intersection_area / union_area;
 
         return intersection_over_union > iou_threshold;
     }
-
+    
     // filter boxes below score_threshold
     template <class T>
-    std::priority_queue<std::pair<float, int64_t>> filter_boxes_by_score(
-        T scores, std::size_t score_offset_ind, std::size_t num_boxes, float score_threshold) const
+    std::priority_queue<std::pair<double, int64_t>> filter_boxes_by_score(
+        T scores_start, std::size_t num_boxes, double score_threshold) const
     {
-        std::priority_queue<std::pair<float, int64_t>> boxes_heap;
+        std::priority_queue<std::pair<double, int64_t>> boxes_heap;
         auto insert_to_boxes_heap =
             make_function_output_iterator([&](const auto& x) { boxes_heap.push(x); });
         int64_t box_idx = 0;
         transform_if(
-            scores.begin() + score_offset_ind,
-            scores.begin() + score_offset_ind + num_boxes,
+            scores_start,
+            scores_start + num_boxes,
             insert_to_boxes_heap,
             [&](auto sc) {
                 box_idx++;
@@ -178,7 +178,48 @@ struct nonmaxsuppression
             [&](auto sc) { return std::make_pair(sc, box_idx - 1); });
         return boxes_heap;
     }
+    
+    template <class H, class S>
+    void select_boxes(
+        H& boxes_heap,
+        std::vector<std::pair<double, int64_t>>& selected_boxes_inside_class,
+        std::vector<int64_t>& selected_indices,
+        S batch_boxes_start,
+        std::size_t max_output_boxes_per_class,
+        double iou_threshold,
+        std::size_t batch_idx,
+        std::size_t class_idx
+    ) const
+    {
+        selected_boxes_inside_class.clear();
+        // Get the next box with top score, filter by iou_threshold
+        while(!boxes_heap.empty() &&
+                selected_boxes_inside_class.size() < max_output_boxes_per_class)
+        {
+            // Check with existing selected boxes for this class, remove box if it
+            // exceeds the IOU (Intersection Over Union) threshold
+            const auto next_top_score = boxes_heap.top();
+            bool not_selected         = std::any_of(
+                    selected_boxes_inside_class.begin(),
+                    selected_boxes_inside_class.end(),
+                    [&](auto selected_index) {
+                    return this->suppress_by_iou(
+                            batch_box(batch_boxes_start, next_top_score.second),
+                            batch_box(batch_boxes_start, selected_index.second),
+                            iou_threshold);
+                    });
 
+            if(not not_selected)
+            {
+                selected_boxes_inside_class.push_back(next_top_score);
+                selected_indices.push_back(batch_idx);
+                selected_indices.push_back(class_idx);
+                selected_indices.push_back(next_top_score.second);
+            }
+            boxes_heap.pop();
+        }
+    }
+    
     argument compute(const shape& output_shape, std::vector<argument> args) const
     {
         argument result{output_shape};
@@ -189,61 +230,34 @@ struct nonmaxsuppression
         {
             return result;
         }
-        float iou_threshold   = (args.size() > 3) ? (args.at(3).at<float>()) : 0.0f;
-        float score_threshold = (args.size() > 4) ? (args.at(4).at<float>()) : 0.0f;
+        double iou_threshold   = (args.size() > 3) ? (args.at(3).at<double>()) : 0.0f;
+        double score_threshold = (args.size() > 4) ? (args.at(4).at<double>()) : 0.0f;
 
         result.visit([&](auto output) {
             visit_all(args[0], args[1])([&](auto boxes, auto scores) {
-                std::fill(output.begin(), output.end(), 0);
-                const auto& lens       = scores.get_shape().lens();
-                const auto num_batches = lens[0];
-                const auto num_classes = lens[1];
-                const auto num_boxes   = lens[2];
-                // boxes of a class with NMS applied [score, index]
-                std::vector<std::pair<float, int64_t>> selected_boxes_inside_class;
-                std::vector<int64_t> selected_indices;
-                selected_boxes_inside_class.reserve(output_shape.elements());
-                // iterate over batches and classes
-                shape comp_s{shape::float_type, {num_batches, num_classes}};
-                shape_for_each(comp_s, [&](auto idx) {
-                    auto batch_idx = idx[0];
-                    auto class_idx = idx[1];
-                    // index offset for this class
-                    std::size_t score_offset_ind =
-                        (batch_idx * num_classes + class_idx) * num_boxes;
-                    // index to first value of this batch
-                    std::size_t batch_boxes_ind = batch_idx * num_boxes * 4;
-                    auto boxes_heap =
-                        filter_boxes_by_score(scores, score_offset_ind, num_boxes, score_threshold);
-                    selected_boxes_inside_class.clear();
-                    // Get the next box with top score, filter by iou_threshold
-                    while(!boxes_heap.empty() &&
-                          selected_boxes_inside_class.size() < max_output_boxes_per_class)
-                    {
-                        // Check with existing selected boxes for this class, remove box if it
-                        // exceeds the IOU (Intersection Over Union) threshold
-                        const auto next_top_score = boxes_heap.top();
-                        bool not_selected         = std::any_of(
-                            selected_boxes_inside_class.begin(),
-                            selected_boxes_inside_class.end(),
-                            [&](auto selected_index) {
-                                return this->suppress_by_iou(
-                                    batch_box(boxes, batch_boxes_ind, next_top_score.second),
-                                    batch_box(boxes, batch_boxes_ind, selected_index.second),
-                                    iou_threshold);
-                            });
-
-                        if(not not_selected)
-                        {
-                            selected_boxes_inside_class.push_back(next_top_score);
-                            selected_indices.push_back(batch_idx);
-                            selected_indices.push_back(class_idx);
-                            selected_indices.push_back(next_top_score.second);
-                        }
-                        boxes_heap.pop();
-                    }
-                });
-                std::copy(selected_indices.begin(), selected_indices.end(), output.begin());
+                    std::fill(output.begin(), output.end(), 0);
+                    const auto& lens       = scores.get_shape().lens();
+                    const auto num_batches = lens[0];
+                    const auto num_classes = lens[1];
+                    const auto num_boxes   = lens[2];
+                    // boxes of a class with NMS applied [score, index]
+                    std::vector<std::pair<double, int64_t>> selected_boxes_inside_class;
+                    std::vector<int64_t> selected_indices;
+                    selected_boxes_inside_class.reserve(output_shape.elements());
+                    // iterate over batches and classes
+                    shape comp_s{shape::double_type, {num_batches, num_classes}};
+                    shape_for_each(comp_s, [&](auto idx) {
+                            auto batch_idx = idx[0];
+                            auto class_idx = idx[1];
+                            // index offset for this class
+                            auto scores_start = scores.begin() + (batch_idx * num_classes + class_idx) * num_boxes;
+                            // iterator to first value of this batch
+                            auto batch_boxes_start = boxes.begin() + batch_idx * num_boxes * 4;
+                            auto boxes_heap =
+                            filter_boxes_by_score(scores_start, num_boxes, score_threshold);
+                            select_boxes(boxes_heap, selected_boxes_inside_class, selected_indices, batch_boxes_start, max_output_boxes_per_class, iou_threshold, batch_idx, class_idx);
+                    });
+                    std::copy(selected_indices.begin(), selected_indices.end(), output.begin());
             });
         });