3d47813f1e
eye diagram working for all signals again... A problem: I found that the Tx eyes look weird (i.e. PAM-8) vibecoded ML_MLSE_GPU ... a bit faster :-)
526 lines
23 KiB
Matlab
526 lines
23 KiB
Matlab
classdef ML_MLSE_GPU < handle
|
|
% ---------------------------------------------------------------------
|
|
% W. Lanneer and Y. Lefevre,
|
|
% “Machine Learning-Based Pre-Equalizers for Maximum Likelihood
|
|
% Sequence Estimation in High-Speed PONs,” EUSIPCO 2023
|
|
% ---------------------------------------------------------------------
|
|
% This implementation reproduces the closed-loop ML-based
|
|
% pre-equalizer training for MLSE, supporting both training and
|
|
% detection (decision-directed) modes.
|
|
% ---------------------------------------------------------------------
|
|
|
|
properties
|
|
sps
|
|
order
|
|
e
|
|
e_tr
|
|
error
|
|
|
|
len_tr
|
|
mu_tr
|
|
epochs_tr
|
|
|
|
dd_mode
|
|
mu_dd
|
|
epochs_dd
|
|
|
|
adaptive_mu
|
|
|
|
constellation
|
|
L
|
|
alpha
|
|
DIR
|
|
DIR_flip
|
|
trellis_states
|
|
traceback_depth
|
|
delta
|
|
|
|
% --- New: GPU Support ---
|
|
use_gpu
|
|
|
|
% Internal variables
|
|
S
|
|
Nf
|
|
nStates
|
|
nFeasible
|
|
combs
|
|
first_sym
|
|
last_sym
|
|
valid
|
|
valid_to_idx
|
|
valid_from_idx
|
|
w
|
|
|
|
% Fast lookup
|
|
nSym
|
|
key_table
|
|
trans_index
|
|
true_to_state_idx
|
|
|
|
% Debug metrics
|
|
ber = []
|
|
ce = ones(1,1)
|
|
end
|
|
|
|
methods
|
|
function obj = ML_MLSE_GPU(options)
|
|
arguments(Input)
|
|
options.sps = 2;
|
|
options.order = 15;
|
|
options.len_tr = 4096;
|
|
options.mu_tr = 0.001;
|
|
options.epochs_tr = 5;
|
|
options.dd_mode = 1;
|
|
options.mu_dd = 1e-5;
|
|
options.epochs_dd = 5;
|
|
options.adaptive_mu = 1;
|
|
options.delta = 0;
|
|
options.traceback_depth = 1024;
|
|
options.L = 1;
|
|
options.use_gpu = 0; % Default: CPU
|
|
end
|
|
|
|
fn = fieldnames(options);
|
|
for n = 1:numel(fn)
|
|
obj.(fn{n}) = options.(fn{n});
|
|
end
|
|
|
|
% Check GPU availability
|
|
if obj.use_gpu
|
|
if gpuDeviceCount() > 0
|
|
% GPU is available, keeping use_gpu = 1
|
|
else
|
|
warning('GPU requested but not available. Falling back to CPU.');
|
|
obj.use_gpu = 0;
|
|
end
|
|
end
|
|
|
|
obj.e = zeros(obj.order,1);
|
|
obj.error = 0;
|
|
end
|
|
|
|
% ==============================================================
|
|
% PROCESS
|
|
% ==============================================================
|
|
function [X,X_viterbi] = process(obj, X, D)
|
|
% Normalize input RMS
|
|
X = X.normalize("mode","rms");
|
|
obj.constellation = sort(unique(D.signal),'ascend');
|
|
obj.nSym = numel(obj.constellation);
|
|
|
|
if length(X)/length(D) ~= obj.sps
|
|
warning('Signal length does not fit to reference!');
|
|
end
|
|
|
|
% --- Parameters
|
|
obj.S = obj.nSym;
|
|
obj.Nf = obj.order * obj.sps;
|
|
obj.nStates = obj.S^obj.L;
|
|
obj.nFeasible = obj.nStates * obj.S;
|
|
|
|
% --- Trellis mapping
|
|
obj.trellis_states = reshape(obj.constellation,1,[]);
|
|
pre_comb_mat = repmat(obj.trellis_states, obj.L, 1);
|
|
pre_comb_cell = mat2cell(pre_comb_mat, ones(1,obj.L), size(pre_comb_mat,2));
|
|
obj.combs = fliplr(combvec(pre_comb_cell{:}).');
|
|
obj.first_sym = obj.combs(:,1);
|
|
obj.last_sym = obj.combs(:,end);
|
|
obj.nStates = size(obj.combs,1);
|
|
|
|
% --- Valid transitions
|
|
obj.valid = false(obj.nStates);
|
|
for from = 1:obj.nStates
|
|
for to = 1:obj.nStates
|
|
if all(obj.combs(to,2:end) == obj.combs(from,1:end-1))
|
|
obj.valid(to,from) = true;
|
|
end
|
|
end
|
|
end
|
|
[obj.valid_to_idx,obj.valid_from_idx] = find(obj.valid);
|
|
|
|
% --- Initialize weights
|
|
if isempty(obj.w) || any(size(obj.w) ~= [obj.Nf+1,obj.nFeasible])
|
|
obj.w = randn(obj.Nf+1,obj.nFeasible);
|
|
obj.w = zeros(obj.Nf+1,obj.nFeasible);
|
|
end
|
|
|
|
% --- Fast lookup tables
|
|
[~, sym_idx_mat] = ismember(obj.combs, obj.constellation);
|
|
key_vals = 1 + sum((sym_idx_mat - 1) .* (obj.nSym .^ (0:obj.L-1)), 2);
|
|
max_key = obj.nSym^obj.L;
|
|
obj.key_table = zeros(max_key,1,'uint32');
|
|
obj.key_table(key_vals) = 1:obj.nStates;
|
|
|
|
obj.trans_index = sparse(obj.nStates,obj.nStates);
|
|
for i = 1:length(obj.valid_from_idx)
|
|
f = obj.valid_from_idx(i);
|
|
t = obj.valid_to_idx(i);
|
|
obj.trans_index(t,f) = i;
|
|
end
|
|
|
|
% ==============================================================
|
|
% TRAINING
|
|
% ==============================================================
|
|
fprintf('\n--- Training mode ---\n');
|
|
% Always run training on CPU to avoid loop latency on GPU
|
|
obj.equalize(X.signal, D.signal, obj.mu_tr, obj.epochs_tr, obj.len_tr, true);
|
|
obj.e_tr = obj.e;
|
|
|
|
% ==============================================================
|
|
% DECISION-DIRECTED / TESTING
|
|
% ==============================================================
|
|
fprintf('--- Decision-directed / detection mode ---\n');
|
|
|
|
x_sig = X.signal;
|
|
d_sig = D.signal;
|
|
|
|
% Move to GPU only for inference if requested
|
|
if obj.use_gpu
|
|
try
|
|
x_sig = gpuArray(single(x_sig));
|
|
if isa(obj.w, 'double')
|
|
obj.w = gpuArray(single(obj.w));
|
|
end
|
|
catch
|
|
warning('Failed to move data to GPU. Falling back to CPU.');
|
|
obj.use_gpu = 0;
|
|
end
|
|
end
|
|
|
|
[y, y_vit] = obj.equalize(x_sig, d_sig, obj.mu_dd, obj.epochs_dd, X.length, false);
|
|
|
|
% Gather results back to CPU
|
|
if obj.use_gpu
|
|
y = gather(y);
|
|
y_vit = gather(y_vit);
|
|
obj.w = gather(obj.w);
|
|
end
|
|
|
|
X_viterbi = X;
|
|
X.signal = y;
|
|
X_viterbi.signal = y_vit;
|
|
end
|
|
|
|
% ==============================================================
|
|
% EQUALIZE
|
|
% ==============================================================
|
|
function [y,y_ref] = equalize(obj,x,d,mu,epochs,N,training)
|
|
debug = 1;
|
|
showPlots = 1;
|
|
|
|
% Ensure basic types
|
|
if obj.use_gpu && ~isa(x, 'gpuArray') && training % Only force gpu for training if desired, but here we focus on inference
|
|
% For training, we keep existing flow for now.
|
|
end
|
|
|
|
y = zeros(N,1);
|
|
% On GPU, y should be gpuArray if we build it there, but we return it at the end.
|
|
if obj.use_gpu && ~training
|
|
y = gpuArray.zeros(N,1);
|
|
end
|
|
|
|
nSymbols = ceil(N/obj.sps);
|
|
|
|
for epoch = 1:epochs
|
|
pm = zeros(obj.nStates,1);
|
|
pred = zeros(nSymbols,obj.nStates,'uint32');
|
|
% pred is large uint32, GPU support for uint32 exists but sometimes limited.
|
|
% We'll keep pred on CPU for Viterbi path storage or gather per chunk if needed.
|
|
|
|
pm_sto = nan(obj.nStates,nSymbols,'like',pm);
|
|
CE_accum = 0;
|
|
|
|
start_sample = 1;
|
|
end_sample = N;
|
|
start_symbol = 1 + floor((start_sample - 1)/obj.sps);
|
|
|
|
% --- initialize true state
|
|
if numel(d) >= obj.L && start_symbol >= obj.L
|
|
init_seq = d(start_symbol-obj.L+1:start_symbol);
|
|
key_init = obj.seq2key(init_seq);
|
|
true_to_state_idx = obj.key_table(key_init);
|
|
if true_to_state_idx==0, true_to_state_idx=1; end
|
|
else
|
|
true_to_state_idx = uint32(1);
|
|
end
|
|
|
|
% =========================================================
|
|
% INFERENCE OPTIMIZATION (Vectorized Branch Metrics)
|
|
% =========================================================
|
|
run_vectorized = ~training && obj.use_gpu;
|
|
|
|
% --- Pre-calculate symbol indices for fast key generation (Training Only)
|
|
% This avoids slow ismember() calls inside the loop
|
|
d_indices = [];
|
|
if training
|
|
% Map d to 0..M-1 indices once
|
|
[~, d_indices] = ismember(d, obj.constellation);
|
|
d_indices = d_indices - 1; % 0-based
|
|
end
|
|
|
|
if run_vectorized
|
|
% 1. Construct Sliding Window Input Matrix
|
|
% Windows corresponding to symbol centers: start_sample:sps:end_sample
|
|
% Each window is [x(sample-Nf+1+delta : sample+delta); 1]
|
|
|
|
% Create index matrix
|
|
num_syms = numel(start_sample:obj.sps:end_sample);
|
|
samples_idx = start_sample + (0:num_syms-1)*obj.sps; % [1 x nSymbols]
|
|
|
|
% Indices for window: relative -Nf+1+delta to +delta
|
|
rel_idx = (-obj.Nf + 1 + obj.delta : obj.delta)'; % [Nf x 1]
|
|
|
|
% Full index matrix (implicit expansion)
|
|
idx_mat = rel_idx + samples_idx; % [Nf x nSymbols]
|
|
|
|
% Handle boundary padding (clamping indices)
|
|
% Since x is gpuArray, indexing with clamp is efficient if rewritten
|
|
% But MATLAB indexing x(idx_mat) with OOB indices is tricky vectorized.
|
|
% Faster approach: Clamp indices to [1, length(x)] and mask zero-pads.
|
|
|
|
mask_valid = (idx_mat >= 1) & (idx_mat <= length(x));
|
|
idx_clamped = max(1, min(length(x), idx_mat));
|
|
|
|
X_windows = x(idx_clamped); % [Nf x nSymbols]
|
|
X_windows = X_windows .* mask_valid; % Zero pad out-of-bounds
|
|
|
|
% Add bias row
|
|
X_windows = [X_windows; ones(1, num_syms, 'like', x)]; % [(Nf+1) x nSymbols]
|
|
|
|
% 2. Bulk Compute Branch Metrics
|
|
% C_all: [nFeasible x nSymbols] = ( [(Nf+1) x nFeasible] )' * [(Nf+1) x nSymbols]
|
|
% obj.w usually (Nf+1)xFeasible
|
|
C_all = obj.w' * X_windows;
|
|
|
|
% 3. Bring Metrics to CPU for Viterbi
|
|
% Processing Viterbi on CPU is often faster than serial kernel launches on GPU
|
|
C_all_cpu = gather(C_all);
|
|
|
|
% Pre-computation done. Now loop for Viterbi (Add-Compare-Select)
|
|
|
|
% Prepare CPU variables
|
|
pm = gather(pm);
|
|
pm_sto = gather(pm_sto);
|
|
pred = gather(pred);
|
|
v_tilde_mat = inf(obj.nStates, obj.nStates);
|
|
|
|
% Loop over symbols (pure CPU Viterbi)
|
|
for sym_i = 1:num_syms
|
|
% Current branch metrics for all edges
|
|
c_hat_curr = C_all_cpu(:, sym_i); % [nFeasible x 1]
|
|
|
|
% ACS Update
|
|
pm = pm - min(pm);
|
|
v_tilde = pm(obj.valid_from_idx) + c_hat_curr;
|
|
|
|
v_tilde_mat(obj.valid) = v_tilde;
|
|
[pm_next, pred(sym_i,:)] = min(v_tilde_mat, [], 2);
|
|
|
|
pm_next = pm_next - min(pm_next);
|
|
pm = pm_next;
|
|
pm_sto(:,sym_i) = pm;
|
|
end
|
|
|
|
symbol = num_syms; % Update for traceback
|
|
|
|
else
|
|
% =========================================================
|
|
% STANDARD SEQUENTIAL LOOP (Training / CPU Inference)
|
|
% =========================================================
|
|
pow_vec = (obj.nSym .^ (0:obj.L-1)).'; % Pre-calc powers for keygen
|
|
|
|
for sample = start_sample:obj.sps:end_sample
|
|
symbol = (sample - start_sample)/obj.sps + 1;
|
|
sym_idx = start_symbol + (symbol - 1);
|
|
|
|
% --- Observation window (with delta)
|
|
i1 = sample - obj.Nf + 1 + obj.delta;
|
|
i2 = sample + obj.delta;
|
|
buf = x(max(1,i1):min(length(x),i2));
|
|
padL = max(0,1 - i1);
|
|
padR = max(0,i2 - length(x));
|
|
yk = [zeros(padL,1); buf(:); zeros(padR,1)];
|
|
yk = [yk;1];
|
|
|
|
% --- Branch metrics
|
|
c_hat = (yk.' * obj.w).';
|
|
pm = pm - min(pm);
|
|
v_tilde = pm(obj.valid_from_idx) + c_hat;
|
|
|
|
% --- allocate once
|
|
if epoch==1 && symbol==1
|
|
obj.true_to_state_idx = ones(ceil(N/obj.sps),1,'uint32');
|
|
end
|
|
|
|
% --- previous "to" becomes "from"
|
|
if symbol>1
|
|
true_from_state_idx = obj.true_to_state_idx(symbol-1);
|
|
else
|
|
true_from_state_idx = 1;
|
|
end
|
|
|
|
% --- compute or reuse "to" state
|
|
if epoch==1
|
|
if sym_idx>=obj.L
|
|
% OPTIMIZED KEY GENERATION: Use pre-calculated indices
|
|
% key_to = obj.seq2key(d(sym_idx-obj.L+1:sym_idx));
|
|
|
|
% Extract subsequence of indices (flip needed as per seq2key logic?)
|
|
% seq2key does: ismember(flip(seq)...).
|
|
% d_indices is 0-based index of d.
|
|
% We want indices of d(sym_idx-obj.L+1 : sym_idx)
|
|
|
|
d_sub = d_indices(sym_idx-obj.L+1 : sym_idx);
|
|
% seq2key flips the sequence.
|
|
% So we need to efficiently calculate scalar key from d_sub.
|
|
% key = 1 + sum( flip(d_sub) .* pow );
|
|
% Let's do it manually to be fast
|
|
|
|
key_to = 1 + sum(flip(d_sub) .* pow_vec);
|
|
|
|
state_idx = obj.key_table(key_to);
|
|
if state_idx==0
|
|
state_idx = true_from_state_idx;
|
|
end
|
|
obj.true_to_state_idx(symbol) = state_idx;
|
|
else
|
|
obj.true_to_state_idx(symbol) = true_from_state_idx;
|
|
end
|
|
end
|
|
true_to_state_idx = obj.true_to_state_idx(symbol);
|
|
|
|
% --- fast Dirac creation
|
|
dirac = zeros(obj.nFeasible,1,'like',x); % inherit type (gpu or cpu)
|
|
trans_idx = obj.trans_index(true_to_state_idx,true_from_state_idx);
|
|
if trans_idx~=0
|
|
dirac(trans_idx)=1;
|
|
end
|
|
|
|
% --- ensure valid (from,to)
|
|
if ~any(dirac)
|
|
mask = obj.valid_from_idx==true_from_state_idx & ...
|
|
obj.valid_to_idx ==true_to_state_idx;
|
|
if any(mask)
|
|
dirac(mask) = 1;
|
|
else
|
|
idx = find(obj.valid_from_idx==true_from_state_idx,1,'first');
|
|
dirac(idx) = 1;
|
|
obj.true_to_state_idx(symbol) = obj.valid_to_idx(idx);
|
|
end
|
|
end
|
|
|
|
% ===================================================================
|
|
% TRAINING MODE (weight update)
|
|
% ===================================================================
|
|
if training
|
|
% --- Softmax and CE
|
|
v_shift = -(v_tilde - min(v_tilde));
|
|
v_shift = min(v_shift,100);
|
|
expv = exp(v_shift);
|
|
p = expv./(sum(expv)+eps);
|
|
CE_symbol(symbol) = -log(p(dirac==1)+eps);
|
|
|
|
% --- CE smoothing and adaptive μ
|
|
if sym_idx>obj.L
|
|
CE_smooth(symbol)=0.01*CE_symbol(symbol)+0.99*CE_symbol(symbol-1);
|
|
else
|
|
CE_smooth(symbol)=CE_symbol(symbol);
|
|
end
|
|
CE_accum=CE_accum+CE_symbol(symbol);
|
|
|
|
% --- Gradient update
|
|
dmp=(dirac-p)';
|
|
dL_Dw=(yk).*dmp;
|
|
if sym_idx>=obj.L
|
|
if obj.adaptive_mu
|
|
mu_eff=CE_smooth(symbol);
|
|
mu_eff=max(min(mu_eff,0.2),1e-4);
|
|
else
|
|
mu_eff=mu;
|
|
end
|
|
obj.w=obj.w - mu_eff.*dL_Dw;
|
|
end
|
|
end
|
|
|
|
% ===================================================================
|
|
% DECODING MODE (Viterbi only)
|
|
% ===================================================================
|
|
% Compare-Select (always executed)
|
|
vmat=inf(obj.nStates,obj.nStates);
|
|
vmat(obj.valid)=v_tilde;
|
|
[pm_next,pred(symbol,:)]=min(vmat,[],2);
|
|
pm_next=pm_next-min(pm_next);
|
|
pm=pm_next;
|
|
pm_sto(:,symbol)=pm;
|
|
end
|
|
end
|
|
|
|
% --- Traceback
|
|
[~,s_end]=min(pm);
|
|
vpath=zeros(symbol,1,'uint32');
|
|
vpath(symbol)=s_end;
|
|
for n=symbol:-1:2
|
|
vpath(n-1)=pred(n,vpath(n));
|
|
end
|
|
|
|
y_ref=d(start_symbol:end);
|
|
y=obj.first_sym(vpath);
|
|
|
|
% --- BER/CE reporting and plots
|
|
if training
|
|
err=sum(y~=y_ref(1:length(y)));
|
|
ser=err/length(y);
|
|
try
|
|
ref_bits=PAMmapper(obj.S,0).demap(y_ref(1:length(y)));
|
|
eq_bits=PAMmapper(obj.S,0).demap(y);
|
|
[~,~,ber,~]=calc_ber(ref_bits,eq_bits,"skip_front",10,"skip_end",10,"returnErrorLocation",1);
|
|
fprintf('Epoch %d - BER: %.2e\n',epoch,ber);
|
|
obj.ber(epoch)=ber;
|
|
catch
|
|
fprintf('Epoch %d - SER: %.2e\n',epoch,ser);
|
|
obj.ber(epoch)=ser;
|
|
end
|
|
obj.ce(epoch)=CE_accum/symbol;
|
|
|
|
if debug && mod(epoch,10)==1 && showPlots
|
|
figure(10);clf
|
|
subplot(3,2,1:2);
|
|
if obj.use_gpu, wm=gather(obj.w); else, wm=obj.w; end
|
|
imagesc(wm);axis xy;colorbar;title('Filter W');
|
|
subplot(3,2,3);
|
|
vtilde_mat=NaN(obj.nStates,obj.nStates);
|
|
vtilde_mat(obj.valid)=gather(v_tilde); % gather just in case
|
|
imagesc(vtilde_mat);axis xy;colorbar;title('Path Metrics (v\_tilde)');
|
|
subplot(3,2,4);
|
|
plot(1:symbol,gather(pm_sto));title('Path Metric Evolution');
|
|
subplot(3,2,5);hold on;
|
|
scatter(1:symbol,CE_symbol,1,'.');
|
|
scatter(1:symbol,CE_smooth,1,'.');
|
|
title('Cross Entropy');
|
|
subplot(3,2,6);hold on;
|
|
yyaxis left
|
|
scatter(1:length(obj.ce),obj.ce,10,'s','filled');
|
|
ylabel('Cross Entropy');
|
|
yyaxis right
|
|
scatter(1:length(obj.ber),obj.ber,10,'d','filled');
|
|
set(gca,'YScale','log');
|
|
ylabel('BER (log)');
|
|
xlabel('Epoch');grid on;
|
|
title('Convergence');
|
|
drawnow;
|
|
end
|
|
end
|
|
end
|
|
end
|
|
|
|
% ==============================================================
|
|
% Helper: Sequence → key (always scalar)
|
|
% ==============================================================
|
|
function key = seq2key(obj, seq)
|
|
[~, idx] = ismember(flip(seq), obj.constellation);
|
|
pow = (obj.nSym .^ (0:obj.L-1)).';
|
|
key = 1 + sum((idx(:) - 1) .* pow);
|
|
end
|
|
end
|
|
end
|