Function forward_impl

dnn/src/naive/roi_pooling/opr_impl.cpp:13–64

template <typename T, typename Pooler>
void forward_impl(
        _megdnn_tensor_in src, _megdnn_tensor_in rois, _megdnn_tensor_in dst,
        _megdnn_tensor_out index, float spatial_scale) {
    auto C = src.layout.shape[1], IH = src.layout.shape[2], IW = src.layout.shape[3];
    auto OH = dst.layout.shape[2], OW = dst.layout.shape[3];

    auto total_nr_elem = dst.layout.total_nr_elems();
    auto pooled_height = OH, pooled_width = OW;
    auto height = IH, width = IW;
    auto channels = C;
    for (size_t i = 0; i < total_nr_elem; ++i) {
        int pw = i % pooled_width;
        int ph = (i / pooled_width) % pooled_height;
        int c = (i / pooled_width / pooled_height) % channels;
        int n = i / pooled_width / pooled_height / channels;
        auto rois_ptr = rois.ptr<T>() + n * 5;
        int roi_batch_ind = rois_ptr[0];
        int roi_start_w = round(rois_ptr[1] * spatial_scale);
        int roi_start_h = round(rois_ptr[2] * spatial_scale);
        int roi_end_w = round(rois_ptr[3] * spatial_scale);
        int roi_end_h = round(rois_ptr[4] * spatial_scale);
        // Force malformed ROIs to be 1x1
        int roi_width = std::max(roi_end_w - roi_start_w + 1, 1);
        int roi_height = std::max(roi_end_h - roi_start_h + 1, 1);
        float bin_size_h =
                static_cast<float>(roi_height) / static_cast<float>(pooled_height);
        float bin_size_w =
                static_cast<float>(roi_width) / static_cast<float>(pooled_width);

        int hstart = static_cast<int>(floor(static_cast<float>(ph) * bin_size_h));
        int wstart = static_cast<int>(floor(static_cast<float>(pw) * bin_size_w));
        int hend = static_cast<int>(ceil(static_cast<float>(ph + 1) * bin_size_h));
        int wend = static_cast<int>(ceil(static_cast<float>(pw + 1) * bin_size_w));
        // Add roi offsets and clip to input boundaries
        hstart = std::min<int>(std::max(hstart + roi_start_h, 0), height);
        hend = std::min<int>(std::max(hend + roi_start_h, 0), height);
        wstart = std::min<int>(std::max(wstart + roi_start_w, 0), width);
        wend = std::min<int>(std::max(wend + roi_start_w, 0), width);

        Pooler pooler;
        auto feat_map_ptr =
                src.ptr<T>() + (roi_batch_ind * channels + c) * height * width;
        for (int h = hstart; h < hend; ++h) {
            for (int w = wstart; w < wend; ++w) {
                int bottom_i = h * width + w;
                pooler.feed(feat_map_ptr[bottom_i], bottom_i);
            }
        }
        pooler.writeback_val(dst.ptr<T>()[i]);
        pooler.writeback_idx(index.ptr<dt_int32>()[i]);
    }
}

template <typename T>
void forward(