395 lines
13 KiB
Metal
395 lines
13 KiB
Metal
//
|
|
// reduction.metal
|
|
// Metal kernels for NMF reduction operations
|
|
//
|
|
// This file contains GPU kernels for parallel reduction operations
|
|
// used in Non-negative Matrix Factorization on Apple Silicon.
|
|
//
|
|
|
|
#include <metal_stdlib>
|
|
using namespace metal;
|
|
|
|
/// Parameters for reduction operations
|
|
struct ReductionParams {
|
|
uint total_elements; // Total number of elements to reduce
|
|
uint block_size; // Threads per threadgroup
|
|
};
|
|
|
|
/// Threadgroup size for reductions (must match dispatch)
|
|
constant uint REDUCTION_BLOCK_SIZE = 256;
|
|
|
|
/// First-pass reduction: compute squared values and partial sums
|
|
/// Used for Frobenius norm calculation: ||A||_F = sqrt(sum(a_ij^2))
|
|
///
|
|
/// @param data [in] Input data array
|
|
/// @param partial_sums [out] Partial sum for each threadgroup
|
|
/// @param params [in] Reduction parameters
|
|
/// @param tid Thread index within threadgroup
|
|
/// @param gid Global thread index
|
|
/// @param threadgroup_id Threadgroup index
|
|
/// @param shared Threadgroup shared memory
|
|
kernel void nmf_frobenius_norm_squared(
|
|
constant float* data [[buffer(0)]],
|
|
device float* partial_sums [[buffer(1)]],
|
|
constant ReductionParams& params [[buffer(2)]],
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
uint gid [[thread_position_in_grid]],
|
|
uint threadgroup_id [[threadgroup_position_in_grid]],
|
|
threadgroup float* shared [[threadgroup(0)]]
|
|
) {
|
|
// Load data and compute square
|
|
float value = 0.0f;
|
|
if (gid < params.total_elements) {
|
|
value = data[gid];
|
|
value = value * value;
|
|
}
|
|
shared[tid] = value;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Parallel reduction in shared memory
|
|
for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
shared[tid] += shared[tid + s];
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
// Write result for this threadgroup
|
|
if (tid == 0) {
|
|
partial_sums[threadgroup_id] = shared[0];
|
|
}
|
|
}
|
|
|
|
/// Second-pass reduction: sum partial sums from first pass
|
|
/// This kernel is used to finalize multi-block reductions
|
|
///
|
|
/// @param partial_sums [in/out] Partial sums to reduce
|
|
/// @param num_blocks [in] Number of partial sums
|
|
/// @param tid Thread index within threadgroup
|
|
/// @param shared Threadgroup shared memory
|
|
kernel void nmf_reduce_partial_sums(
|
|
device float* partial_sums [[buffer(0)]],
|
|
constant uint& num_blocks [[buffer(1)]],
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
threadgroup float* shared [[threadgroup(0)]]
|
|
) {
|
|
// Load partial sums
|
|
float value = 0.0f;
|
|
if (tid < num_blocks) {
|
|
value = partial_sums[tid];
|
|
}
|
|
shared[tid] = value;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Parallel reduction
|
|
for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
|
|
if (tid < s && tid + s < num_blocks) {
|
|
shared[tid] += shared[tid + s];
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
// Write final result
|
|
if (tid == 0) {
|
|
partial_sums[0] = shared[0];
|
|
}
|
|
}
|
|
|
|
/// Sum reduction (first pass)
|
|
///
|
|
/// @param data [in] Input data array
|
|
/// @param partial_sums [out] Partial sum for each threadgroup
|
|
/// @param params [in] Reduction parameters
|
|
/// @param tid Thread index within threadgroup
|
|
/// @param gid Global thread index
|
|
/// @param threadgroup_id Threadgroup index
|
|
/// @param shared Threadgroup shared memory
|
|
kernel void nmf_sum_reduction(
|
|
constant float* data [[buffer(0)]],
|
|
device float* partial_sums [[buffer(1)]],
|
|
constant ReductionParams& params [[buffer(2)]],
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
uint gid [[thread_position_in_grid]],
|
|
uint threadgroup_id [[threadgroup_position_in_grid]],
|
|
threadgroup float* shared [[threadgroup(0)]]
|
|
) {
|
|
// Load data
|
|
float value = 0.0f;
|
|
if (gid < params.total_elements) {
|
|
value = data[gid];
|
|
}
|
|
shared[tid] = value;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Parallel reduction
|
|
for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
shared[tid] += shared[tid + s];
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
// Write result for this threadgroup
|
|
if (tid == 0) {
|
|
partial_sums[threadgroup_id] = shared[0];
|
|
}
|
|
}
|
|
|
|
/// Max reduction (first pass)
|
|
///
|
|
/// @param data [in] Input data array
|
|
/// @param partial_max [out] Partial max for each threadgroup
|
|
/// @param params [in] Reduction parameters
|
|
/// @param tid Thread index within threadgroup
|
|
/// @param gid Global thread index
|
|
/// @param threadgroup_id Threadgroup index
|
|
/// @param shared Threadgroup shared memory
|
|
kernel void nmf_max_reduction(
|
|
constant float* data [[buffer(0)]],
|
|
device float* partial_max [[buffer(1)]],
|
|
constant ReductionParams& params [[buffer(2)]],
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
uint gid [[thread_position_in_grid]],
|
|
uint threadgroup_id [[threadgroup_position_in_grid]],
|
|
threadgroup float* shared [[threadgroup(0)]]
|
|
) {
|
|
// Load data
|
|
float value = -INFINITY;
|
|
if (gid < params.total_elements) {
|
|
value = data[gid];
|
|
}
|
|
shared[tid] = value;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Parallel max reduction
|
|
for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
shared[tid] = max(shared[tid], shared[tid + s]);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
// Write result for this threadgroup
|
|
if (tid == 0) {
|
|
partial_max[threadgroup_id] = shared[0];
|
|
}
|
|
}
|
|
|
|
/// Max reduction second pass
|
|
///
|
|
/// @param partial_max [in/out] Partial max values to reduce
|
|
/// @param num_blocks [in] Number of partial max values
|
|
/// @param tid Thread index within threadgroup
|
|
/// @param shared Threadgroup shared memory
|
|
kernel void nmf_reduce_partial_max(
|
|
device float* partial_max [[buffer(0)]],
|
|
constant uint& num_blocks [[buffer(1)]],
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
threadgroup float* shared [[threadgroup(0)]]
|
|
) {
|
|
// Load partial max values
|
|
float value = -INFINITY;
|
|
if (tid < num_blocks) {
|
|
value = partial_max[tid];
|
|
}
|
|
shared[tid] = value;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Parallel max reduction
|
|
for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
|
|
if (tid < s && tid + s < num_blocks) {
|
|
shared[tid] = max(shared[tid], shared[tid + s]);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
// Write final result
|
|
if (tid == 0) {
|
|
partial_max[0] = shared[0];
|
|
}
|
|
}
|
|
|
|
/// Min reduction (first pass)
|
|
///
|
|
/// @param data [in] Input data array
|
|
/// @param partial_min [out] Partial min for each threadgroup
|
|
/// @param params [in] Reduction parameters
|
|
/// @param tid Thread index within threadgroup
|
|
/// @param gid Global thread index
|
|
/// @param threadgroup_id Threadgroup index
|
|
/// @param shared Threadgroup shared memory
|
|
kernel void nmf_min_reduction(
|
|
constant float* data [[buffer(0)]],
|
|
device float* partial_min [[buffer(1)]],
|
|
constant ReductionParams& params [[buffer(2)]],
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
uint gid [[thread_position_in_grid]],
|
|
uint threadgroup_id [[threadgroup_position_in_grid]],
|
|
threadgroup float* shared [[threadgroup(0)]]
|
|
) {
|
|
// Load data
|
|
float value = INFINITY;
|
|
if (gid < params.total_elements) {
|
|
value = data[gid];
|
|
}
|
|
shared[tid] = value;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Parallel min reduction
|
|
for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
shared[tid] = min(shared[tid], shared[tid + s]);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
// Write result for this threadgroup
|
|
if (tid == 0) {
|
|
partial_min[threadgroup_id] = shared[0];
|
|
}
|
|
}
|
|
|
|
/// Convergence check: compute max absolute difference
|
|
/// Used to check NMF convergence: max|H_new - H_old| < tolerance
|
|
///
|
|
/// @param new_data [in] New matrix values
|
|
/// @param old_data [in] Old matrix values
|
|
/// @param partial_max_diff [out] Partial max differences
|
|
/// @param params [in] Reduction parameters
|
|
/// @param tid Thread index within threadgroup
|
|
/// @param gid Global thread index
|
|
/// @param threadgroup_id Threadgroup index
|
|
/// @param shared Threadgroup shared memory
|
|
kernel void nmf_max_abs_diff(
|
|
constant float* new_data [[buffer(0)]],
|
|
constant float* old_data [[buffer(1)]],
|
|
device float* partial_max_diff [[buffer(2)]],
|
|
constant ReductionParams& params [[buffer(3)]],
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
uint gid [[thread_position_in_grid]],
|
|
uint threadgroup_id [[threadgroup_position_in_grid]],
|
|
threadgroup float* shared [[threadgroup(0)]]
|
|
) {
|
|
// Compute absolute difference
|
|
float diff = 0.0f;
|
|
if (gid < params.total_elements) {
|
|
diff = abs(new_data[gid] - old_data[gid]);
|
|
}
|
|
shared[tid] = diff;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Parallel max reduction
|
|
for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
shared[tid] = max(shared[tid], shared[tid + s]);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
// Write result for this threadgroup
|
|
if (tid == 0) {
|
|
partial_max_diff[threadgroup_id] = shared[0];
|
|
}
|
|
}
|
|
|
|
/// Reconstruction error: compute ||V - WH||_F^2
|
|
/// Used to measure NMF quality: smaller is better
|
|
///
|
|
/// @param v_data [in] Original matrix V
|
|
/// @param wh_data [in] Reconstructed matrix W*H
|
|
/// @param partial_error [out] Partial squared errors
|
|
/// @param params [in] Reduction parameters
|
|
/// @param tid Thread index within threadgroup
|
|
/// @param gid Global thread index
|
|
/// @param threadgroup_id Threadgroup index
|
|
/// @param shared Threadgroup shared memory
|
|
kernel void nmf_reconstruction_error(
|
|
constant float* v_data [[buffer(0)]],
|
|
constant float* wh_data [[buffer(1)]],
|
|
device float* partial_error [[buffer(2)]],
|
|
constant ReductionParams& params [[buffer(3)]],
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
uint gid [[thread_position_in_grid]],
|
|
uint threadgroup_id [[threadgroup_position_in_grid]],
|
|
threadgroup float* shared [[threadgroup(0)]]
|
|
) {
|
|
// Compute squared difference
|
|
float diff = 0.0f;
|
|
if (gid < params.total_elements) {
|
|
diff = v_data[gid] - wh_data[gid];
|
|
diff = diff * diff;
|
|
}
|
|
shared[tid] = diff;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Parallel reduction
|
|
for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
shared[tid] += shared[tid + s];
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
// Write result for this threadgroup
|
|
if (tid == 0) {
|
|
partial_error[threadgroup_id] = shared[0];
|
|
}
|
|
}
|
|
|
|
/// KL divergence: compute D_KL(V || WH)
|
|
/// Alternative measure for NMF quality
|
|
/// D_KL = sum(V * log(V/WH) - V + WH)
|
|
///
|
|
/// @param v_data [in] Original matrix V
|
|
/// @param wh_data [in] Reconstructed matrix W*H
|
|
/// @param partial_div [out] Partial divergence values
|
|
/// @param epsilon [in] Small value for numerical stability
|
|
/// @param params [in] Reduction parameters
|
|
/// @param tid Thread index within threadgroup
|
|
/// @param gid Global thread index
|
|
/// @param threadgroup_id Threadgroup index
|
|
/// @param shared Threadgroup shared memory
|
|
kernel void nmf_kl_divergence(
|
|
constant float* v_data [[buffer(0)]],
|
|
constant float* wh_data [[buffer(1)]],
|
|
device float* partial_div [[buffer(2)]],
|
|
constant float& epsilon [[buffer(3)]],
|
|
constant ReductionParams& params [[buffer(4)]],
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
uint gid [[thread_position_in_grid]],
|
|
uint threadgroup_id [[threadgroup_position_in_grid]],
|
|
threadgroup float* shared [[threadgroup(0)]]
|
|
) {
|
|
// Compute KL divergence term
|
|
float div = 0.0f;
|
|
if (gid < params.total_elements) {
|
|
float v = v_data[gid] + epsilon;
|
|
float wh = wh_data[gid] + epsilon;
|
|
div = v * log(v / wh) - v + wh;
|
|
}
|
|
shared[tid] = div;
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
// Parallel reduction
|
|
for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
shared[tid] += shared[tid + s];
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
// Write result for this threadgroup
|
|
if (tid == 0) {
|
|
partial_div[threadgroup_id] = shared[0];
|
|
}
|
|
}
|