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