toy 0.8.0

data/vendor/ggml/tests/test-conv1d-dw-c2.cpp

@@ -0,0 +1,243 @@
+#include "ggml.h"
+#include "ggml-cpu.h"
+#include "ggml-alloc.h"
+#include "ggml-backend.h"
+
+#ifdef GGML_USE_CUDA
+#include "ggml-cuda.h"
+#endif
+
+#ifdef GGML_USE_METAL
+#include "ggml-metal.h"
+#endif
+
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+#include <fstream>
+#include <map>
+#include <string>
+#include <vector>
+
+static void ggml_log_callback_default(ggml_log_level level, const char * text, void * user_data) {
+    (void) level;
+    (void) user_data;
+    fputs(text, stderr);
+    fflush(stderr);
+}
+
+struct test_model {
+    struct ggml_tensor * weight;
+    struct ggml_tensor * input;
+    ggml_backend_t backend = NULL;
+    ggml_backend_buffer_t buffer;
+    struct ggml_context * ctx;
+};
+
+void load_model(test_model & model, bool use_gpu = false) {
+    // create data
+    int K = 3, IC = 2, OC = 2;
+    int IL = 6, N = 1;
+
+    // Initialize adata
+    float weight_data[6] = {10.0f, 20.0f, 30.0f, 0.1f, 0.2f, 0.3f};
+
+    // Convert adata to fp16 format
+    std::vector<ggml_fp16_t> h_weight_data(K * IC);
+    ggml_fp32_to_fp16_row(weight_data, h_weight_data.data(), K * IC);
+
+    // Initialize input data, 2 channels, 6 timesteps, 1 batch
+    float input_data[12] = {
+        1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f,
+        1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f,
+    };
+
+    size_t buffer_size = 0;
+    {
+        buffer_size += K * IC * ggml_type_size(GGML_TYPE_F16); // tensor weight
+        buffer_size += IL * IC * N * ggml_type_size(GGML_TYPE_F32); // tensor input
+        buffer_size += 1024; // overhead
+    }
+
+    printf("%s: ggml tensor size    = %d bytes\n", __func__, (int) sizeof(ggml_tensor));
+    printf("%s: backend buffer size = %0.2f MB\n", __func__, (buffer_size/ 1024.f/ 1024.f));
+
+    ggml_log_set(ggml_log_callback_default, nullptr);
+
+    int num_tensors = 2;
+    struct ggml_init_params params {
+            /*.mem_size   =*/ ggml_tensor_overhead() * num_tensors,
+            /*.mem_buffer =*/ NULL,
+            /*.no_alloc   =*/ true,
+    };
+
+    // initialize the backend
+#ifdef GGML_USE_CUDA
+    if (use_gpu) {
+        fprintf(stderr, "%s: using CUDA backend\n", __func__);
+        model.backend = ggml_backend_cuda_init(0);
+        if (!model.backend) {
+            fprintf(stderr, "%s: ggml_backend_cuda_init() failed\n", __func__);
+        }
+    }
+#endif
+
+#ifdef GGML_USE_METAL
+    if (use_gpu) {
+        fprintf(stderr, "%s: using Metal backend\n", __func__);
+        model.backend = ggml_backend_metal_init();
+        if (!model.backend) {
+            fprintf(stderr, "%s: ggml_backend_metal_init() failed\n", __func__);
+        }
+    }
+#endif
+
+    if(!model.backend) {
+        // fallback to CPU backend
+        model.backend = ggml_backend_cpu_init();
+    }
+
+    model.buffer = ggml_backend_alloc_buffer(model.backend, buffer_size);
+
+    // create context
+    model.ctx = ggml_init(params);
+
+    // create tensors
+    // A Pytorch grouped Conv1d weight parameter is of shape (out_channels, input_channels/groups, kernel_size)
+    model.weight = ggml_new_tensor_3d(model.ctx, GGML_TYPE_F16,  K, 1, IC);
+    model.input = ggml_new_tensor_3d(model.ctx, GGML_TYPE_F32, IL, IC, N);
+
+    // create a allocator
+    ggml_tallocr alloc = ggml_tallocr_new(model.buffer);
+
+    // alloc memory
+    ggml_tallocr_alloc(&alloc, model.weight);
+
+    // load data to buffer
+    if(ggml_backend_is_cpu(model.backend)) {
+        memcpy(model.weight->data, h_weight_data.data(), ggml_nbytes(model.weight));
+    } else {
+        ggml_backend_tensor_set(model.weight, h_weight_data.data(), 0, ggml_nbytes(model.weight));
+    }
+
+    // alloc memory
+    ggml_tallocr_alloc(&alloc, model.input);
+
+    if(ggml_backend_is_cpu(model.backend)
+#ifdef GGML_USE_METAL
+                || ggml_backend_is_metal(model.backend)
+#endif
+    ) {
+        memcpy(model.input->data, input_data, ggml_nbytes(model.input));
+    } else {
+        ggml_backend_tensor_set(model.input, input_data, 0, ggml_nbytes(model.input));
+    }
+}
+
+struct ggml_cgraph * build_graph(const test_model& model) {
+    static size_t buf_size = ggml_tensor_overhead()*GGML_DEFAULT_GRAPH_SIZE + ggml_graph_overhead();
+    static std::vector<uint8_t> buf(buf_size);
+
+    struct ggml_init_params params0 = {
+        /*.mem_size   =*/ buf_size,
+        /*.mem_buffer =*/ buf.data(),
+        /*.no_alloc   =*/ true, // the tensors will be allocated later by ggml_gallocr_alloc_graph()
+    };
+
+    // create a temporally context to build the graph
+    struct ggml_context * ctx0 = ggml_init(params0);
+
+    struct ggml_cgraph  * gf = ggml_new_graph(ctx0);
+
+    int s0 = 3;
+    int p0 = 0;
+    int d0 = 1;
+
+    struct ggml_tensor* conv1d_dw_res = ggml_conv_1d_dw(ctx0, model.weight, model.input, s0, p0, d0);
+    ggml_set_name(conv1d_dw_res, "conv1d_dw_res");
+    ggml_build_forward_expand(gf, conv1d_dw_res);
+
+    // delete the temporally context used to build the graph
+    ggml_free(ctx0);
+    return gf;
+}
+
+struct ggml_cgraph* compute_graph(const test_model & model, ggml_gallocr_t allocr) {
+    struct ggml_cgraph * gf = build_graph(model);
+
+    // allocate tensors
+    ggml_gallocr_alloc_graph(allocr, gf);
+    int n_threads = 1;
+
+    if (ggml_backend_is_cpu(model.backend)) {
+        ggml_backend_cpu_set_n_threads(model.backend, n_threads);
+    }
+
+    ggml_backend_graph_compute(model.backend, gf);
+
+    //ggml_graph_print(gf);
+
+    return gf;
+}
+
+int main(void)
+{
+    ggml_time_init();
+
+    test_model model;
+    load_model(model, true);
+
+    ggml_gallocr_t allocr = NULL;
+
+    {
+        allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend));
+
+        //create the worst case graph for memory usage estimation
+        struct ggml_cgraph * gf = build_graph(model);
+
+        // compute the required memory
+        ggml_gallocr_reserve(allocr, gf);
+        size_t mem_size = ggml_gallocr_get_buffer_size(allocr, 0);
+        fprintf(stderr, "%s: compute buffer size: %.2f MB\n", __func__, mem_size/1024.0f/1024.0f);
+    }
+
+    struct ggml_cgraph * gf_res = compute_graph(model, allocr);
+
+    struct ggml_tensor * conv1d_dw_res = NULL;
+
+    for(int i = 0; i < ggml_graph_n_nodes(gf_res); i++) {
+        if(strcmp(ggml_get_name(ggml_graph_node(gf_res, i)), "conv1d_dw_res") == 0) {
+            conv1d_dw_res = ggml_graph_node(gf_res, i);
+        }
+    }
+
+    std::vector<float> conv2d_data(ggml_nelements(conv1d_dw_res));
+
+    ggml_backend_tensor_get(conv1d_dw_res, conv2d_data.data(), 0, ggml_nbytes(conv1d_dw_res));
+
+    const int n_conv1d_dw_test = 4;
+
+    float expected_conv1d_dw[n_conv1d_dw_test] = {
+        60.0f, 60.0f, 0.6f, 0.6f
+    };
+
+    printf("\nPerforming test:\n");
+
+    bool passed = true;
+    passed = true;
+    for(int i = 0; i < n_conv1d_dw_test; i++) {
+        if(std::abs(conv2d_data[i] - expected_conv1d_dw[i]) > 1e-4) {
+            passed = false;
+            break;
+        }
+    }
+
+    printf("ggml_conv1d (%d): %s\n", (int) ggml_nelements(conv1d_dw_res), passed && (ggml_nelements(conv1d_dw_res) == n_conv1d_dw_test) ? "\033[32mPASSED\033[0m" : "\033[31mFAILED\033[0m");
+    ggml_free(model.ctx);
+
+    ggml_backend_buffer_free(model.buffer);
+    ggml_backend_free(model.backend);
+    ggml_gallocr_free(allocr);
+    return 0;
+}

