/**
* mst-kernel.cl - GPU kernels for MST parsing
*
* Strategy: Batch process many sentences in parallel.
* Each work-item handles one sentence's edge weight lookups.
*
* For MST itself, we use a simple O(V²) approach per sentence
* since sentences are small (12 words max).
*/
/**
* batch_mi_lookup - Look up MI values for many word pairs at once
*
* The MI matrix is stored as a flattened 2D array: mi[w1 * vocab_size + w2]
* We look up pairs in parallel.
*
* @param mi_matrix Flattened MI matrix (vocab_size × vocab_size)
* @param word_indices Word index pairs: [w1_0, w2_0, w1_1, w2_1, ...]
* @param mi_values Output MI values
* @param vocab_size Size of vocabulary
* @param sz Number of pairs to look up
*/
__kernel void batch_mi_lookup(
__global const double* mi_matrix,
__global const int* word_indices,
__global double* mi_values,
const int vocab_size,
const unsigned long sz)
{
size_t i = get_global_id(0);
if (i >= sz) return;
int w1 = word_indices[2 * i];
int w2 = word_indices[2 * i + 1];
// Handle out-of-vocab words
if (w1 < 0 || w2 < 0 || w1 >= vocab_size || w2 >= vocab_size) {
mi_values[i] = 0.0;
return;
}
// Look up MI value (symmetric, but we store both directions)
mi_values[i] = mi_matrix[w1 * vocab_size + w2];
}
/**
* prim_mst_batch - Run Prim's MST on multiple sentences in parallel
*
* Each work-item handles one sentence.
* Sentences are stored as word indices with -1 padding.
*
* @param sentences Word indices for each sentence [max_len × n_sentences]
* @param sentence_lens Actual length of each sentence
* @param mi_matrix Flattened MI matrix
* @param vocab_size Vocabulary size
* @param max_len Maximum sentence length (e.g., 12)
* @param mst_edges Output: MST edges per sentence [max_edges × n_sentences]
* @param mst_edge_count Output: Number of MST edges per sentence
* @param sz Number of sentences
*/
__kernel void prim_mst_batch(
__global const int* sentences,
__global const int* sentence_lens,
__global const double* mi_matrix,
const int vocab_size,
const int max_len,
__global int* mst_edges,
__global int* mst_edge_count,
const unsigned long sz)
{
size_t sent_idx = get_global_id(0);
if (sent_idx >= sz) return;
int n = sentence_lens[sent_idx];
if (n < 2) {
mst_edge_count[sent_idx] = 0;
return;
}
// Get this sentence's words
__global const int* words = sentences + sent_idx * max_len;
__global int* edges = mst_edges + sent_idx * (max_len - 1) * 2;
// Prim's algorithm: track which nodes are in tree
bool in_tree[32]; // Assumes max_len <= 32
for (int i = 0; i < n; i++) in_tree[i] = false;
in_tree[0] = true;
int edge_count = 0;
// Add n-1 edges
for (int iter = 0; iter < n - 1; iter++) {
double best_weight = -1e30;
int best_from = -1;
int best_to = -1;
// Find max weight edge from tree to non-tree
for (int from = 0; from < n; from++) {
if (!in_tree[from]) continue;
int w1 = words[from];
if (w1 < 0) continue;
for (int to = 0; to < n; to++) {
if (in_tree[to]) continue;
int w2 = words[to];
if (w2 < 0) continue;
// Look up MI
double weight = mi_matrix[w1 * vocab_size + w2];
if (weight > best_weight) {
best_weight = weight;
best_from = from;
best_to = to;
}
}
}
if (best_from >= 0) {
// Add edge
edges[edge_count * 2] = best_from;
edges[edge_count * 2 + 1] = best_to;
edge_count++;
in_tree[best_to] = true;
}
}
mst_edge_count[sent_idx] = edge_count;
}
/**
* count_sections_batch - Count sections from MST edges
*
* For each word, count how many left and right connections it has.
* This creates a histogram for Section types.
*
* @param sentences Word indices [max_len × n_sentences]
* @param sentence_lens Sentence lengths
* @param mst_edges MST edges [(max_len-1) × 2 × n_sentences]
* @param mst_edge_count Number of edges per sentence
* @param max_len Max sentence length
* @param section_counts Atomic counters for each section type [vocab × max_connectors × 2]
* @param sz Number of sentences
*/
__kernel void count_sections_batch(
__global const int* sentences,
__global const int* sentence_lens,
__global const int* mst_edges,
__global const int* mst_edge_count,
const int max_len,
__global atomic_int* section_counts,
const int max_connectors,
const unsigned long sz)
{
size_t sent_idx = get_global_id(0);
if (sent_idx >= sz) return;
int n = sentence_lens[sent_idx];
__global const int* words = sentences + sent_idx * max_len;
__global const int* edges = mst_edges + sent_idx * (max_len - 1) * 2;
int n_edges = mst_edge_count[sent_idx];
// For each edge, update section counts for both endpoints
for (int e = 0; e < n_edges; e++) {
int pos1 = edges[e * 2];
int pos2 = edges[e * 2 + 1];
int word1 = words[pos1];
int word2 = words[pos2];
if (word1 < 0 || word2 < 0) continue;
// Determine direction based on position
if (pos1 < pos2) {
// word1 is left of word2
// word1 gets a right connector
atomic_fetch_add(§ion_counts[word1 * max_connectors * 2 + 1], 1);
// word2 gets a left connector
atomic_fetch_add(§ion_counts[word2 * max_connectors * 2 + 0], 1);
} else {
// word2 is left of word1
atomic_fetch_add(§ion_counts[word2 * max_connectors * 2 + 1], 1);
atomic_fetch_add(§ion_counts[word1 * max_connectors * 2 + 0], 1);
}
}
}