Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,226 @@
|
||||
//
|
||||
// elementwise.metal
|
||||
// Metal kernels for elementwise NMF operations
|
||||
//
|
||||
// This file contains GPU kernels for elementwise operations
|
||||
// used in Non-negative Matrix Factorization on Apple Silicon.
|
||||
//
|
||||
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
/// Parameters for elementwise operations
|
||||
struct ElementwiseParams {
|
||||
uint total_elements; // Total number of elements to process
|
||||
float value; // Scalar value for operations that need it
|
||||
};
|
||||
|
||||
/// Non-negativity clamping kernel
|
||||
/// Ensures all elements are >= 0
|
||||
///
|
||||
/// @param data [in/out] Array to clamp
|
||||
/// @param params [in] Parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_clamp_non_negative(
|
||||
device float* data [[buffer(0)]],
|
||||
constant ElementwiseParams& params [[buffer(1)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
data[gid] = max(0.0f, data[gid]);
|
||||
}
|
||||
|
||||
/// Clamp to range kernel
|
||||
/// Clamps all elements to [min_val, max_val]
|
||||
///
|
||||
/// @param data [in/out] Array to clamp
|
||||
/// @param min_val [in] Minimum value
|
||||
/// @param max_val [in] Maximum value
|
||||
/// @param params [in] Parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_clamp_range(
|
||||
device float* data [[buffer(0)]],
|
||||
constant float& min_val [[buffer(1)]],
|
||||
constant float& max_val [[buffer(2)]],
|
||||
constant ElementwiseParams& params [[buffer(3)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
data[gid] = clamp(data[gid], min_val, max_val);
|
||||
}
|
||||
|
||||
/// Element-wise multiplication (Hadamard product)
|
||||
/// result = a * b
|
||||
///
|
||||
/// @param result [out] Output array
|
||||
/// @param a [in] First input array
|
||||
/// @param b [in] Second input array
|
||||
/// @param params [in] Parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_elementwise_mul(
|
||||
device float* result [[buffer(0)]],
|
||||
constant float* a [[buffer(1)]],
|
||||
constant float* b [[buffer(2)]],
|
||||
constant ElementwiseParams& params [[buffer(3)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
result[gid] = a[gid] * b[gid];
|
||||
}
|
||||
|
||||
/// Element-wise division with epsilon for stability
|
||||
/// result = a / (b + epsilon)
|
||||
///
|
||||
/// @param result [out] Output array
|
||||
/// @param a [in] Numerator array
|
||||
/// @param b [in] Denominator array
|
||||
/// @param epsilon [in] Small value to prevent division by zero
|
||||
/// @param params [in] Parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_elementwise_div(
|
||||
device float* result [[buffer(0)]],
|
||||
constant float* a [[buffer(1)]],
|
||||
constant float* b [[buffer(2)]],
|
||||
constant float& epsilon [[buffer(3)]],
|
||||
constant ElementwiseParams& params [[buffer(4)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
result[gid] = a[gid] / (b[gid] + epsilon);
|
||||
}
|
||||
|
||||
/// Element-wise subtraction
|
||||
/// result = a - b
|
||||
///
|
||||
/// @param result [out] Output array
|
||||
/// @param a [in] First input array
|
||||
/// @param b [in] Second input array
|
||||
/// @param params [in] Parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_elementwise_sub(
|
||||
device float* result [[buffer(0)]],
|
||||
constant float* a [[buffer(1)]],
|
||||
constant float* b [[buffer(2)]],
|
||||
constant ElementwiseParams& params [[buffer(3)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
result[gid] = a[gid] - b[gid];
|
||||
}
|
||||
|
||||
/// Scalar multiplication
|
||||
/// result = a * scalar
|
||||
///
|
||||
/// @param result [out] Output array
|
||||
/// @param a [in] Input array
|
||||
/// @param scalar [in] Scalar multiplier
|
||||
/// @param params [in] Parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_scalar_mul(
|
||||
device float* result [[buffer(0)]],
|
||||
constant float* a [[buffer(1)]],
|
||||
constant float& scalar [[buffer(2)]],
|
||||
constant ElementwiseParams& params [[buffer(3)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
result[gid] = a[gid] * scalar;
|
||||
}
|
||||
|
||||
/// Element-wise square
|
||||
/// result = a^2
|
||||
///
|
||||
/// @param result [out] Output array
|
||||
/// @param a [in] Input array
|
||||
/// @param params [in] Parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_elementwise_square(
|
||||
device float* result [[buffer(0)]],
|
||||
constant float* a [[buffer(1)]],
|
||||
constant ElementwiseParams& params [[buffer(2)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
float val = a[gid];
|
||||
result[gid] = val * val;
|
||||
}
|
||||
|
||||
/// AXPY operation: y = alpha * x + y
|
||||
/// Common linear algebra operation
|
||||
///
|
||||
/// @param y [in/out] Result/second input array
|
||||
/// @param x [in] First input array
|
||||
/// @param alpha [in] Scalar multiplier
|
||||
/// @param params [in] Parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_axpy(
|
||||
device float* y [[buffer(0)]],
|
||||
constant float* x [[buffer(1)]],
|
||||
constant float& alpha [[buffer(2)]],
|
||||
constant ElementwiseParams& params [[buffer(3)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
y[gid] = alpha * x[gid] + y[gid];
|
||||
}
|
||||
|
||||
/// Initialize array with random values in [0, 1)
|
||||
/// Uses a simple hash-based PRNG for reproducibility
|
||||
///
|
||||
/// @param data [out] Output array
|
||||
/// @param seed [in] Random seed
|
||||
/// @param params [in] Parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_random_init(
|
||||
device float* data [[buffer(0)]],
|
||||
constant uint& seed [[buffer(1)]],
|
||||
constant ElementwiseParams& params [[buffer(2)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
// Simple hash-based PRNG (xorshift)
|
||||
uint state = gid + seed;
|
||||
state ^= state << 13;
|
||||
state ^= state >> 17;
|
||||
state ^= state << 5;
|
||||
|
||||
// Convert to float in [0, 1)
|
||||
data[gid] = float(state) / float(0xFFFFFFFF);
|
||||
}
|
||||
|
||||
/// Initialize with scaled random values for NMF
|
||||
/// Initializes to sqrt(mean(V) / k) * random
|
||||
/// where k is the rank of the factorization
|
||||
///
|
||||
/// @param data [out] Output array
|
||||
/// @param scale [in] Scale factor (typically sqrt(mean(V)/k))
|
||||
/// @param seed [in] Random seed
|
||||
/// @param params [in] Parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_scaled_random_init(
|
||||
device float* data [[buffer(0)]],
|
||||
constant float& scale [[buffer(1)]],
|
||||
constant uint& seed [[buffer(2)]],
|
||||
constant ElementwiseParams& params [[buffer(3)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
// Simple hash-based PRNG (xorshift)
|
||||
uint state = gid + seed;
|
||||
state ^= state << 13;
|
||||
state ^= state >> 17;
|
||||
state ^= state << 5;
|
||||
|
||||
// Convert to float in [0, scale)
|
||||
data[gid] = scale * float(state) / float(0xFFFFFFFF);
|
||||
}
|
||||
@@ -0,0 +1,145 @@
|
||||
//
|
||||
// nmf_update.metal
|
||||
// Metal kernels for NMF multiplicative update operations
|
||||
//
|
||||
// This file contains GPU kernels for Non-negative Matrix Factorization
|
||||
// on Apple Silicon using Metal Shading Language (MSL).
|
||||
//
|
||||
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
/// Parameters for NMF matrix update operations
|
||||
struct NMFUpdateParams {
|
||||
uint rows; // Number of rows in the matrix
|
||||
uint cols; // Number of columns in the matrix
|
||||
float epsilon; // Small value to prevent division by zero
|
||||
uint total_elements; // Total elements (rows * cols)
|
||||
};
|
||||
|
||||
/// Multiplicative update kernel for H matrix
|
||||
/// H = H * (W^T V) / (W^T W H + epsilon)
|
||||
///
|
||||
/// This kernel implements the multiplicative update rule for the H matrix
|
||||
/// in standard NMF decomposition: V ≈ W * H
|
||||
///
|
||||
/// @param h_data [in/out] The H matrix to update
|
||||
/// @param wt_v_data [in] W^T V (numerator term)
|
||||
/// @param wtw_h_data [in] W^T W H (denominator base)
|
||||
/// @param params [in] Update parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_update_h(
|
||||
device float* h_data [[buffer(0)]],
|
||||
constant float* wt_v_data [[buffer(1)]],
|
||||
constant float* wtw_h_data [[buffer(2)]],
|
||||
constant NMFUpdateParams& params [[buffer(3)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
float numerator = wt_v_data[gid];
|
||||
float denominator = wtw_h_data[gid] + params.epsilon;
|
||||
float current_h = h_data[gid];
|
||||
|
||||
// Multiplicative update: H_new = H * numerator / denominator
|
||||
float updated_value = current_h * numerator / denominator;
|
||||
|
||||
// Ensure non-negativity (clamp to 0)
|
||||
h_data[gid] = max(0.0f, updated_value);
|
||||
}
|
||||
|
||||
/// Multiplicative update kernel for W matrix
|
||||
/// W = W * (V H^T) / (W H H^T + epsilon)
|
||||
///
|
||||
/// This kernel implements the multiplicative update rule for the W matrix
|
||||
/// in standard NMF decomposition: V ≈ W * H
|
||||
///
|
||||
/// @param w_data [in/out] The W matrix to update
|
||||
/// @param v_ht_data [in] V H^T (numerator term)
|
||||
/// @param w_hht_data [in] W H H^T (denominator base)
|
||||
/// @param params [in] Update parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_update_w(
|
||||
device float* w_data [[buffer(0)]],
|
||||
constant float* v_ht_data [[buffer(1)]],
|
||||
constant float* w_hht_data [[buffer(2)]],
|
||||
constant NMFUpdateParams& params [[buffer(3)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
float numerator = v_ht_data[gid];
|
||||
float denominator = w_hht_data[gid] + params.epsilon;
|
||||
float current_w = w_data[gid];
|
||||
|
||||
// Multiplicative update: W_new = W * numerator / denominator
|
||||
float updated_value = current_w * numerator / denominator;
|
||||
|
||||
// Ensure non-negativity (clamp to 0)
|
||||
w_data[gid] = max(0.0f, updated_value);
|
||||
}
|
||||
|
||||
/// Fused multiplicative update for H matrix with stability checks
|
||||
/// This version includes additional numerical stability measures
|
||||
///
|
||||
/// @param h_data [in/out] The H matrix to update
|
||||
/// @param wt_v_data [in] W^T V (numerator term)
|
||||
/// @param wtw_h_data [in] W^T W H (denominator base)
|
||||
/// @param params [in] Update parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_update_h_stable(
|
||||
device float* h_data [[buffer(0)]],
|
||||
constant float* wt_v_data [[buffer(1)]],
|
||||
constant float* wtw_h_data [[buffer(2)]],
|
||||
constant NMFUpdateParams& params [[buffer(3)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
float numerator = wt_v_data[gid];
|
||||
float denominator = wtw_h_data[gid] + params.epsilon;
|
||||
float current_h = h_data[gid];
|
||||
|
||||
// Multiplicative update with stability
|
||||
float ratio = numerator / denominator;
|
||||
|
||||
// Clamp ratio to prevent exploding updates
|
||||
ratio = clamp(ratio, 0.0f, 10.0f);
|
||||
|
||||
float updated_value = current_h * ratio;
|
||||
|
||||
// Ensure non-negativity and prevent very small values
|
||||
h_data[gid] = max(params.epsilon, updated_value);
|
||||
}
|
||||
|
||||
/// Fused multiplicative update for W matrix with stability checks
|
||||
///
|
||||
/// @param w_data [in/out] The W matrix to update
|
||||
/// @param v_ht_data [in] V H^T (numerator term)
|
||||
/// @param w_hht_data [in] W H H^T (denominator base)
|
||||
/// @param params [in] Update parameters
|
||||
/// @param gid Thread index
|
||||
kernel void nmf_update_w_stable(
|
||||
device float* w_data [[buffer(0)]],
|
||||
constant float* v_ht_data [[buffer(1)]],
|
||||
constant float* w_hht_data [[buffer(2)]],
|
||||
constant NMFUpdateParams& params [[buffer(3)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= params.total_elements) return;
|
||||
|
||||
float numerator = v_ht_data[gid];
|
||||
float denominator = w_hht_data[gid] + params.epsilon;
|
||||
float current_w = w_data[gid];
|
||||
|
||||
// Multiplicative update with stability
|
||||
float ratio = numerator / denominator;
|
||||
|
||||
// Clamp ratio to prevent exploding updates
|
||||
ratio = clamp(ratio, 0.0f, 10.0f);
|
||||
|
||||
float updated_value = current_w * ratio;
|
||||
|
||||
// Ensure non-negativity and prevent very small values
|
||||
w_data[gid] = max(params.epsilon, updated_value);
|
||||
}
|
||||
@@ -0,0 +1,394 @@
|
||||
//
|
||||
// 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];
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user