data/vendor/ggml/tests/test-conv1d.cpp

@@ -0,0 +1,289 @@
+#include "ggml.h"
+#include "ggml-cpu.h"
+#include "ggml-alloc.h"
+#include "ggml-backend.h"
+
+#ifdef GGML_USE_CUDA
+#include "ggml-cuda.h"
+#endif
+
+#ifdef GGML_USE_METAL
+#include "ggml-metal.h"
+#endif
+
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+#include <fstream>
+#include <map>
+#include <string>
+#include <vector>
+
+static void ggml_log_callback_default(ggml_log_level level, const char * text, void * user_data) {
+    (void) level;
+    (void) user_data;
+    fputs(text, stderr);
+    fflush(stderr);
+}
+
+struct test_model {
+    struct ggml_tensor * a;
+    struct ggml_tensor * b;
+    ggml_backend_t backend = NULL;
+    ggml_backend_buffer_t buffer;
+    struct ggml_context * ctx;
+};
+
+void load_model(test_model & model, bool use_gpu = false) {
+    // create data
+    int K = 3, IC = 10, OC = 10;
+    int IL = 8, N = 1;
+
+    // Initialize adata
+    std::vector<float> adata(K * IC * OC);
+    for (int i = 0; i < K * IC * OC; i++) {
+        adata[i] = 4.5f;
+    }
+
+    // Convert adata to fp16 format
+    std::vector<ggml_fp16_t> hadata(K * IC * OC);
+    ggml_fp32_to_fp16_row(adata.data(), hadata.data(), K * IC * OC);
+
+    // Initialize bdata
+    std::vector<float> bdata(IL * IC * N);
+    for (int i = 0; i < IL * IC * N; i++) {
+        bdata[i] = 2.5f;
+    }
+
+    size_t buffer_size = 0;
+    {
+        buffer_size += K * IC * OC * ggml_type_size(GGML_TYPE_F16); // tensor a
+        buffer_size += IL * IC * N * ggml_type_size(GGML_TYPE_F32); // tensor b
+        buffer_size += 1024; // overhead
+    }
+
+    printf("%s: ggml tensor size    = %d bytes\n", __func__, (int) sizeof(ggml_tensor));
+    printf("%s: backend buffer size = %0.2f MB\n", __func__, (buffer_size/ 1024.f/ 1024.f));
+
+    ggml_log_set(ggml_log_callback_default, nullptr);
+
+    int num_tensors = 2;
+    struct ggml_init_params params {
+            /*.mem_size   =*/ ggml_tensor_overhead() * num_tensors,
+            /*.mem_buffer =*/ NULL,
+            /*.no_alloc   =*/ true,
+    };
+
+    // initialize the backend
+#ifdef GGML_USE_CUDA
+    if (use_gpu) {
+        fprintf(stderr, "%s: using CUDA backend\n", __func__);
+        model.backend = ggml_backend_cuda_init(0);
+        if (!model.backend) {
+            fprintf(stderr, "%s: ggml_backend_cuda_init() failed\n", __func__);
+        }
+    }
+#endif
+
+#ifdef GGML_USE_METAL
+    if (use_gpu) {
+        fprintf(stderr, "%s: using Metal backend\n", __func__);
+        model.backend = ggml_backend_metal_init();
+        if (!model.backend) {
+            fprintf(stderr, "%s: ggml_backend_metal_init() failed\n", __func__);
+        }
+    }
+#endif
+
+    if(!model.backend) {
+        // fallback to CPU backend
+        model.backend = ggml_backend_cpu_init();
+    }
+
+    model.buffer = ggml_backend_alloc_buffer(model.backend, buffer_size);
+
+    // create context
+    model.ctx = ggml_init(params);
+
+    // create tensors
+    model.a = ggml_new_tensor_3d(model.ctx, GGML_TYPE_F16,  K, IC, OC);
+    model.b = ggml_new_tensor_3d(model.ctx, GGML_TYPE_F32, IL, IC, N);
+
+    // create a allocator
+    ggml_tallocr alloc = ggml_tallocr_new(model.buffer);
+
+    // alloc memory
+    ggml_tallocr_alloc(&alloc, model.a);
+
+    // load data to buffer
+    if(ggml_backend_is_cpu(model.backend)) {
+        memcpy(model.a->data, hadata.data(), ggml_nbytes(model.a));
+    } else {
+        ggml_backend_tensor_set(model.a, hadata.data(), 0, ggml_nbytes(model.a));
+    }
+
+    // alloc memory
+    ggml_tallocr_alloc(&alloc, model.b);
+
+    if(ggml_backend_is_cpu(model.backend)
+#ifdef GGML_USE_METAL
+                || ggml_backend_is_metal(model.backend)
+#endif
+    ) {
+        memcpy(model.b->data, bdata.data(), ggml_nbytes(model.b));
+    } else {
+        ggml_backend_tensor_set(model.b, bdata.data(), 0, ggml_nbytes(model.b));
+    }
+}
+
+struct ggml_cgraph * build_graph(const test_model& model) {
+    static size_t buf_size = ggml_tensor_overhead()*GGML_DEFAULT_GRAPH_SIZE + ggml_graph_overhead();
+    static std::vector<uint8_t> buf(buf_size);
+
+    struct ggml_init_params params0 = {
+        /*.mem_size   =*/ buf_size,
+        /*.mem_buffer =*/ buf.data(),
+        /*.no_alloc   =*/ true, // the tensors will be allocated later by ggml_gallocr_alloc_graph()
+    };
+
+    // create a temporally context to build the graph
+    struct ggml_context * ctx0 = ggml_init(params0);
+
+    struct ggml_cgraph  * gf = ggml_new_graph(ctx0);
+
+    int s0 = 1;
+    int p0 = 1;
+    int d0 = 1;
+
+    // split conv1d in fundamental methods for test unit
+    struct ggml_tensor* im2col_0 = ggml_im2col(ctx0, model.a, model.b, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F16);
+    ggml_set_name(im2col_0, "im2col_res");
+    ggml_build_forward_expand(gf, im2col_0);
+
+    struct ggml_tensor* conv1d_res = ggml_conv_1d(ctx0, model.a, model.b, s0, p0, d0);
+    ggml_set_name(conv1d_res, "conv1d_res");
+    ggml_build_forward_expand(gf, conv1d_res);
+
+    // delete the temporally context used to build the graph
+    ggml_free(ctx0);
+    return gf;
+}
+
+struct ggml_cgraph* compute_graph(const test_model & model, ggml_gallocr_t allocr) {
+    struct ggml_cgraph * gf = build_graph(model);
+
+    // allocate tensors
+    ggml_gallocr_alloc_graph(allocr, gf);
+    int n_threads = 1;
+
+    if (ggml_backend_is_cpu(model.backend)) {
+        ggml_backend_cpu_set_n_threads(model.backend, n_threads);
+    }
+
+    ggml_backend_graph_compute(model.backend, gf);
+
+    //ggml_graph_print(gf);
+
+    return gf;
+}
+
+int main(void)
+{
+    ggml_time_init();
+
+    test_model model;
+    load_model(model, true);
+
+    ggml_gallocr_t allocr = NULL;
+
+    {
+        allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend));
+
+        //create the worst case graph for memory usage estimation
+        struct ggml_cgraph * gf = build_graph(model);
+
+        // compute the required memory
+        ggml_gallocr_reserve(allocr, gf);
+        size_t mem_size = ggml_gallocr_get_buffer_size(allocr, 0);
+        fprintf(stderr, "%s: compute buffer size: %.2f MB\n", __func__, mem_size/1024.0f/1024.0f);
+    }
+
+    struct ggml_cgraph * gf_res = compute_graph(model, allocr);
+
+     struct ggml_tensor * im2col_res = NULL;
+    struct ggml_tensor * conv1d_res = NULL;
+
+    for(int i = 0; i < ggml_graph_n_nodes(gf_res); i++) {
+        if(strcmp(ggml_get_name(ggml_graph_node(gf_res, i)), "im2col_res") == 0) {
+            im2col_res = ggml_graph_node(gf_res, i);
+        } else if(strcmp(ggml_get_name(ggml_graph_node(gf_res, i)), "conv1d_res") == 0) {
+            conv1d_res = ggml_graph_node(gf_res, i);
+        }
+    }
+
+    std::vector<uint16_t> im2col_data(ggml_nelements(im2col_res));
+    std::vector<float> conv2d_data(ggml_nelements(conv1d_res));
+
+    ggml_backend_tensor_get(im2col_res, im2col_data.data(), 0, ggml_nbytes(im2col_res));
+    ggml_backend_tensor_get(conv1d_res, conv2d_data.data(), 0, ggml_nbytes(conv1d_res));
+
+    const int n_conv1d_test = 80;
+    const int n_im2col_test = 240;
+
+    float expected_conv1d[n_conv1d_test] = {
+        225.00f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 225.00f,
+        225.00f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 225.00f,
+        225.00f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 225.00f,
+        225.00f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 225.00f,
+        225.00f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 225.00f,
+        225.00f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 225.00f,
+        225.00f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 225.00f,
+        225.00f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 225.00f,
+        225.00f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 225.00f,
+        225.00f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 337.50f, 225.00f
+    };
+    // first im2col test
+
+    uint16_t expected_im2col[n_conv1d_test] = {
+        0, 16640, 16640, 0, 16640, 16640, 0, 16640,
+        16640, 0, 16640, 16640, 0, 16640, 16640, 0,
+        16640, 16640, 0, 16640, 16640, 0, 16640, 16640,
+        0, 16640, 16640, 0, 16640, 16640, 16640, 16640,
+        16640, 16640, 16640, 16640, 16640, 16640, 16640, 16640,
+        16640, 16640, 16640, 16640, 16640, 16640, 16640, 16640,
+        16640, 16640, 16640, 16640, 16640, 16640, 16640, 16640,
+        16640, 16640, 16640, 16640, 16640, 16640, 16640, 16640,
+        16640, 16640, 16640, 16640, 16640, 16640, 16640, 16640,
+        16640, 16640, 16640, 16640, 16640, 16640, 16640, 16640
+    };
+
+    printf("\nPerforming test:\n");
+
+    bool passed = true;
+    for(int i = 0; i < n_conv1d_test; i++) {
+        if(
+            im2col_data[i] != expected_im2col[i]) {
+            passed = false;
+            break;
+        }
+    }
+
+    printf("ggml_im2col (%d): %s\n", (int) ggml_nelements(im2col_res), passed && (ggml_nelements(im2col_res) == n_im2col_test) ? "\033[32mPASSED\033[0m" : "\033[31mFAILED\033[0m");
+
+    passed = true;
+    for(int i = 0; i < n_conv1d_test; i++) {
+        if(conv2d_data[i] != expected_conv1d[i]) {
+            passed = false;
+            break;
+        }
+    }
+
+    printf("ggml_conv1d (%d): %s\n", (int) ggml_nelements(conv1d_res), passed && (ggml_nelements(conv1d_res) == n_conv1d_test) ? "\033[32mPASSED\033[0m" : "\033[31mFAILED\033[0m");
+    ggml_free(model.ctx);
+
+    ggml_backend_buffer_free(model.buffer);
+    ggml_backend_free(model.backend);
+    ggml_gallocr_free(allocr);
+    return 0;
+}

