Merge pull request #28897 from Lurie97:fix_diag

core: implement OpenCL kernel for UMat::diag to avoid aliasing races

modules/core/src/opencl/copyset.cl

@@ -98,6 +98,50 @@ __kernel void copyToMask(__global const uchar * srcptr, int src_step, int src_of
     }
 }
 
+#elif defined(SET_DIAG)
+
+#ifdef cl_khr_fp64
+#pragma OPENCL EXTENSION cl_khr_fp64 : enable
+#endif
+
+__kernel void setDiag(__global uchar* dst, int dst_step, int dst_offset,
+                      int dst_rows, int dst_cols,
+                      __global const uchar* src, int src_step, int src_offset,
+                      int len)
+{
+    int x = get_global_id(0);
+    int y = get_global_id(1);
+
+    if (x < dst_cols && y < dst_rows)
+    {
+        int elem_size = (int)sizeof(T1) * cn;
+        int dst_idx = mad24(y, dst_step, mad24(x, elem_size, dst_offset));
+
+        if (x == y && x < len)
+        {
+#if IS_ROW_VECTOR
+            int src_idx = src_offset + x * elem_size;
+#else
+            int src_idx = mad24(x, src_step, src_offset);
+#endif
+            __global T1* dst_ptr = (__global T1*)(dst + dst_idx);
+            __global const T1* src_ptr = (__global const T1*)(src + src_idx);
+
+            #pragma unroll
+            for (int c = 0; c < cn; ++c)
+                dst_ptr[c] = src_ptr[c];
+        }
+        else
+        {
+            __global T1* dst_ptr = (__global T1*)(dst + dst_idx);
+
+            #pragma unroll
+            for (int c = 0; c < cn; ++c)
+                dst_ptr[c] = (T1)0;
+        }
+    }
+}
+
 #else
 
 #ifndef dstST

modules/core/src/umatrix.cpp

@@ -993,10 +993,54 @@ UMat UMat::reshape(int new_cn, int new_rows) const
     return hdr;
 }
 
+#ifdef HAVE_OPENCL
+namespace {
+static bool ocl_setDiag(const UMat& d, UMat& m, int len)
+{
+    int cn = d.channels();
+    int depth = d.depth();
+
+    if (depth == CV_64F && !ocl::Device::getDefault().doubleFPConfig())
+        return false;
+
+    String opts = format("-D SET_DIAG -D T1=%s -D cn=%d -D IS_ROW_VECTOR=%d",
+                         ocl::memopTypeToStr(depth),
+                         cn,
+                         (d.rows == 1) ? 1 : 0);
+
+    ocl::Kernel k("setDiag", ocl::core::copyset_oclsrc, opts);
+    if (k.empty())
+        return false;
+
+    k.args(ocl::KernelArg::WriteOnly(m),
+           ocl::KernelArg::ReadOnlyNoSize(d),
+           len);
+
+    size_t globalsize[2] = { (size_t)len, (size_t)len };
+    return k.run(2, globalsize, NULL, false);
+}
+}
+#endif
+
 UMat UMat::diag(const UMat& d, UMatUsageFlags usageFlags)
 {
-    CV_Assert( d.cols == 1 || d.rows == 1 );
+    CV_INSTRUMENT_REGION();
+
+    CV_Assert(d.cols == 1 || d.rows == 1);
     int len = d.rows + d.cols - 1;
+
+#ifdef HAVE_OPENCL
+    if (ocl::useOpenCL())
+    {
+        UMat m(len, len, d.type(), usageFlags);
+        if (ocl_setDiag(d, m, len))
+        {
+            CV_IMPL_ADD(CV_IMPL_OCL);
+            return m;
+        }
+    }
+#endif
+
     UMat m(len, len, d.type(), Scalar(0), usageFlags);
     UMat md = m.diag();
     if( d.cols == 1 )