libs/vulkan-compute/include/vk_similarity.h

146 lines

1/* vk_similarity.h - Vulkan-accelerated similarity matrix computation
2 *
3 * Provides GPU-accelerated cosine similarity calculation for poem embeddings
4 * using triangular individual files format for storage efficiency.
5 */
6
7#ifndef VK_SIMILARITY_H
8#define VK_SIMILARITY_H
9
10#include "vk_compute.h"
11#include <stdint.h>
12
13#ifdef __cplusplus
14extern "C" {
15#endif
16
17/* {{{ Similarity computation context
18 */
19
20typedef struct VkSimilarityContext VkSimilarityContext;
21
22/**
23 * Initialize similarity computation context
24 * @param ctx Vulkan compute context
25 * @param embeddings All poem embeddings (num_poems × 768 floats)
26 * @param num_poems Total number of poems
27 * @param embedding_dim Dimension of embeddings (768)
28 * @return Similarity context or NULL on error
29 */
30VkSimilarityContext* vks_init(VkComputeContext* ctx,
31 const float* embeddings,
32 uint32_t num_poems,
33 uint32_t embedding_dim);
34
35/**
36 * Clean up similarity computation context
37 * @param sim_ctx Similarity context to destroy
38 */
39void vks_destroy(VkSimilarityContext* sim_ctx);
40
41/* }}} */
42
43/* The single-poem and sequential batch entry points were removed with Issue
44 * 9-002b -- the parallel full-matrix dispatch below is the only path. */
45
46/* {{{ Parallel full-matrix computation (Issue 9-002 original design)
47 */
48
49/**
50 * Compute ALL similarities in a single GPU dispatch
51 * This is the correct parallel implementation per Issue 9-002 design.
52 *
53 * @param sim_ctx Similarity context
54 * @param output_triangular Output buffer for triangular matrix
55 * Size: num_poems * (num_poems - 1) / 2 floats
56 * Layout: For pair (i,j) where i<j:
57 * index = i * num_poems - i*(i+1)/2 + (j-i-1)
58 * @return VKC_SUCCESS or error code
59 *
60 * Performance: Single dispatch with ~173K workgroups, completes in seconds
61 * vs sequential approach which takes 70+ minutes.
62 */
63VkComputeResult vks_compute_all_similarities_parallel(
64 VkSimilarityContext* sim_ctx,
65 float* output_triangular);
66
67/**
68 * Get the size of the triangular similarity matrix in floats
69 * @param num_poems Number of poems
70 * @return Size in floats: num_poems * (num_poems - 1) / 2
71 */
72static inline uint64_t vks_triangular_size(uint32_t num_poems) {
73 return ((uint64_t)num_poems * (num_poems - 1)) / 2;
74}
75
76/**
77 * Get the linear index for a pair (i, j) in the triangular matrix
78 * @param i First poem index (must be < j)
79 * @param j Second poem index (must be > i)
80 * @param num_poems Total number of poems
81 * @return Linear index into triangular buffer
82 */
83static inline uint64_t vks_triangular_index(uint32_t i, uint32_t j, uint32_t num_poems) {
84 return (uint64_t)i * num_poems - ((uint64_t)i * (i + 1)) / 2 + (j - i - 1);
85}
86
87/* }}} */
88
89/* {{{ Parallel file I/O with pthreads
90 */
91
92/**
93 * Write similarity files in parallel using pthreads
94 * This avoids Lua/effil serialization overhead by keeping all data in C.
95 *
96 * @param triangular_buffer Pre-computed triangular similarity matrix from GPU
97 * @param num_poems Total number of poems
98 * @param poem_indices Array mapping array index (0-based) to poem_index for each poem
99 * @param poem_ids Array of poem IDs (as strings, for metadata)
100 * @param output_dir Directory to write files to (e.g., "assets/embeddings/model/similarities")
101 * @param num_threads Number of worker threads (1-64)
102 * @return VKC_SUCCESS or error code
103 *
104 * File format: poem_index_{N}.json with sorted similarities
105 * Uses atomic task counter - threads grab next poem index when ready
106 */
107VkComputeResult vks_write_similarity_files_parallel(
108 const float* triangular_buffer,
109 uint32_t num_poems,
110 const uint32_t* poem_indices,
111 const char** poem_ids,
112 const char* output_dir,
113 uint32_t num_threads);
114
115/**
116 * Write similarity rankings cache file in parallel
117 * Generates the pre-sorted rankings cache that HTML generation uses.
118 *
119 * @param triangular_buffer Pre-computed triangular similarity matrix from GPU
120 * @param num_poems Total number of poems
121 * @param poem_indices Array mapping array index (0-based) to poem_index for each poem
122 * @param cache_file Path to output cache file (e.g., "assets/.../similarity_rankings_cache.json")
123 * @param num_threads Number of worker threads for parallel sorting
124 * @param top_k Keep only the top-K nearest neighbours per poem (0 = keep all).
125 * Caps the on-disk JSON and the RAM table the HTML stage parses it
126 * into (Issue 10-057).
127 * @return VKC_SUCCESS or error code
128 *
129 * Output format: JSON with rankings[poem_index] = [sorted_neighbor_indices...]
130 */
131VkComputeResult vks_write_rankings_cache_parallel(
132 const float* triangular_buffer,
133 uint32_t num_poems,
134 const uint32_t* poem_indices,
135 const char* cache_file,
136 uint32_t num_threads,
137 uint32_t top_k);
138
139/* }}} */
140
141#ifdef __cplusplus
142}
143#endif
144
145#endif /* VK_SIMILARITY_H */
146