data/vendor/ggml/tests/test-conv2d-dw.cpp

@@ -0,0 +1,153 @@
+#include <ggml.h>
+#include <ggml-cpu.h>
+#include <ggml-alloc.h>
+#include <ggml-backend.h>
+#include <ggml-cpp.h>
+
+#include <cassert>
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+#include <vector>
+
+std::vector<float> f32_range(int n, float start, float end) {
+    std::vector<float> values(n);
+    float step = (end - start) / n;
+    for (int i = 0; i < n; i++) {
+        values[i] = start + i * step;
+    }
+    return values;
+}
+
+// Most straightforward implementation without any optimizations
+std::vector<float> conv_2d_dw_reference(
+        int src_w, int src_h, const float * src_data,
+        int knl_w, int knl_h, const float * knl_data,
+        int channels, int batch, int stride, int pad, int dilation) {
+
+    int dst_w = (src_w + 2 * pad - dilation * (knl_w - 1) - 1) / stride + 1;
+    int dst_h = (src_h + 2 * pad - dilation * (knl_h - 1) - 1) / stride + 1;
+    std::vector<float> dst_data(dst_w * dst_h * channels * batch);
+
+    for (int b = 0; b < batch; b++) {
+        const float * src_base = src_data + b * src_w * src_h * channels;
+        float * dst_base = dst_data.data() + b * dst_w * dst_h * channels;
+        for (int c = 0; c < channels; c++) {
+            for (int y = 0; y < dst_h; y++) {
+                for (int x = 0; x < dst_w; x++) {
+                    float sum = 0;
+                    for (int knl_y = 0; knl_y < knl_h; knl_y++) {
+                        for (int knl_x = 0; knl_x < knl_w; knl_x++) {
+                            int src_x = x * stride + knl_x * dilation - pad;
+                            int src_y = y * stride + knl_y * dilation - pad;
+                            if (src_x >= 0 && src_x < src_w && src_y >= 0 && src_y < src_h) {
+                                sum += src_base[c * src_w * src_h + src_y * src_w + src_x] *
+                                       knl_data[c * knl_w * knl_h + knl_y * knl_w + knl_x];
+                            }
+                        }
+                    }
+                    dst_base[c * dst_w * dst_h + y * dst_w + x] = sum;
+                }
+            }
+        }
+    }
+    return dst_data;
+}
+
+bool check_equal(const std::vector<float> & result, const std::vector<float> & expected) {
+    if (result.size() != expected.size()) {
+        printf("result.size() = %d, expected.size() = %d\n", (int)result.size(), (int)expected.size());
+        return false;
+    }
+    for (int i = 0; i < result.size(); i++) {
+        if(std::abs(result[i] - expected[i]) > 1e-5) {
+            printf("result[%d] %f != %f expected[%d]\n", i, result[i], expected[i], i);
+            return false;
+        }
+    }
+    return true;
+}
+
+bool test_conv_2d_dw(
+        int channels,
+        int kernel_size,
+        int stride,
+        int pad,
+        int dilation,
+        bool contiguous_channels) {
+    ggml_time_init();
+
+    const int batch = 2;
+    const int src_w = 8;
+    const int src_h = 6;
+    const int knl_w = kernel_size;
+    const int knl_h = kernel_size;
+
+    ggml_init_params params {
+        /*.mem_size   =*/ 64 * ggml_tensor_overhead() + ggml_graph_overhead(),
+        /*.mem_buffer =*/ NULL,
+        /*.no_alloc   =*/ true
+    };
+
+    ggml_context_ptr ctx_ptr{ggml_init(params)};
+    ggml_context * ctx = ctx_ptr.get();
+    ggml_cgraph * gf = ggml_new_graph(ctx);
+
+    // Build graph
+    ggml_tensor * src_input = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, src_w, src_h, channels, batch);
+    ggml_tensor * knl_input = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, knl_w, knl_h, 1, channels);
+    ggml_tensor * src = src_input;
+    ggml_tensor * knl = knl_input;
+    if (contiguous_channels) {
+        // Convert tensor to [C, W, H, N] layout in memory, then permute strides back to [W, H, C, N]
+        src = ggml_cont(ctx, ggml_permute(ctx, src, 1, 2, 0, 3));
+        src = ggml_permute(ctx, src, 2, 0, 1, 3);
+        knl = ggml_cont(ctx, ggml_permute(ctx, knl, 2, 3, 1, 0));
+        knl = ggml_permute(ctx, knl, 3, 2, 0, 1);
+    }
+    ggml_tensor * res = ggml_conv_2d_dw_direct(
+        ctx, knl, src, stride, stride, pad, pad, dilation, dilation);
+    if (contiguous_channels) {
+        res = ggml_cont(ctx, res);
+    }
+    ggml_build_forward_expand(gf, res);
+
+    // Create backend & allocate buffers
+    ggml_backend_ptr backend_ptr{ggml_backend_cpu_init()};
+    ggml_backend_t backend = backend_ptr.get();
+    ggml_backend_cpu_set_n_threads(backend, 2);
+    ggml_backend_buffer_ptr buffer{ggml_backend_alloc_ctx_tensors(ctx, backend)};
+
+    std::vector<float> src_values = f32_range(ggml_nelements(src), -1.f, 1.f);
+    std::vector<float> knl_values = f32_range(ggml_nelements(knl), -1.f, 1.f);
+    ggml_backend_tensor_set(src_input, src_values.data(), 0, ggml_nbytes(src));
+    ggml_backend_tensor_set(knl_input, knl_values.data(), 0, ggml_nbytes(knl));
+
+    ggml_backend_graph_compute(backend, gf);
+
+    std::vector<float> res_values(ggml_nelements(res));
+    ggml_backend_tensor_get(res, res_values.data(), 0, ggml_nbytes(res));
+
+    std::vector<float> expected = conv_2d_dw_reference(
+        src_w, src_h, src_values.data(),
+        knl_w, knl_h, knl_values.data(),
+        channels, batch, stride, pad, dilation);
+
+    bool passed = check_equal(res_values, expected);
+
+    printf("ggml_conv_2d_dw(channels=%d, kernel=%dx%d, stride=%d, pad=%d, dilation=%d, layout=%s): %s\n",
+        channels, kernel_size, kernel_size, stride, pad, dilation, contiguous_channels ? "CWHN" : "WHCN",
+        passed ? "\033[32mPASSED\033[0m" : "\033[31mFAILED\033[0m");
+    return passed;
+}
+
+int main(int argc, char ** argv) {
+    bool passed = true;
+    passed = test_conv_2d_dw(3, 1, 1, 0, 1, false) && passed;
+    passed = test_conv_2d_dw(3, 1, 1, 0, 1, true) && passed;
+    passed = test_conv_2d_dw(42, 3, 2, 1, 1, false) && passed;
+    passed = test_conv_2d_dw(42, 3, 2, 1, 1, true) && passed;
+    passed = test_conv_2d_dw(8, 5, 1, 2, 2, false) && passed;
+    passed = test_conv_2d_dw(8, 5, 1, 2, 2, true) && passed;
+    return passed ? 0 : 1;
+}