Modularize and cleanHEADmain

6 files changed, 282 insertions, 179 deletions

diff --git a/.gitignore b/.gitignore
index ccc36f7..42a0e36 100644
--- a/.gitignore
+++ b/.gitignore
@@ -68,3 +68,4 @@ vcpkg_installed/
 # test output & cache
 Testing/
 .cache/
+vendor/
diff --git a/CMakeLists.txt b/CMakeLists.txt
new file mode 100644
index 0000000..8ac5561
--- /dev/null
+++ b/CMakeLists.txt
@@ -0,0 +1,20 @@
+cmake_minimum_required(VERSION 3.21)
+project(kuwahara LANGUAGES CXX HIP)
+
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+
+find_package(hip REQUIRED)
+
+add_executable(kuwahara
+    src/main.cpp
+    src/kernels/pyramid.hip
+)
+
+target_include_directories(kuwahara PRIVATE
+    src/
+    vendor/
+)
+
+target_link_libraries(kuwahara PRIVATE hip::host)
diff --git a/main.cpp b/main.cpp
deleted file mode 100644
index bd03139..0000000
--- a/main.cpp
+++ /dev/null
@@ -1,179 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <source_location>
-#include <stdexcept>
-#include <string>
-#include <vector>
-#include <cmath>
-#include <algorithm>
-
-
-inline void hip_check(
-    hipError_t err,
-    std::source_location loc = std::source_location::current())
-{
-    if (err != hipSuccess) {
-        throw std::runtime_error(
-            std::string(hipGetErrorString(err)) +
-            " at " + loc.file_name() + ":" + std::to_string(loc.line()));
-    }
-}
-
-
-struct PyramidLevel {
-    int width = 0;
-    int height = 0;
-    float4* color_smax = nullptr; /* rgb s_max*/ 
-    float4* tensor = nullptr;      /* x=E, y=F, z=G, w= */ 
-};
-
-struct KuwaharaPyramid {
-    std::vector<PyramidLevel> levels; /*  0 = finest, back() = coarsest  */
-};
-
-struct KuwaharaParams {
-    float radius;
-    float alpha;      
-    int   num_sectors; /* 4 or 8 */
-    float q;
-    float tau_w;
-    float p_s, p_d, tau_v;
-};
-
-/* allocations */
-
-PyramidLevel make_level(int width, int height) {
-    PyramidLevel lvl{width, height};
-    hip_check(hipMalloc(&lvl.color_smax, sizeof(float4) * width * height));
-    hip_check(hipMalloc(&lvl.tensor,     sizeof(float4) * width * height));
-    return lvl;
-}
-
-void free_level(PyramidLevel& lvl) {
-    hipFree(lvl.color_smax);
-    hipFree(lvl.tensor);
-    lvl.color_smax = nullptr;
-    lvl.tensor = nullptr;
-}
-
-KuwaharaPyramid make_pyramid(int base_width, int base_height, int num_levels) {
-    KuwaharaPyramid pyr;
-    pyr.levels.reserve(num_levels);
-    int w = base_width, h = base_height;
-    for (int i = 0; i < num_levels; ++i) {
-        pyr.levels.push_back(make_level(w, h));
-        w = std::max(1, w / 2);
-        h = std::max(1, h / 2);
-    }
-    return pyr;
-}
-
-void destroy_pyramid(KuwaharaPyramid& pyr) {
-    for (auto& lvl : pyr.levels) {
-        free_level(lvl);
-    } 
-    pyr.levels.clear();
-}
-
-/* Lanczos3 downsample */
-
-__device__ __forceinline__ float lanczos3(float x) {
-    const float a = 3.0f;
-    if (x == 0.0f) return 1.0f;
-    if (x <= -a || x >= a) {
-        return 0.0f;
-    }
-    const float px = M_PI * x;
-    return (a * sinf(px) * sinf(px / a)) / (px * px);
-}
-
-__device__ float4 lanczos3_sample_axis(
-    const float4* src, int src_len, int src_stride, int src_base, int dst_idx, float scale)
-{
-    float center = (dst_idx + 0.5f) * scale - 0.5f;
-    int icenter = static_cast<int>(floorf(center));
-
-    float filter_scale = fmaxf(scale, 1.0f);
-    int radius = static_cast<int>(ceilf(3.0f * filter_scale));
-
-    float4 sum = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
-    float wsum = 0.0f;
-
-    for (int i = icenter - radius; i <= icenter + radius; ++i) {
-        int clamped = min(max(i, 0), src_len - 1);
-        float w = lanczos3((i - center) / filter_scale);
-        const float4& s = src[src_base + clamped * src_stride];
-        sum.x += s.x * w;
-        sum.y += s.y * w;
-        sum.z += s.z * w;
-        sum.w += s.w * w;
-        wsum += w;
-    }
-
-    if (wsum > 1e-6f) {
-        sum.x /= wsum;
-        sum.y /= wsum;
-        sum.z /= wsum;
-        sum.w /= wsum;
-    }
-    return sum;
-}
-
-__global__ void downsample_horizontal(
-    const float4* src, int src_w, int h,
-    float4* dst, int dst_w, float scale_x)
-{
-    int x = blockIdx.x * blockDim.x + threadIdx.x;
-    int y = blockIdx.y * blockDim.y + threadIdx.y;
-    if (x >= dst_w || y >= h) {
-        return;
-    }
-
-    int src_row_base = y * src_w;
-    dst[y * dst_w + x] =
-        lanczos3_sample_axis(src, src_w, 1, src_row_base, x, scale_x);
-}
-
-__global__ void downsample_vertical(
-    const float4* src, int w, int src_h,
-    float4* dst, int dst_h, float scale_y)
-{
-    int x = blockIdx.x * blockDim.x + threadIdx.x;
-    int y = blockIdx.y * blockDim.y + threadIdx.y;
-    if (x >= w || y >= dst_h) return;
-
-    dst[y * w + x] =
-        lanczos3_sample_axis(src, src_h, w, x, y, scale_y);
-}
-
-void downsample_level(const PyramidLevel& src, PyramidLevel& dst) {
-    float scale_x = static_cast<float>(src.width) / dst.width;
-    float scale_y = static_cast<float>(src.height) / dst.height;
-
-    float4* tmp;
-    hip_check(hipMalloc(&tmp, sizeof(float4) * dst.width * src.height));
-
-    dim3 block(16, 16);
-
-    dim3 grid_h((dst.width + block.x - 1) / block.x,
-                (src.height + block.y - 1) / block.y);
-    downsample_horizontal<<<grid_h, block>>>(
-        src.color_smax, src.width, src.height, tmp, dst.width, scale_x);
-    hip_check(hipGetLastError());
-
-    dim3 grid_v((dst.width + block.x - 1) / block.x,
-                (dst.height + block.y - 1) / block.y);
-    downsample_vertical<<<grid_v, block>>>(
-        tmp, dst.width, src.height, dst.color_smax, dst.height, scale_y);
-    hip_check(hipGetLastError());
-
-    hipFree(tmp);
-}
-
-void build_pyramid(KuwaharaPyramid& pyr) {
-    for (size_t i = 1; i < pyr.levels.size(); ++i) {
-        downsample_level(pyr.levels[i - 1], pyr.levels[i]);
-    }
-}
-
-
-int main() { }
diff --git a/src/kernels/pyramid.hip b/src/kernels/pyramid.hip
new file mode 100644
index 0000000..80d9838
--- /dev/null
+++ b/src/kernels/pyramid.hip
@@ -0,0 +1,122 @@
+#include "pipeline.h"
+#include <algorithm>
+
+PyramidLevel make_level(int width, int height) {
+    PyramidLevel lvl{width, height};
+    hip_check(hipMalloc(&lvl.color_smax, sizeof(float4) * width * height));
+    hip_check(hipMalloc(&lvl.tensor,     sizeof(float4) * width * height));
+    return lvl;
+}
+
+void free_level(PyramidLevel& lvl) {
+    hipFree(lvl.color_smax);
+    hipFree(lvl.tensor);
+    lvl.color_smax = nullptr;
+    lvl.tensor     = nullptr;
+}
+
+Pyramid make_pyramid(int base_width, int base_height, int num_levels) {
+    Pyramid pyr;
+    pyr.levels.reserve(num_levels);
+    int w = base_width, h = base_height;
+    for (int i = 0; i < num_levels; ++i) {
+        pyr.levels.push_back(make_level(w, h));
+        w = std::max(1, w / 2);
+        h = std::max(1, h / 2);
+    }
+    return pyr;
+}
+
+void destroy_pyramid(Pyramid& pyr) {
+    for (auto& lvl : pyr.levels) free_level(lvl);
+    pyr.levels.clear();
+}
+
+/* Lanczos3 */
+
+__device__ __forceinline__ float lanczos3(float x) {
+    const float a = 3.0f;
+    if (x == 0.0f) return 1.0f;
+    if (x <= -a || x >= a) return 0.0f;
+    const float px = M_PI * x;
+    return (a * sinf(px) * sinf(px / a)) / (px * px);
+}
+
+__device__ float4 lanczos3_sample_axis(
+    const float4* src, int src_len, int src_stride, int src_base,
+    int dst_idx, float scale)
+{
+    float center  = (dst_idx + 0.5f) * scale - 0.5f;
+    int icenter   = static_cast<int>(floorf(center));
+    float fscale  = fmaxf(scale, 1.0f);
+    int radius    = static_cast<int>(ceilf(3.0f * fscale));
+
+    float4 sum = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+    float wsum = 0.0f;
+
+    for (int i = icenter - radius; i <= icenter + radius; ++i) {
+        int c = min(max(i, 0), src_len - 1);
+        float w = lanczos3((i - center) / fscale);
+        float4 s = src[src_base + c * src_stride];
+        sum.x += s.x * w; sum.y += s.y * w;
+        sum.z += s.z * w; sum.w += s.w * w;
+        wsum += w;
+    }
+
+    if (wsum > 1e-6f) {
+        float inv = 1.0f / wsum;
+        sum.x *= inv; sum.y *= inv; sum.z *= inv; sum.w *= inv;
+    }
+    return sum;
+}
+
+__global__ void downsample_horizontal(
+    const float4* src, int src_w, int h,
+    float4* dst, int dst_w, float scale_x)
+{
+    int x = blockIdx.x * blockDim.x + threadIdx.x;
+    int y = blockIdx.y * blockDim.y + threadIdx.y;
+    if (x >= dst_w || y >= h) return;
+
+    dst[y * dst_w + x] =
+        lanczos3_sample_axis(src, src_w, 1, y * src_w, x, scale_x);
+}
+
+__global__ void downsample_vertical(
+    const float4* src, int w, int src_h,
+    float4* dst, int dst_h, float scale_y)
+{
+    int x = blockIdx.x * blockDim.x + threadIdx.x;
+    int y = blockIdx.y * blockDim.y + threadIdx.y;
+    if (x >= w || y >= dst_h) return;
+
+    dst[y * w + x] =
+        lanczos3_sample_axis(src, src_h, w, x, y, scale_y);
+}
+
+static void downsample_level(const PyramidLevel& src, PyramidLevel& dst) {
+    float scale_x = static_cast<float>(src.width)  / dst.width;
+    float scale_y = static_cast<float>(src.height) / dst.height;
+
+    float4* tmp;
+    hip_check(hipMalloc(&tmp, sizeof(float4) * dst.width * src.height));
+
+    dim3 block(16, 16);
+
+    dim3 grid_h((dst.width  + 15) / 16, (src.height + 15) / 16);
+    downsample_horizontal<<<grid_h, block>>>(
+        src.color_smax, src.width, src.height, tmp, dst.width, scale_x);
+    hip_check(hipGetLastError());
+
+    dim3 grid_v((dst.width  + 15) / 16, (dst.height + 15) / 16);
+    downsample_vertical<<<grid_v, block>>>(
+        tmp, dst.width, src.height, dst.color_smax, dst.height, scale_y);
+    hip_check(hipGetLastError());
+
+    hipFree(tmp);
+}
+
+void build_pyramid(Pyramid& pyr) {
+    for (size_t i = 1; i < pyr.levels.size(); ++i)
+        downsample_level(pyr.levels[i - 1], pyr.levels[i]);
+}
diff --git a/src/main.cpp b/src/main.cpp
new file mode 100644
index 0000000..6526f82
--- /dev/null
+++ b/src/main.cpp
@@ -0,0 +1,87 @@
+#define STB_IMAGE_IMPLEMENTATION
+#include "stb_image.h"
+#define STB_IMAGE_WRITE_IMPLEMENTATION
+#include "stb_image_write.h"
+
+#include "pipeline.h"
+#include <hip/hip_runtime.h>
+#include <cstdio>
+#include <string>
+#include <vector>
+
+static float4* upload_image(const char* path, int& out_width, int& out_height) {
+    int w, h, channels;
+    unsigned char* data = stbi_load(path, &w, &h, &channels, 4);
+    if (!data) throw std::runtime_error(std::string("Failed to load: ") + path);
+
+    out_width = w; out_height = h;
+    const int n = w * h;
+
+    std::vector<float4> host(n);
+    for (int i = 0; i < n; ++i) {
+        host[i] = make_float4(
+            data[4*i + 0] / 255.0f,
+            data[4*i + 1] / 255.0f,
+            data[4*i + 2] / 255.0f,
+            0.0f);
+    }
+    stbi_image_free(data);
+
+    float4* dev;
+    hip_check(hipMalloc(&dev, sizeof(float4) * n));
+    hip_check(hipMemcpy(dev, host.data(), sizeof(float4) * n, hipMemcpyHostToDevice));
+    return dev;
+}
+
+static void dump_level(const PyramidLevel& lvl, const char* path) {
+    const int n = lvl.width * lvl.height;
+    std::vector<float4> host(n);
+    hip_check(hipMemcpy(host.data(), lvl.color_smax,
+                        sizeof(float4) * n, hipMemcpyDeviceToHost));
+
+    std::vector<unsigned char> bytes(n * 3);
+    for (int i = 0; i < n; ++i) {
+        bytes[3*i + 0] = static_cast<unsigned char>(
+            std::min(host[i].x, 1.0f) * 255.0f);
+        bytes[3*i + 1] = static_cast<unsigned char>(
+            std::min(host[i].y, 1.0f) * 255.0f);
+        bytes[3*i + 2] = static_cast<unsigned char>(
+            std::min(host[i].z, 1.0f) * 255.0f);
+    }
+    stbi_write_png(path, lvl.width, lvl.height, 3, bytes.data(), lvl.width * 3);
+    printf("Wrote %s (%dx%d)\n", path, lvl.width, lvl.height);
+}
+
+int main(int argc, char* argv[]) {
+    if (argc < 2) {
+        fprintf(stderr, "Usage: %s <image>\n", argv[0]);
+        return 1;
+    }
+
+    try {
+        int w, h;
+        float4* src = upload_image(argv[1], w, h);
+
+        const int num_levels = 5;
+        Pyramid pyr = make_pyramid(w, h, num_levels);
+
+        hip_check(hipMemcpy(pyr.levels[0].color_smax, src,
+                            sizeof(float4) * w * h, hipMemcpyDeviceToDevice));
+        hipFree(src);
+
+        build_pyramid(pyr);
+
+        for (int i = 0; i < static_cast<int>(pyr.levels.size()); ++i) {
+            std::string path = "level_" + std::to_string(i) + ".png";
+            dump_level(pyr.levels[i], path.c_str());
+        }
+
+        destroy_pyramid(pyr);
+
+    } catch (const std::exception& e) {
+        fprintf(stderr, "Error: %s\n", e.what());
+        return 1;
+    }
+
+    return 0;
+}
diff --git a/src/pipeline.h b/src/pipeline.h
new file mode 100644
index 0000000..93a7ea9
--- /dev/null
+++ b/src/pipeline.h
@@ -0,0 +1,52 @@
+#pragma once
+
+#include <hip/hip_runtime.h>
+#include <stdexcept>
+#include <string>
+#include <vector>
+#include <source_location>
+
+
+inline void hip_check(
+    hipError_t err,
+    std::source_location loc = std::source_location::current())
+{
+    if (err != hipSuccess) {
+        throw std::runtime_error(
+            std::string(hipGetErrorString(err)) +
+            " at " + loc.file_name() + ":" + std::to_string(loc.line()));
+    }
+}
+
+
+struct PyramidLevel {
+    int width = 0, height = 0;
+    float4* color_smax = nullptr; // rgb = color, a = s_max
+    float4* tensor     = nullptr; // x=E, y=F, z=G, w=unused
+};
+
+struct Pyramid {
+    std::vector<PyramidLevel> levels; // 0 = finest, back() = coarsest
+};
+
+struct Params {
+    float radius      = 5.0f;
+    float alpha       = 1.0f; // eccentricity tuning
+    int   num_sectors = 8;
+    float q           = 8.0f;
+    float tau_w       = 0.02f;
+    float p_s         = 0.5f;
+    float p_d         = 1.25f;
+    float tau_v       = 0.1f;
+};
+
+
+PyramidLevel make_level(int width, int height);
+void         free_level(PyramidLevel& lvl);
+Pyramid make_pyramid(int base_width, int base_height, int num_levels);
+void            destroy_pyramid(Pyramid& pyr);
+
+void build_pyramid (Pyramid& pyr);
+void compute_tensor(Pyramid& pyr, const Params& p);
+void smooth_tensor (Pyramid& pyr, const Params& p);
+void run_pipeline  (Pyramid& pyr, const Params& p);