chore/faster-ocl-bfmatcher-crosscheck

The existing 2-pass cross-check implementation requires both forward and reverse distance computations on GPU, followed by CPU-side filtering. This adds a single-pass kernel using int64 atomics to eliminate the reverse pass. The BruteForceMatch_CrossCheckMatch kernel performs forward matching and tracks best train->query mappings in one GPU pass, reducing memory transfers. Falls back to 2-pass approach on devices without cl_khr_int64_base_atomics extension. Tested on NVIDIA RTX 5060 Ti: 1.65x-1.97x faster than 2-pass GPU, 3.49x-7.33x faster than CPU at 640x480 to 1920x1080 resolutions.
2026-07-21 19:33:03 +04:00 · 2026-04-18 13:29:41 +05:30
parent a5295015b2
commit ab0def4854
2 changed files with 248 additions and 13 deletions
@@ -74,7 +74,7 @@ static void ensureSizeIsEnough(int rows, int cols, int type, UMat &m)
 }

 static bool ocl_matchSingle(InputArray query, InputArray train,
-        UMat &trainIdx, UMat &distance, int distType)
+        UMat &trainIdx, UMat *distance, int distType)
 {
    if (query.empty() || train.empty())
        return false;
@@ -83,7 +83,8 @@ static bool ocl_matchSingle(InputArray query, InputArray train,
    const int query_cols = query.cols();

    ensureSizeIsEnough(1, query_rows, CV_32S, trainIdx);
-    ensureSizeIsEnough(1, query_rows, CV_32F, distance);
+    if (distance)
+        ensureSizeIsEnough(1, query_rows, CV_32F, *distance);

    ocl::Device devDef = ocl::Device::getDefault();

@@ -117,7 +118,10 @@ static bool ocl_matchSingle(InputArray query, InputArray train,
    idx = k.set(idx, ocl::KernelArg::PtrReadOnly(uquery));
    idx = k.set(idx, ocl::KernelArg::PtrReadOnly(utrain));
    idx = k.set(idx, ocl::KernelArg::PtrWriteOnly(trainIdx));
-    idx = k.set(idx, ocl::KernelArg::PtrWriteOnly(distance));
+    UMat dummyDist;
+    if (!distance)
+        ensureSizeIsEnough(1, query_rows, CV_32F, dummyDist);
+    idx = k.set(idx, ocl::KernelArg::PtrWriteOnly(distance ? *distance : dummyDist));
    idx = k.set(idx, uquery.rows);
    idx = k.set(idx, uquery.cols);
    idx = k.set(idx, utrain.rows);
@@ -734,7 +738,7 @@ Ptr<DescriptorMatcher> BFMatcher::clone( bool emptyTrainData ) const
 static bool ocl_match(InputArray query, InputArray _train, std::vector< std::vector<DMatch> > &matches, int dstType)
 {
    UMat trainIdx, distance;
-    if (!ocl_matchSingle(query, _train, trainIdx, distance, dstType))
+    if (!ocl_matchSingle(query, _train, trainIdx, &distance, dstType))
        return false;
    if (!ocl_matchDownload(trainIdx, distance, matches))
        return false;
@@ -753,22 +757,119 @@ static bool ocl_knnMatch(InputArray query, InputArray _train, std::vector< std::
    return true;
 }

-// Run match in both directions on GPU, cross-check filter on CPU.
-// Reuses ocl_matchSingle so no new kernel is needed.
+static bool ocl_matchWithCrossCheckSinglePass(InputArray query, InputArray train,
+        std::vector< std::vector<DMatch> >& matches, int dstType)
+{
+    if (query.empty() || train.empty())
+        return false;
+
+    const int query_rows = query.rows();
+    const int train_rows = train.rows();
+
+    UMat fwdIdx, fwdDist;
+    ensureSizeIsEnough(1, query_rows, CV_32S, fwdIdx);
+    ensureSizeIsEnough(1, query_rows, CV_32F, fwdDist);
+
+    UMat revBest;
+    ensureSizeIsEnough(1, train_rows, CV_32SC2, revBest);
+    revBest.setTo(Scalar::all(-1));
+
+    ocl::Device devDef = ocl::Device::getDefault();
+
+    UMat uquery = query.getUMat(), utrain = train.getUMat();
+    int kercn = 1;
+    if (devDef.isIntel() &&
+        (0 == (uquery.step % 4)) && (0 == (uquery.cols % 4)) && (0 == (uquery.offset % 4)) &&
+        (0 == (utrain.step % 4)) && (0 == (utrain.cols % 4)) && (0 == (utrain.offset % 4)))
+        kercn = 4;
+
+    int block_size = 16;
+    int max_desc_len = 0;
+    bool is_cpu = devDef.type() == ocl::Device::TYPE_CPU;
+    if (uquery.cols <= 64)
+        max_desc_len = 64 / kercn;
+    else if (uquery.cols <= 128 && !is_cpu)
+        max_desc_len = 128 / kercn;
+
+    int depth = uquery.depth();
+    cv::String opts;
+    opts = cv::format("-D T=%s -D TN=%s -D kercn=%d %s -D DIST_TYPE=%d -D BLOCK_SIZE=%d -D MAX_DESC_LEN=%d -D HAVE_INT64_ATOMICS",
+        ocl::typeToStr(depth), ocl::typeToStr(CV_MAKETYPE(depth, kercn)), kercn, depth == CV_32F ? "-D T_FLOAT" : "", dstType, block_size, max_desc_len);
+    ocl::Kernel k("BruteForceMatch_CrossCheckMatch", ocl::features2d::brute_force_match_oclsrc, opts);
+    if(k.empty())
+        return false;
+
+    size_t globalSize[] = {((size_t)uquery.size().height + block_size - 1) / block_size * block_size, (size_t)block_size};
+    size_t localSize[] = {(size_t)block_size, (size_t)block_size};
+
+    int idx = 0;
+    idx = k.set(idx, ocl::KernelArg::PtrReadOnly(uquery));
+    idx = k.set(idx, ocl::KernelArg::PtrReadOnly(utrain));
+    idx = k.set(idx, ocl::KernelArg::PtrWriteOnly(fwdIdx));
+    idx = k.set(idx, ocl::KernelArg::PtrWriteOnly(fwdDist));
+    idx = k.set(idx, ocl::KernelArg::PtrWriteOnly(revBest));
+    idx = k.set(idx, uquery.rows);
+    idx = k.set(idx, uquery.cols);
+    idx = k.set(idx, utrain.rows);
+    idx = k.set(idx, utrain.cols);
+    idx = k.set(idx, (int)(uquery.step / sizeof(float)));
+
+    if (!k.run(2, globalSize, localSize, false))
+        return false;
+
+    Mat fwdIdxCPU = fwdIdx.getMat(ACCESS_READ);
+    Mat fwdDistCPU = fwdDist.getMat(ACCESS_READ);
+    Mat revBestCPU = revBest.getMat(ACCESS_READ);
+
+    if (fwdIdxCPU.empty() || fwdDistCPU.empty() || revBestCPU.empty())
+        return false;
+
+    const int* fwd = fwdIdxCPU.ptr<int>();
+    const float* dist = fwdDistCPU.ptr<float>();
+    const int* revBestData = revBestCPU.ptr<int>();
+
+    matches.clear();
+    matches.reserve(query_rows);
+
+    for (int q = 0; q < query_rows; ++q)
+    {
+        int t = fwd[q];
+        if (t >= 0 && t < train_rows)
+        {
+            uint64_t packed_val;
+            memcpy(&packed_val, &revBestData[2 * t], sizeof(uint64_t));
+            int revQueryIdx = (int)(packed_val & 0xFFFFFFFF);
+            if (revQueryIdx == q)
+            {
+                matches.push_back(std::vector<DMatch>(1, DMatch(q, t, 0, dist[q])));
+            }
+        }
+    }
+    return true;
+}
+
 static bool ocl_matchWithCrossCheck(InputArray query, InputArray train,
        std::vector< std::vector<DMatch> >& matches, int dstType)
 {
-    // Forward pass: for each query descriptor find nearest in train
+    if (query.empty() || train.empty())
+        return false;
+
+    ocl::Device devDef = ocl::Device::getDefault();
+
+    if (devDef.isExtensionSupported("cl_khr_int64_base_atomics"))
+    {
+        if (ocl_matchWithCrossCheckSinglePass(query, train, matches, dstType))
+            return true;
+    }
+
    UMat fwdIdx, fwdDist;
-    if (!ocl_matchSingle(query, train, fwdIdx, fwdDist, dstType))
+    if (!ocl_matchSingle(query, train, fwdIdx, &fwdDist, dstType))
        return false;

-    // Reverse pass: for each train descriptor find nearest in query
-    UMat revIdx, revDist;
-    if (!ocl_matchSingle(train, query, revIdx, revDist, dstType))
+    UMat revIdx;
+    if (!ocl_matchSingle(train, query, revIdx, nullptr, dstType))
        return false;

-    // Download index arrays (1 x N ints each — cheap)
    Mat fwdIdxCPU = fwdIdx.getMat(ACCESS_READ);
    Mat revIdxCPU = revIdx.getMat(ACCESS_READ);
    Mat fwdDistCPU = fwdDist.getMat(ACCESS_READ);
@@ -557,4 +557,138 @@ __kernel void BruteForceMatch_knnMatch(
        bestTrainIdx[queryIdx] = (int2)(myBestTrainIdx1, myBestTrainIdx2);
        bestDistance[queryIdx] = (float2)(myBestDistance1, myBestDistance2);
    }
-}
+}
+
+#ifdef HAVE_INT64_ATOMICS
+#pragma OPENCL EXTENSION cl_khr_int64_base_atomics : enable
+
+__kernel void BruteForceMatch_CrossCheckMatch(
+    __global T *query,
+    __global T *train,
+    __global int *bestTrainIdx,
+    __global float *bestDistance,
+    __global ulong *revBest,
+    int query_rows,
+    int query_cols,
+    int train_rows,
+    int train_cols,
+    int step
+)
+{
+    const int lidx = get_local_id(0);
+    const int lidy = get_local_id(1);
+    const int groupidx = get_group_id(0);
+
+    const int queryIdx = mad24(BLOCK_SIZE, groupidx, lidy);
+    const int queryOffset = min(queryIdx, query_rows - 1) * step;
+    __global TN *query_vec = (__global TN *)(query + queryOffset);
+    query_cols /= kercn;
+
+    __local float sharebuffer[SHARED_MEM_SZ];
+    __local value_type *s_query = (__local value_type *)sharebuffer;
+
+#if 0 < MAX_DESC_LEN
+    __local value_type *s_train = (__local value_type *)sharebuffer + BLOCK_SIZE * MAX_DESC_LEN;
+    #pragma unroll
+    for (int i = 0; i < MAX_DESC_LEN / BLOCK_SIZE; i++)
+    {
+        const int loadx = mad24(BLOCK_SIZE, i, lidx);
+        s_query[mad24(MAX_DESC_LEN, lidy, loadx)] = loadx < query_cols ? query_vec[loadx] : 0;
+    }
+#else
+    __local value_type *s_train = (__local value_type *)sharebuffer + BLOCK_SIZE_ODD * BLOCK_SIZE;
+    const int s_query_i = mad24(BLOCK_SIZE_ODD, lidy, lidx);
+    const int s_train_i = mad24(BLOCK_SIZE_ODD, lidx, lidy);
+#endif
+
+    float myBestDistance = MAX_FLOAT;
+    int myBestTrainIdx = -1;
+
+    for (int t = 0, endt = (train_rows + BLOCK_SIZE - 1) / BLOCK_SIZE; t < endt; t++)
+    {
+        result_type result = 0;
+
+        const int trainOffset = min(mad24(BLOCK_SIZE, t, lidy), train_rows - 1) * step;
+        __global TN *train_vec = (__global TN *)(train + trainOffset);
+#if 0 < MAX_DESC_LEN
+        #pragma unroll
+        for (int i = 0; i < MAX_DESC_LEN / BLOCK_SIZE; i++)
+        {
+            const int loadx = mad24(BLOCK_SIZE, i, lidx);
+            s_train[mad24(BLOCK_SIZE, lidx, lidy)] = loadx < train_cols ? train_vec[loadx] : 0;
+
+            barrier(CLK_LOCAL_MEM_FENCE);
+
+            result += reduce_multi_block(s_query, s_train, i, lidx, lidy);
+
+            barrier(CLK_LOCAL_MEM_FENCE);
+        }
+#else
+        for (int i = 0, endq = (query_cols + BLOCK_SIZE - 1) / BLOCK_SIZE; i < endq; i++)
+        {
+            const int loadx = mad24(BLOCK_SIZE, i, lidx);
+            if (loadx < query_cols)
+            {
+                s_query[s_query_i] = query_vec[loadx];
+                s_train[s_train_i] = train_vec[loadx];
+            }
+            else
+            {
+                s_query[s_query_i] = 0;
+                s_train[s_train_i] = 0;
+            }
+
+            barrier(CLK_LOCAL_MEM_FENCE);
+
+            result += reduce_block_match(s_query, s_train, lidx, lidy);
+
+            barrier(CLK_LOCAL_MEM_FENCE);
+        }
+#endif
+        result = DIST_RES(result);
+
+        const int trainIdx = mad24(BLOCK_SIZE, t, lidx);
+
+        if (queryIdx < query_rows && trainIdx < train_rows)
+        {
+            if (result < myBestDistance)
+            {
+                myBestDistance = result;
+                myBestTrainIdx = trainIdx;
+            }
+
+            uint dist_bits = as_uint(result);
+            ulong packed = ((ulong)dist_bits << 32) | (ulong)queryIdx;
+            atom_min(&revBest[trainIdx], packed);
+        }
+    }
+
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    __local float *s_distance = (__local float *)sharebuffer;
+    __local int *s_trainIdx = (__local int *)(sharebuffer + BLOCK_SIZE_ODD * BLOCK_SIZE);
+
+    s_distance += lidy * BLOCK_SIZE_ODD;
+    s_trainIdx += lidy * BLOCK_SIZE_ODD;
+    s_distance[lidx] = myBestDistance;
+    s_trainIdx[lidx] = myBestTrainIdx;
+
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    #pragma unroll
+    for (int k = 0; k < BLOCK_SIZE; k++)
+    {
+        if (myBestDistance > s_distance[k])
+        {
+            myBestDistance = s_distance[k];
+            myBestTrainIdx = s_trainIdx[k];
+        }
+    }
+
+    if (queryIdx < query_rows && lidx == 0)
+    {
+        bestTrainIdx[queryIdx] = myBestTrainIdx;
+        bestDistance[queryIdx] = myBestDistance;
+    }
+}
+#endif