1 files changed, 106 insertions, 0 deletions
diff --git a/src/kernels/tensor.hip b/src/kernels/tensor.hip
new file mode 100644
index 0000000..d97bb5e
--- /dev/null
+++ b/src/kernels/tensor.hip
@@ -0,0 +1,106 @@
+#include "pipeline.h"
+
+/* Structure tensor from Jähne's rotationally optimized Sobel derivatives */
+
+__global__ void structure_tensor(
+    const float4* src, float4* dst, int w, int h)
+{
+    int x = blockIdx.x * blockDim.x + threadIdx.x;
+    int y = blockIdx.y * blockDim.y + threadIdx.y;
+    if (x >= w || y >= h) return;
+
+    const float p1 = 0.183f;
+    const float p2 = 1.0f - 2.0f * p1;
+
+    int xm = max(x - 1, 0), xp = min(x + 1, w - 1);
+    int ym = max(y - 1, 0), yp = min(y + 1, h - 1);
+
+    float4 c_mm = src[ym * w + xm], c_0m = src[ym * w + x], c_pm = src[ym * w + xp];
+    float4 c_m0 = src[y  * w + xm],                         c_p0 = src[y  * w + xp];
+    float4 c_mp = src[yp * w + xm], c_0p = src[yp * w + x], c_pp = src[yp * w + xp];
+
+    float3 gx, gy;
+    gx.x = 0.5f * (p1 * (c_pm.x - c_mm.x) + p2 * (c_p0.x - c_m0.x) + p1 * (c_pp.x - c_mp.x));
+    gx.y = 0.5f * (p1 * (c_pm.y - c_mm.y) + p2 * (c_p0.y - c_m0.y) + p1 * (c_pp.y - c_mp.y));
+    gx.z = 0.5f * (p1 * (c_pm.z - c_mm.z) + p2 * (c_p0.z - c_m0.z) + p1 * (c_pp.z - c_mp.z));
+    gy.x = 0.5f * (p1 * (c_mp.x - c_mm.x) + p2 * (c_0p.x - c_0m.x) + p1 * (c_pp.x - c_pm.x));
+    gy.y = 0.5f * (p1 * (c_mp.y - c_mm.y) + p2 * (c_0p.y - c_0m.y) + p1 * (c_pp.y - c_pm.y));
+    gy.z = 0.5f * (p1 * (c_mp.z - c_mm.z) + p2 * (c_0p.z - c_0m.z) + p1 * (c_pp.z - c_pm.z));
+
+    dst[y * w + x] = make_float4(
+        gx.x * gx.x + gx.y * gx.y + gx.z * gx.z,
+        gx.x * gy.x + gx.y * gy.y + gx.z * gy.z,
+        gy.x * gy.x + gy.y * gy.y + gy.z * gy.z,
+        0.0f);
+}
+
+__device__ float4 gaussian_sample_axis(
+    const float4* src, int src_len, int src_stride, int src_base,
+    int center, float sigma)
+{
+    int radius = static_cast<int>(ceilf(3.0f * sigma));
+    float inv_two_sigma2 = 1.0f / (2.0f * sigma * sigma);
+
+    float4 sum = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+    float wsum = 0.0f;
+
+    for (int i = -radius; i <= radius; ++i) {
+        int c = min(max(center + i, 0), src_len - 1);
+        float wt = expf(-(i * i) * inv_two_sigma2);
+        float4 s = src[src_base + c * src_stride];
+        sum.x += s.x * wt; sum.y += s.y * wt;
+        sum.z += s.z * wt; sum.w += s.w * wt;
+        wsum += wt;
+    }
+
+    float inv = 1.0f / wsum;
+    sum.x *= inv; sum.y *= inv; sum.z *= inv; sum.w *= inv;
+    return sum;
+}
+
+__global__ void gaussian_horizontal(
+    const float4* src, float4* dst, int w, int h, float sigma)
+{
+    int x = blockIdx.x * blockDim.x + threadIdx.x;
+    int y = blockIdx.y * blockDim.y + threadIdx.y;
+    if (x >= w || y >= h) return;
+
+    dst[y * w + x] = gaussian_sample_axis(src, w, 1, y * w, x, sigma);
+}
+
+__global__ void gaussian_vertical(
+    const float4* src, float4* dst, int w, int h, float sigma)
+{
+    int x = blockIdx.x * blockDim.x + threadIdx.x;
+    int y = blockIdx.y * blockDim.y + threadIdx.y;
+    if (x >= w || y >= h) return;
+
+    dst[y * w + x] = gaussian_sample_axis(src, h, w, x, y, sigma);
+}
+
+void compute_tensor(PyramidLevel& lvl) {
+    dim3 block(16, 16);
+    dim3 grid((lvl.width + 15) / 16, (lvl.height + 15) / 16);
+
+    structure_tensor<<<grid, block>>>(
+        lvl.color_smax, lvl.tensor, lvl.width, lvl.height);
+    hip_check(hipGetLastError());
+}
+
+void smooth_tensor(PyramidLevel& lvl, const Params& p) {
+    float4* tmp;
+    hip_check(hipMalloc(&tmp, sizeof(float4) * lvl.width * lvl.height));
+
+    dim3 block(16, 16);
+    dim3 grid((lvl.width + 15) / 16, (lvl.height + 15) / 16);
+
+    gaussian_horizontal<<<grid, block>>>(
+        lvl.tensor, tmp, lvl.width, lvl.height, p.rho);
+    hip_check(hipGetLastError());
+
+    gaussian_vertical<<<grid, block>>>(
+        tmp, lvl.tensor, lvl.width, lvl.height, p.rho);
+    hip_check(hipGetLastError());
+
+    hipFree(tmp);
+}