From 321156578f08ae6478185e278a536ddc06492e49 Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Sat, 13 Jun 2026 16:54:26 +0200 Subject: [PATCH 01/19] Add CompressionType enum --- .../runtime/compress/CompressionType.java | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) create mode 100644 src/main/java/org/apache/sysds/runtime/compress/CompressionType.java diff --git a/src/main/java/org/apache/sysds/runtime/compress/CompressionType.java b/src/main/java/org/apache/sysds/runtime/compress/CompressionType.java new file mode 100644 index 00000000000..ba6a7b9851d --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/compress/CompressionType.java @@ -0,0 +1,43 @@ +package org.apache.sysds.runtime.compress; + +/** + * Enumeration of supported compression techniques for federated learning. + * Used for configuration, serialization, and technique selection. + * + * @author Nirvan C. Udaysingh Jhurree + */ +public enum CompressionType { + + /** TopK sparsification: keep largest-magnitude elements only */ + TOPK("topk", "Top-K Sparsification"), + + /** Probabilistic quantization: reduce precision with stochastic rounding */ + PROBABILISTIC_QUANTIZATION("prob_quant", "Probabilistic Quantization"), + + /** 1-bit compressed sensing: sign-only transmission + iterative reconstruction */ + ONE_BIT_CS("1bit_cs", "1-Bit Compressed Sensing"), + + /** No compression (passthrough) */ + NONE("none", "No Compression"); + + private final String id; + private final String description; + + CompressionType(String id, String description) { + this.id = id; + this.description = description; + } + + public String getId() { return id; } + public String getDescription() { return description; } + + /** Parse from string identifier (case-insensitive) */ + public static CompressionType fromString(String text) { + for (CompressionType type : CompressionType.values()) { + if (type.id.equalsIgnoreCase(text)) { + return type; + } + } + throw new IllegalArgumentException("Unknown compression type: " + text); + } +} \ No newline at end of file From 36d7263cc9ab7bbd6f3a024f1313ad2d8d3cb533 Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Sat, 13 Jun 2026 16:58:36 +0200 Subject: [PATCH 02/19] Add CompressionException and DecompressionException --- .../exceptions/CompressionException.java | 19 +++++++++++++++++++ .../exceptions/DecompressionException.java | 19 +++++++++++++++++++ 2 files changed, 38 insertions(+) create mode 100644 src/main/java/org/apache/sysds/runtime/compress/exceptions/CompressionException.java create mode 100644 src/main/java/org/apache/sysds/runtime/compress/exceptions/DecompressionException.java diff --git a/src/main/java/org/apache/sysds/runtime/compress/exceptions/CompressionException.java b/src/main/java/org/apache/sysds/runtime/compress/exceptions/CompressionException.java new file mode 100644 index 00000000000..da62873aa7d --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/compress/exceptions/CompressionException.java @@ -0,0 +1,19 @@ +package org.apache.sysds.runtime.compress.exceptions; + +/** + * Exception thrown when matrix compression fails. + * + * @author Nirvan C. UdaysinghJhurree + */ +public class CompressionException extends Exception { + + private static final long serialVersionUID = 1L; + + public CompressionException(String message) { + super(message); + } + + public CompressionException(String message, Throwable cause) { + super(message, cause); + } +} \ No newline at end of file diff --git a/src/main/java/org/apache/sysds/runtime/compress/exceptions/DecompressionException.java b/src/main/java/org/apache/sysds/runtime/compress/exceptions/DecompressionException.java new file mode 100644 index 00000000000..ff060588200 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/compress/exceptions/DecompressionException.java @@ -0,0 +1,19 @@ +package org.apache.sysds.runtime.compress.exceptions; + +/** + * Exception thrown when matrix decompression fails. + * + * @author Nirvan C. Udaysingh Jhurree + */ +public class DecompressionException extends Exception { + + private static final long serialVersionUID = 1L; + + public DecompressionException(String message) { + super(message); + } + + public DecompressionException(String message, Throwable cause) { + super(message, cause); + } +} \ No newline at end of file From 23d7022e0f785b6ed8969556f73f3e5d6e1022cf Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Sat, 13 Jun 2026 17:02:07 +0200 Subject: [PATCH 03/19] Add CompressedMatrix container class --- .../runtime/compress/CompressedMatrix.java | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 src/main/java/org/apache/sysds/runtime/compress/CompressedMatrix.java diff --git a/src/main/java/org/apache/sysds/runtime/compress/CompressedMatrix.java b/src/main/java/org/apache/sysds/runtime/compress/CompressedMatrix.java new file mode 100644 index 00000000000..a24b5255d1c --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/compress/CompressedMatrix.java @@ -0,0 +1,64 @@ +package org.apache.sysds.runtime.compress; + +import java.io.Serializable; + +/** + * Generic container for compressed matrix data. + * Stores the compressed representation along with metadata + * needed for decompression and size estimation. + * + * @author Nirvan C. Udaysingh Jhurree + */ +public class CompressedMatrix implements Serializable { + + private static final long serialVersionUID = 1L; + + private final CompressionType type; + private final int numRows; + private final int numCols; + private final Object compressedData; // Technique-specific data + private final double compressionRatio; + private final byte[] metadata; // Optional: scaling factors, etc. + + public CompressedMatrix(CompressionType type, int numRows, int numCols, + Object compressedData, double compressionRatio) { + this(type, numRows, numCols, compressedData, compressionRatio, null); + } + + public CompressedMatrix(CompressionType type, int numRows, int numCols, + Object compressedData, double compressionRatio, + byte[] metadata) { + this.type = type; + this.numRows = numRows; + this.numCols = numCols; + this.compressedData = compressedData; + this.compressionRatio = compressionRatio; + this.metadata = metadata; + } + + public CompressionType getType() { return type; } + public int getNumRows() { return numRows; } + public int getNumCols() { return numCols; } + public Object getCompressedData() { return compressedData; } + public double getCompressionRatio() { return compressionRatio; } + public byte[] getMetadata() { return metadata; } + + /** Estimate original size in bytes (8 bytes per double) */ + public long estimateOriginalSizeBytes() { + return (long) numRows * numCols * 8; + } + + /** Estimate compressed size in bytes */ + public long getCompressedSizeBytes() { + if (compressedData instanceof byte[]) { + return ((byte[]) compressedData).length; + } + return 0; + } + + @Override + public String toString() { + return String.format("CompressedMatrix[%s, %dx%d, ratio=%.2fx]", + type.getId(), numRows, numCols, compressionRatio); + } +} \ No newline at end of file From ec8f5df0149073642dbec57734a6158f94949d7a Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Sat, 13 Jun 2026 17:06:12 +0200 Subject: [PATCH 04/19] Add MatrixCompressor interface --- .../runtime/compress/MatrixCompressor.java | 44 +++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 src/main/java/org/apache/sysds/runtime/compress/MatrixCompressor.java diff --git a/src/main/java/org/apache/sysds/runtime/compress/MatrixCompressor.java b/src/main/java/org/apache/sysds/runtime/compress/MatrixCompressor.java new file mode 100644 index 00000000000..88d235e9d2f --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/compress/MatrixCompressor.java @@ -0,0 +1,44 @@ +package org.apache.sysds.runtime.compress; + +import org.apache.sysds.runtime.compress.exceptions.CompressionException; +import org.apache.sysds.runtime.compress.exceptions.DecompressionException; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; + +/** + * Interface for matrix compression techniques in federated learning. + * All compressors must implement compress/decompress operations. + * + * @author Nirvan C. Udaysingh Jhurree + */ +public interface MatrixCompressor { + + /** + * Compress a matrix block for transmission. + * @param input The source matrix to compress + * @return CompressedMatrix containing compressed data and metadata + * @throws CompressionException if compression fails + */ + CompressedMatrix compress(MatrixBlock input) throws CompressionException; + + /** + * Decompress a compressed matrix back to MatrixBlock. + * @param compressed The compressed data to decompress + * @return Reconstructed MatrixBlock (may be approximate) + * @throws DecompressionException if decompression fails + */ + MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException; + + /** + * Get the compression technique identifier. + * @return CompressionType enum value + */ + CompressionType getCompressionType(); + + /** + * Estimate the compression ratio achieved. + * Higher is better (e.g. 10.0 means 10x smaller). + */ + default double estimateCompressionRatio(long originalSize, long compressedSize) { + return compressedSize == 0 ? Double.MAX_VALUE : (double) originalSize / compressedSize; + } +} \ No newline at end of file From 58e592dfa837c3e84cf4e207c6a54d93c0e0078d Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Sat, 13 Jun 2026 17:09:01 +0200 Subject: [PATCH 05/19] Add TopKData holder class --- .../sysds/runtime/compress/TopK/TopKData.java | 44 +++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 src/main/java/org/apache/sysds/runtime/compress/TopK/TopKData.java diff --git a/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKData.java b/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKData.java new file mode 100644 index 00000000000..509653f7cd0 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKData.java @@ -0,0 +1,44 @@ +package org.apache.sysds.runtime.compress.TopK; + +import java.io.Serializable; + +/** + * Immutable container for TopK-compressed matrix data. + * Stores only the K largest-magnitude elements with their positions, + * designed for efficient serialization across federated workers. + * + * @author Nirvan C. Udaysingh Jhurree + */ +public class TopKData implements Serializable { + + private static final long serialVersionUID = 1L; + + public final int[] indices; // Linear indices of kept elements (row*numCols + col) + public final double[] values; // Corresponding original values + public final int numCols; // Needed for index → (row, col) conversion + + public TopKData(int[] indices, double[] values, int numCols) { + if (indices.length != values.length) { + throw new IllegalArgumentException( + "Indices and values arrays must have the same length"); + } + this.indices = indices.clone(); // Defensive copy + this.values = values.clone(); + this.numCols = numCols; + } + + /** Number of kept elements */ + public int size() { + return indices.length; + } + + /** Estimate serialized size in bytes (4 bytes per int + 8 bytes per double) */ + public long estimateSizeBytes() { + return (long) indices.length * 12 + 64; // +64 for object headers + } + + @Override + public String toString() { + return String.format("TopKData[k=%d, numCols=%d]", indices.length, numCols); + } +} \ No newline at end of file From 7ef87da09adf5d40c87415fd97cb586ae65281b4 Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Sat, 13 Jun 2026 17:11:13 +0200 Subject: [PATCH 06/19] Add TopKCompressor implementation --- .../runtime/compress/TopK/TopKCompressor.java | 225 ++++++++++++++++++ 1 file changed, 225 insertions(+) create mode 100644 src/main/java/org/apache/sysds/runtime/compress/TopK/TopKCompressor.java diff --git a/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKCompressor.java b/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKCompressor.java new file mode 100644 index 00000000000..d178d6cdc26 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKCompressor.java @@ -0,0 +1,225 @@ +package org.apache.sysds.runtime.compress.TopK; + +import org.apache.sysds.runtime.compress.CompressedMatrix; +import org.apache.sysds.runtime.compress.CompressionType; +import org.apache.sysds.runtime.compress.MatrixCompressor; +import org.apache.sysds.runtime.compress.exceptions.CompressionException; +import org.apache.sysds.runtime.compress.exceptions.DecompressionException; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.PriorityQueue; + +/** + * TopK Sparsification Compressor. + * + * Keeps only the K largest-magnitude elements in the matrix, + * setting all others to zero. Optimal for gradient sparsification + * in federated learning where most gradient values are near zero. + * + * Compression ratio: approximately 1/sparsityRatio + * e.g. sparsityRatio=0.01 keeps 1% of elements → ~100x compression + * + * @author Nirvan Jhurree + */ +public class TopKCompressor implements MatrixCompressor { + + private final double sparsityRatio; // Fraction of elements to keep (0, 1] + private final boolean useHeap; // Use min-heap for O(n log k) selection + + /** + * @param sparsityRatio Fraction of elements to retain e.g. 0.01 = keep top 1% + * @param useHeap If true, use priority queue (faster for large matrices) + */ + public TopKCompressor(double sparsityRatio, boolean useHeap) { + if (sparsityRatio <= 0 || sparsityRatio > 1) { + throw new IllegalArgumentException("sparsityRatio must be in (0, 1]"); + } + this.sparsityRatio = sparsityRatio; + this.useHeap = useHeap; + } + + /** Default constructor: uses heap-based selection */ + public TopKCompressor(double sparsityRatio) { + this(sparsityRatio, true); + } + + @Override + public CompressedMatrix compress(MatrixBlock input) throws CompressionException { + try { + int numRows = input.getNumRows(); + int numCols = input.getNumColumns(); + int totalElements = numRows * numCols; + int k = (int) Math.max(1, Math.ceil(totalElements * sparsityRatio)); + + // If k covers everything, no compression needed + if (k >= totalElements) { + return new CompressedMatrix( + CompressionType.TOPK, numRows, numCols, input, 1.0); + } + + // Extract all non-zero elements with their linear indices + List elements = extractElements(input, numRows, numCols); + + // If fewer non-zeros than k, keep all of them + List topK = (elements.size() <= k) + ? new ArrayList<>(elements) + : selectTopK(elements, k); + + // Pack into TopKData + TopKData data = convertToTopKData(topK, numCols); + + double ratio = calculateCompressionRatio(totalElements, topK.size()); + + return new CompressedMatrix( + CompressionType.TOPK, numRows, numCols, data, ratio); + + } catch (Exception e) { + throw new CompressionException("TopK compression failed: " + e.getMessage(), e); + } + } + + @Override + public MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException { + try { + // Handle passthrough case (no compression was applied) + if (compressed.getCompressedData() instanceof MatrixBlock) { + return (MatrixBlock) compressed.getCompressedData(); + } + + TopKData data = (TopKData) compressed.getCompressedData(); + MatrixBlock result = new MatrixBlock( + compressed.getNumRows(), + compressed.getNumCols(), + true // Start sparse + ); + result.allocateSparseRowsBlock(); + + // Place values back at their original positions + for (int i = 0; i < data.indices.length; i++) { + int linearIdx = data.indices[i]; + int row = linearIdx / data.numCols; + int col = linearIdx % data.numCols; + result.setValue(row, col, data.values[i]); + } + + result.examSparsity(); + return result; + + } catch (ClassCastException e) { + throw new DecompressionException("Invalid compressed data type for TopK", e); + } catch (Exception e) { + throw new DecompressionException("TopK decompression failed: " + e.getMessage(), e); + } + } + + @Override + public CompressionType getCompressionType() { + return CompressionType.TOPK; + } + + // ----------------------------------------------------------------------- + // Private helpers + // ----------------------------------------------------------------------- + + /** + * Extract all non-zero elements with their linear indices. + * Handles both dense and sparse MatrixBlock representations. + */ + private List extractElements(MatrixBlock input, int numRows, int numCols) { + List elements = new ArrayList<>(); + + if (input.isInSparseFormat()) { + // Sparse: iterate only over non-zero entries + for (int i = 0; i < numRows; i++) { + if (input.getSparseBlock() == null) continue; + if (input.getSparseBlock().isEmpty(i)) continue; + int[] rowIndices = input.getSparseBlock().indexes(i); + double[] rowValues = input.getSparseBlock().values(i); + int rowSize = input.getSparseBlock().size(i); + for (int j = 0; j < rowSize; j++) { + double val = rowValues[j]; + if (val != 0.0) { + int linearIdx = i * numCols + rowIndices[j]; + elements.add(new Element(linearIdx, val, Math.abs(val))); + } + } + } + } else { + // Dense: iterate all elements, skip zeros + double[] denseBlock = input.getDenseBlockValues(); + if (denseBlock != null) { + for (int i = 0; i < denseBlock.length; i++) { + if (denseBlock[i] != 0.0) { + elements.add(new Element(i, denseBlock[i], Math.abs(denseBlock[i]))); + } + } + } + } + return elements; + } + + /** + * Select top K elements by absolute value. + * Uses min-heap for O(n log k) when useHeap=true, + * or full sort O(n log n) otherwise. + */ + private List selectTopK(List elements, int k) { + if (useHeap) { + PriorityQueue minHeap = new PriorityQueue<>( + k, Comparator.comparingDouble(e -> e.absValue) + ); + for (Element e : elements) { + if (minHeap.size() < k) { + minHeap.offer(e); + } else if (e.absValue > minHeap.peek().absValue) { + minHeap.poll(); + minHeap.offer(e); + } + } + List result = new ArrayList<>(minHeap); + result.sort(Comparator.comparingInt(e -> e.index)); + return result; + } else { + elements.sort((a, b) -> Double.compare(b.absValue, a.absValue)); + return new ArrayList<>(elements.subList(0, k)); + } + } + + private TopKData convertToTopKData(List topK, int numCols) { + int[] indices = new int[topK.size()]; + double[] values = new double[topK.size()]; + for (int i = 0; i < topK.size(); i++) { + indices[i] = topK.get(i).index; + values[i] = topK.get(i).value; + } + return new TopKData(indices, values, numCols); + } + + private double calculateCompressionRatio(int total, int kept) { + if (kept == 0) return Double.MAX_VALUE; + // Original: total * 8 bytes (doubles) + // Compressed: kept * 12 bytes (int index + double value) + long originalBytes = (long) total * 8; + long compressedBytes = (long) kept * 12; + return (double) originalBytes / compressedBytes; + } + + // ----------------------------------------------------------------------- + // Inner class: element tracking during compression + // ----------------------------------------------------------------------- + + private static class Element { + final int index; // Linear index: row * numCols + col + final double value; // Original value + final double absValue; // Absolute value for magnitude comparison + + Element(int index, double value, double absValue) { + this.index = index; + this.value = value; + this.absValue = absValue; + } + } +} \ No newline at end of file From ec3f6be8de8999a5bac08bc6d355af6f8354660b Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Sat, 13 Jun 2026 17:17:27 +0200 Subject: [PATCH 07/19] Fix MatrixBlock API: setValue -> set in TopKCompressor --- .../org/apache/sysds/runtime/compress/TopK/TopKCompressor.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKCompressor.java b/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKCompressor.java index d178d6cdc26..a8de26db3a6 100644 --- a/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKCompressor.java +++ b/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKCompressor.java @@ -102,7 +102,8 @@ public MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionE int linearIdx = data.indices[i]; int row = linearIdx / data.numCols; int col = linearIdx % data.numCols; - result.setValue(row, col, data.values[i]); + //result.setValue(row, col, data.values[i]); + result.set(row, col, data.values[i]); } result.examSparsity(); From b58b601092c842f04570cc7d05f369082eaab949 Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Mon, 29 Jun 2026 12:43:57 +0200 Subject: [PATCH 08/19] Add QuantizedData holder class --- .../compress/Quantization/QuantizedData.java | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 src/main/java/org/apache/sysds/runtime/compress/Quantization/QuantizedData.java diff --git a/src/main/java/org/apache/sysds/runtime/compress/Quantization/QuantizedData.java b/src/main/java/org/apache/sysds/runtime/compress/Quantization/QuantizedData.java new file mode 100644 index 00000000000..632066e9854 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/compress/Quantization/QuantizedData.java @@ -0,0 +1,57 @@ +package org.apache.sysds.runtime.compress.Quantization; + +import java.io.Serializable; + +/** + * Immutable container for probabilistically quantized matrix data. + * Stores quantized byte indices and the scaling parameters needed + * to reconstruct approximate original values on decompression. + * + * + */ +public class QuantizedData implements Serializable { + + private static final long serialVersionUID = 1L; + + public final byte[] quantizedValues; // Quantized level indices + public final double min; // Original minimum value + public final double max; // Original maximum value + public final int levels; // Number of quantization levels (2^bits) + public final int bitsPerValue; // Bits used per element + public final int numRows; + public final int numCols; + + public QuantizedData(byte[] quantizedValues, double min, double max, + int levels, int bitsPerValue, int numRows, int numCols) { + this.quantizedValues = quantizedValues.clone(); // Defensive copy + this.min = min; + this.max = max; + this.levels = levels; + this.bitsPerValue = bitsPerValue; + this.numRows = numRows; + this.numCols = numCols; + } + + /** Number of quantized elements */ + public int size() { + return quantizedValues.length; + } + + /** Estimate serialized size in bytes */ + public long estimateSizeBytes() { + return quantizedValues.length + 64; // +64 for scalar fields and headers + } + + /** Reconstruct a double value from a quantized level index */ + public double reconstructValue(byte levelIndex) { + if (max - min < 1e-10) return min; // Constant matrix + int idx = levelIndex & 0xFF; // Treat byte as unsigned + return min + (idx / (double)(levels - 1)) * (max - min); + } + + @Override + public String toString() { + return String.format("QuantizedData[%dx%d, levels=%d, bits=%d, min=%.4f, max=%.4f]", + numRows, numCols, levels, bitsPerValue, min, max); + } +} \ No newline at end of file From 136ff91645040d57fa8c747b0a7854b5e7483dfd Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Mon, 29 Jun 2026 12:58:56 +0200 Subject: [PATCH 09/19] Add ProbabilisticQuantizationCompressor implementation --- .../ProbabilisticQuantizationCompressor.java | 158 ++++++++++++++++++ 1 file changed, 158 insertions(+) create mode 100644 src/main/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressor.java diff --git a/src/main/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressor.java b/src/main/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressor.java new file mode 100644 index 00000000000..95ce2562b09 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressor.java @@ -0,0 +1,158 @@ +package org.apache.sysds.runtime.compress.Quantization; + +import org.apache.sysds.runtime.compress.CompressedMatrix; +import org.apache.sysds.runtime.compress.CompressionType; +import org.apache.sysds.runtime.compress.MatrixCompressor; +import org.apache.sysds.runtime.compress.exceptions.CompressionException; +import org.apache.sysds.runtime.compress.exceptions.DecompressionException; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; + +import java.util.Random; + +/** + * Probabilistic Quantization Compressor. + * + * Reduces numerical precision using stochastic rounding to maintain + * an unbiased estimator — meaning E[quantized] = original on average. + * This is critical for federated learning convergence guarantees. + * + * Supports 2, 4, or 8 bits per value: + * 2-bit → 4 levels → 16x compression vs 32-bit float + * 4-bit → 16 levels → 8x compression + * 8-bit → 256 levels → 4x compression + * + * + */ +public class ProbabilisticQuantizationCompressor implements MatrixCompressor { + + private final int bitsPerValue; // 2, 4, or 8 + private final Random rng; + + public ProbabilisticQuantizationCompressor(int bitsPerValue) { + if (bitsPerValue != 2 && bitsPerValue != 4 && bitsPerValue != 8) { + throw new IllegalArgumentException("bitsPerValue must be 2, 4, or 8"); + } + this.bitsPerValue = bitsPerValue; + this.rng = new Random(42); // Fixed seed for reproducibility + } + + @Override + public CompressedMatrix compress(MatrixBlock input) throws CompressionException { + try { + int numRows = input.getNumRows(); + int numCols = input.getNumColumns(); + int totalElements = numRows * numCols; + + // Find min and max for normalization + double[] minMax = findMinMax(input, numRows, numCols); + double min = minMax[0]; + double max = minMax[1]; + + int levels = 1 << bitsPerValue; // 2^bits + + // Quantize each element probabilistically + byte[] quantized = new byte[totalElements]; + for (int i = 0; i < numRows; i++) { + for (int j = 0; j < numCols; j++) { + double value = input.get(i, j); + quantized[i * numCols + j] = probabilisticRound(value, min, max, levels); + } + } + + double ratio = 32.0 / bitsPerValue; // vs 32-bit float + + QuantizedData data = new QuantizedData( + quantized, min, max, levels, bitsPerValue, numRows, numCols); + + return new CompressedMatrix( + CompressionType.PROBABILISTIC_QUANTIZATION, + numRows, numCols, data, ratio); + + } catch (Exception e) { + throw new CompressionException( + "Quantization compression failed: " + e.getMessage(), e); + } + } + + @Override + public MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException { + try { + QuantizedData data = (QuantizedData) compressed.getCompressedData(); + MatrixBlock result = new MatrixBlock(data.numRows, data.numCols, false); + result.allocateDenseBlock(); + + for (int i = 0; i < data.numRows; i++) { + for (int j = 0; j < data.numCols; j++) { + byte levelIndex = data.quantizedValues[i * data.numCols + j]; + double value = data.reconstructValue(levelIndex); + result.set(i, j, value); + } + } + + result.examSparsity(); + return result; + + } catch (ClassCastException e) { + throw new DecompressionException( + "Invalid compressed data type for Quantization", e); + } catch (Exception e) { + throw new DecompressionException( + "Quantization decompression failed: " + e.getMessage(), e); + } + } + + @Override + public CompressionType getCompressionType() { + return CompressionType.PROBABILISTIC_QUANTIZATION; + } + + // ----------------------------------------------------------------------- + // Private helpers + // ----------------------------------------------------------------------- + + /** + * Stochastic rounding: for value x between levels q_i and q_{i+1}: + * P(round up) = (x - q_i) / (q_{i+1} - q_i) + * P(round down) = 1 - P(round up) + * This gives E[output] = x (unbiased). + */ + private byte probabilisticRound(double value, double min, double max, int levels) { + // Handle constant matrix edge case + if (max - min < 1e-10) return 0; + + // Normalize to [0, 1] + double normalized = (value - min) / (max - min); + normalized = Math.max(0.0, Math.min(1.0, normalized)); // Clamp + + // Find bounding level indices + double scaled = normalized * (levels - 1); + int lowerIdx = (int) scaled; + int upperIdx = Math.min(lowerIdx + 1, levels - 1); + + if (lowerIdx == upperIdx) { + return (byte) lowerIdx; + } + + // Probabilistic decision + double probUp = scaled - lowerIdx; + return (rng.nextDouble() < probUp) ? (byte) upperIdx : (byte) lowerIdx; + } + + /** Find min and max values across the entire matrix */ + private double[] findMinMax(MatrixBlock input, int numRows, int numCols) { + double min = Double.MAX_VALUE; + double max = -Double.MAX_VALUE; + + for (int i = 0; i < numRows; i++) { + for (int j = 0; j < numCols; j++) { + double val = input.get(i, j); + if (val < min) min = val; + if (val > max) max = val; + } + } + + // Handle all-zero matrix + if (min == Double.MAX_VALUE) { min = 0; max = 0; } + return new double[]{min, max}; + } +} \ No newline at end of file From 4087cbe547938858158c7ad3318f4a0e38470637 Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Mon, 29 Jun 2026 13:03:11 +0200 Subject: [PATCH 10/19] Add CompressionConfig builder --- .../runtime/compress/CompressionConfig.java | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) create mode 100644 src/main/java/org/apache/sysds/runtime/compress/CompressionConfig.java diff --git a/src/main/java/org/apache/sysds/runtime/compress/CompressionConfig.java b/src/main/java/org/apache/sysds/runtime/compress/CompressionConfig.java new file mode 100644 index 00000000000..d2634bda528 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/compress/CompressionConfig.java @@ -0,0 +1,93 @@ +package org.apache.sysds.runtime.compress; + +import java.util.HashMap; +import java.util.Map; + +/** + * Immutable configuration for compression in federated operations. + * Uses the Builder pattern for flexible, readable configuration. + * + * Usage example: + * CompressionConfig config = CompressionConfig.builder() + * .enable(true) + * .withType(CompressionType.TOPK) + * .withSparsity(0.01) + * .build(); + * + * + */ +public class CompressionConfig { + + private final boolean enabled; + private final CompressionType type; + private final Map parameters; + + private CompressionConfig(Builder builder) { + this.enabled = builder.enabled; + this.type = builder.enabled ? builder.type : CompressionType.NONE; + this.parameters = new HashMap<>(builder.parameters); + } + + public boolean isEnabled() { return enabled; } + public CompressionType getType() { return type; } + public Map getParameters() { return new HashMap<>(parameters); } + + /** Convenience getter for sparsity parameter (TopK) */ + public double getSparsity() { + return (double) parameters.getOrDefault("sparsity", 0.01); + } + + /** Convenience getter for bits parameter (Quantization) */ + public int getBits() { + return (int) parameters.getOrDefault("bits", 4); + } + + @Override + public String toString() { + return String.format("CompressionConfig[enabled=%s, type=%s, params=%s]", + enabled, type.getId(), parameters); + } + + // ----------------------------------------------------------------------- + // Builder + // ----------------------------------------------------------------------- + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + private boolean enabled = false; + private CompressionType type = CompressionType.NONE; + private final Map parameters = new HashMap<>(); + + public Builder enable(boolean enabled) { + this.enabled = enabled; + return this; + } + + public Builder withType(CompressionType type) { + this.type = type; + return this; + } + + public Builder withParameter(String key, Object value) { + this.parameters.put(key, value); + return this; + } + + /** Shorthand for TopK sparsity ratio */ + public Builder withSparsity(double sparsity) { + return withParameter("sparsity", sparsity); + } + + /** Shorthand for quantization bit width */ + public Builder withBits(int bits) { + return withParameter("bits", bits); + } + + public CompressionConfig build() { + return new CompressionConfig(this); + } + } +} \ No newline at end of file From 65e35b34b5ade44e667211f3c1ff6ed1da946107 Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Mon, 29 Jun 2026 13:07:44 +0200 Subject: [PATCH 11/19] Add CompressionFactory with PassthroughCompressor --- .../runtime/compress/CompressionFactory.java | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) create mode 100644 src/main/java/org/apache/sysds/runtime/compress/CompressionFactory.java diff --git a/src/main/java/org/apache/sysds/runtime/compress/CompressionFactory.java b/src/main/java/org/apache/sysds/runtime/compress/CompressionFactory.java new file mode 100644 index 00000000000..04c7a982fc2 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/compress/CompressionFactory.java @@ -0,0 +1,95 @@ +package org.apache.sysds.runtime.compress; + +import org.apache.sysds.runtime.compress.TopK.TopKCompressor; +import org.apache.sysds.runtime.compress.Quantization.ProbabilisticQuantizationCompressor; + +/** + * Factory for creating compressor instances from configuration. + * Centralizes compressor instantiation and parameter validation. + * + * Usage: + * CompressionConfig config = CompressionConfig.builder() + * .enable(true) + * .withType(CompressionType.TOPK) + * .withSparsity(0.01) + * .build(); + * MatrixCompressor compressor = CompressionFactory.create(config); + * + * + */ +public class CompressionFactory { + + private CompressionFactory() { + // Utility class — no instantiation + } + + /** + * Create a compressor from a CompressionConfig. + * @param config The compression configuration + * @return A ready-to-use MatrixCompressor + * @throws IllegalArgumentException if the config is invalid + */ + public static MatrixCompressor create(CompressionConfig config) { + if (config == null || !config.isEnabled()) { + return new PassthroughCompressor(); + } + return create(config.getType(), config); + } + + /** + * Create a compressor for a specific type with given config. + */ + public static MatrixCompressor create(CompressionType type, CompressionConfig config) { + switch (type) { + case TOPK: + double sparsity = config.getSparsity(); + return new TopKCompressor(sparsity, true); + + case PROBABILISTIC_QUANTIZATION: + int bits = config.getBits(); + return new ProbabilisticQuantizationCompressor(bits); + + case ONE_BIT_CS: + throw new UnsupportedOperationException( + "1-Bit Compressed Sensing not yet implemented"); + + case NONE: + default: + return new PassthroughCompressor(); + } + } + + // ----------------------------------------------------------------------- + // Passthrough compressor (no-op) for when compression is disabled + // ----------------------------------------------------------------------- + + /** + * No-op compressor: returns the matrix as-is. + * Used when compression is disabled or type is NONE. + */ + private static class PassthroughCompressor implements MatrixCompressor { + + @Override + public CompressedMatrix compress(org.apache.sysds.runtime.matrix.data.MatrixBlock input) + throws org.apache.sysds.runtime.compress.exceptions.CompressionException { + return new CompressedMatrix( + CompressionType.NONE, + input.getNumRows(), + input.getNumColumns(), + input, + 1.0 + ); + } + + @Override + public org.apache.sysds.runtime.matrix.data.MatrixBlock decompress(CompressedMatrix compressed) + throws org.apache.sysds.runtime.compress.exceptions.DecompressionException { + return (org.apache.sysds.runtime.matrix.data.MatrixBlock) compressed.getCompressedData(); + } + + @Override + public CompressionType getCompressionType() { + return CompressionType.NONE; + } + } +} \ No newline at end of file From f01d81f5d027033e6428072922a5501b2099b57e Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Mon, 29 Jun 2026 13:19:49 +0200 Subject: [PATCH 12/19] Add TopKCompressorTest - 9 tests passing --- .../compress/TopK/TopKCompressorTest.java | 157 ++++++++++++++++++ 1 file changed, 157 insertions(+) create mode 100644 src/test/java/org/apache/sysds/runtime/compress/TopK/TopKCompressorTest.java diff --git a/src/test/java/org/apache/sysds/runtime/compress/TopK/TopKCompressorTest.java b/src/test/java/org/apache/sysds/runtime/compress/TopK/TopKCompressorTest.java new file mode 100644 index 00000000000..3fd1fd426f2 --- /dev/null +++ b/src/test/java/org/apache/sysds/runtime/compress/TopK/TopKCompressorTest.java @@ -0,0 +1,157 @@ +package org.apache.sysds.runtime.compress.TopK; + +import org.apache.sysds.runtime.compress.CompressedMatrix; +import org.apache.sysds.runtime.compress.CompressionType; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; +import org.junit.Test; +import static org.junit.Assert.*; + +/** + * Unit tests for TopKCompressor. + * Verifies compression ratio, reconstruction accuracy, + * and correct handling of edge cases. + * + * + */ +public class TopKCompressorTest { + + // ----------------------------------------------------------------------- + // Basic compression / decompression + // ----------------------------------------------------------------------- + + @Test + public void testTopKKeepsLargestElements() throws Exception { + // 3x3 matrix with three distinct magnitudes + MatrixBlock input = new MatrixBlock(3, 3, false); + input.allocateDenseBlock(); + input.set(0, 0, 10.0); // Largest + input.set(1, 1, 5.0); // Medium + input.set(2, 2, 1.0); // Smallest + input.examSparsity(); + + // Keep top 2 of 9 elements (~22% sparsity) + TopKCompressor compressor = new TopKCompressor(0.22); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + // Largest two values must be preserved exactly + assertEquals(10.0, result.get(0, 0), 1e-10); + assertEquals(5.0, result.get(1, 1), 1e-10); + + // Smallest should be zeroed out + assertEquals(0.0, result.get(2, 2), 1e-10); + } + + @Test + public void testCompressionTypeIsTopK() throws Exception { + MatrixBlock input = createDenseMatrix(4, 4, 1.0); + TopKCompressor compressor = new TopKCompressor(0.5); + CompressedMatrix compressed = compressor.compress(input); + assertEquals(CompressionType.TOPK, compressed.getType()); + } + + @Test + public void testDimensionsPreservedAfterDecompression() throws Exception { + MatrixBlock input = createRandomMatrix(10, 20); + TopKCompressor compressor = new TopKCompressor(0.1); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + assertEquals(10, result.getNumRows()); + assertEquals(20, result.getNumColumns()); + } + + // ----------------------------------------------------------------------- + // Compression ratio + // ----------------------------------------------------------------------- + + @Test + public void testCompressionRatioIsPositive() throws Exception { + MatrixBlock input = createRandomMatrix(50, 50); + TopKCompressor compressor = new TopKCompressor(0.01); + CompressedMatrix compressed = compressor.compress(input); + assertTrue("Compression ratio must be > 0", + compressed.getCompressionRatio() > 0); + } + + @Test + public void testLowerSparsityGivesHigherRatio() throws Exception { + MatrixBlock input = createRandomMatrix(100, 100); + + TopKCompressor c1 = new TopKCompressor(0.1); + TopKCompressor c2 = new TopKCompressor(0.01); + + double ratio1 = c1.compress(input).getCompressionRatio(); + double ratio2 = c2.compress(input).getCompressionRatio(); + + assertTrue("1% sparsity should compress more than 10%", ratio2 > ratio1); + } + + // ----------------------------------------------------------------------- + // Edge cases + // ----------------------------------------------------------------------- + + @Test + public void testAllZeroMatrix() throws Exception { + MatrixBlock input = new MatrixBlock(5, 5, false); + input.allocateDenseBlock(); + input.examSparsity(); + + TopKCompressor compressor = new TopKCompressor(0.1); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + // All zeros in → all zeros out + for (int i = 0; i < 5; i++) + for (int j = 0; j < 5; j++) + assertEquals(0.0, result.get(i, j), 1e-10); + } + + @Test + public void testSparsityOfOneKeepsEverything() throws Exception { + MatrixBlock input = createRandomMatrix(5, 5); + TopKCompressor compressor = new TopKCompressor(1.0); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + // With sparsity=1.0, all values should be preserved + for (int i = 0; i < 5; i++) + for (int j = 0; j < 5; j++) + assertEquals(input.get(i, j), result.get(i, j), 1e-10); + } + + @Test(expected = IllegalArgumentException.class) + public void testInvalidSparsityThrowsException() { + new TopKCompressor(0.0); // Must be > 0 + } + + @Test(expected = IllegalArgumentException.class) + public void testSparsityAboveOneThrowsException() { + new TopKCompressor(1.5); // Must be <= 1 + } + + // ----------------------------------------------------------------------- + // Helpers + // ----------------------------------------------------------------------- + + private MatrixBlock createRandomMatrix(int rows, int cols) { + MatrixBlock m = new MatrixBlock(rows, cols, false); + m.allocateDenseBlock(); + java.util.Random rng = new java.util.Random(42); + for (int i = 0; i < rows; i++) + for (int j = 0; j < cols; j++) + m.set(i, j, rng.nextGaussian() * 10); + m.examSparsity(); + return m; + } + + private MatrixBlock createDenseMatrix(int rows, int cols, double fillValue) { + MatrixBlock m = new MatrixBlock(rows, cols, false); + m.allocateDenseBlock(); + for (int i = 0; i < rows; i++) + for (int j = 0; j < cols; j++) + m.set(i, j, fillValue); + m.examSparsity(); + return m; + } +} \ No newline at end of file From 927fc196e5940106bf32ef33a9e09e5e7499cdfa Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Mon, 29 Jun 2026 13:26:43 +0200 Subject: [PATCH 13/19] Add ProbabilisticQuantizationCompressorTest --- ...obabilisticQuantizationCompressorTest.java | 187 ++++++++++++++++++ 1 file changed, 187 insertions(+) create mode 100644 src/test/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressorTest.java diff --git a/src/test/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressorTest.java b/src/test/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressorTest.java new file mode 100644 index 00000000000..a50ae2f3f79 --- /dev/null +++ b/src/test/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressorTest.java @@ -0,0 +1,187 @@ +package org.apache.sysds.runtime.compress.Quantization; + +import org.apache.sysds.runtime.compress.CompressedMatrix; +import org.apache.sysds.runtime.compress.CompressionType; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; +import org.junit.Test; +import static org.junit.Assert.*; + +/** + * Unit tests for ProbabilisticQuantizationCompressor. + * Verifies compression ratio, reconstruction accuracy, + * unbiasedness property, and edge case handling. + * + * + */ +public class ProbabilisticQuantizationCompressorTest { + + // ----------------------------------------------------------------------- + // Basic compression / decompression + // ----------------------------------------------------------------------- + + @Test + public void testCompressionTypeIsProbabilisticQuantization() throws Exception { + MatrixBlock input = createRandomMatrix(4, 4); + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(4); + CompressedMatrix compressed = compressor.compress(input); + assertEquals(CompressionType.PROBABILISTIC_QUANTIZATION, compressed.getType()); + } + + @Test + public void testDimensionsPreservedAfterDecompression() throws Exception { + MatrixBlock input = createRandomMatrix(10, 20); + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(4); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + assertEquals(10, result.getNumRows()); + assertEquals(20, result.getNumColumns()); + } + + @Test + public void testReconstructedValuesWithinOriginalRange() throws Exception { + MatrixBlock input = createRandomMatrix(20, 20); + + // Find original min/max + double origMin = Double.MAX_VALUE; + double origMax = -Double.MAX_VALUE; + for (int i = 0; i < 20; i++) { + for (int j = 0; j < 20; j++) { + double v = input.get(i, j); + if (v < origMin) origMin = v; + if (v > origMax) origMax = v; + } + } + + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(4); + MatrixBlock result = compressor.decompress(compressor.compress(input)); + + // All reconstructed values must be within [min, max] + for (int i = 0; i < 20; i++) { + for (int j = 0; j < 20; j++) { + double v = result.get(i, j); + assertTrue("Value below min: " + v, v >= origMin - 1e-9); + assertTrue("Value above max: " + v, v <= origMax + 1e-9); + } + } + } + + // ----------------------------------------------------------------------- + // Compression ratio + // ----------------------------------------------------------------------- + + @Test + public void testCompressionRatio2Bit() throws Exception { + MatrixBlock input = createRandomMatrix(10, 10); + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(2); + CompressedMatrix compressed = compressor.compress(input); + assertEquals(16.0, compressed.getCompressionRatio(), 1e-10); + } + + @Test + public void testCompressionRatio4Bit() throws Exception { + MatrixBlock input = createRandomMatrix(10, 10); + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(4); + CompressedMatrix compressed = compressor.compress(input); + assertEquals(8.0, compressed.getCompressionRatio(), 1e-10); + } + + @Test + public void testCompressionRatio8Bit() throws Exception { + MatrixBlock input = createRandomMatrix(10, 10); + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(8); + CompressedMatrix compressed = compressor.compress(input); + assertEquals(4.0, compressed.getCompressionRatio(), 1e-10); + } + + @Test + public void testFewerBitsGivesHigherRatio() throws Exception { + MatrixBlock input = createRandomMatrix(20, 20); + + double ratio2bit = new ProbabilisticQuantizationCompressor(2) + .compress(input).getCompressionRatio(); + double ratio8bit = new ProbabilisticQuantizationCompressor(8) + .compress(input).getCompressionRatio(); + + assertTrue("2-bit should compress more than 8-bit", ratio2bit > ratio8bit); + } + + // ----------------------------------------------------------------------- + // Unbiasedness: E[quantized] ≈ original over many runs + // ----------------------------------------------------------------------- + + @Test + public void testUnbiasednessOverManyRuns() throws Exception { + // Single element matrix with value 0.5 (midpoint) + MatrixBlock input = new MatrixBlock(1, 1, false); + input.allocateDenseBlock(); + input.set(0, 0, 5.0); + input.examSparsity(); + + // Run quantization 1000 times and average the results + double sum = 0.0; + int runs = 1000; + for (int r = 0; r < runs; r++) { + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(4); + MatrixBlock result = compressor.decompress(compressor.compress(input)); + sum += result.get(0, 0); + } + double average = sum / runs; + + // Average should be close to original value (unbiased estimator) + assertEquals("Quantization should be unbiased", + 5.0, average, 0.5); // Allow 0.5 tolerance + } + + // ----------------------------------------------------------------------- + // Edge cases + // ----------------------------------------------------------------------- + + @Test + public void testConstantMatrix() throws Exception { + // All same value — min == max edge case + MatrixBlock input = new MatrixBlock(3, 3, false); + input.allocateDenseBlock(); + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) + input.set(i, j, 7.0); + input.examSparsity(); + + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(4); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + // Should not throw; all values should reconstruct to min (7.0 or close) + assertNotNull(result); + assertEquals(3, result.getNumRows()); + assertEquals(3, result.getNumColumns()); + } + + @Test(expected = IllegalArgumentException.class) + public void testInvalidBitsThrowsException() { + new ProbabilisticQuantizationCompressor(3); // Only 2, 4, 8 allowed + } + + // ----------------------------------------------------------------------- + // Helpers + // ----------------------------------------------------------------------- + + private MatrixBlock createRandomMatrix(int rows, int cols) { + MatrixBlock m = new MatrixBlock(rows, cols, false); + m.allocateDenseBlock(); + java.util.Random rng = new java.util.Random(42); + for (int i = 0; i < rows; i++) + for (int j = 0; j < cols; j++) + m.set(i, j, rng.nextGaussian() * 10); + m.examSparsity(); + return m; + } +} \ No newline at end of file From c6e5ec51123e87a11c6be783bf4574fcc941ece5 Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Mon, 29 Jun 2026 13:29:39 +0200 Subject: [PATCH 14/19] Add GitHub Actions CI/CD workflow for compression tests --- .github/workflows/compression-tests.yml | 46 +++++++++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 .github/workflows/compression-tests.yml diff --git a/.github/workflows/compression-tests.yml b/.github/workflows/compression-tests.yml new file mode 100644 index 00000000000..36bca793362 --- /dev/null +++ b/.github/workflows/compression-tests.yml @@ -0,0 +1,46 @@ +name: Compression Tests + +on: + push: + branches: [ feature/compression ] + pull_request: + branches: [ main ] + +jobs: + test: + runs-on: ubuntu-latest + + steps: + - name: Checkout code + uses: actions/checkout@v3 + + - name: Set up Java 17 + uses: actions/setup-java@v3 + with: + java-version: '17' + distribution: 'temurin' + + - name: Cache Maven dependencies + uses: actions/cache@v3 + with: + path: ~/.m2/repository + key: ${{ runner.os }}-maven-${{ hashFiles('**/pom.xml') }} + restore-keys: | + ${{ runner.os }}-maven- + + - name: Build project + run: mvn clean package -Dmaven.test.skip=true -Dmaven.javadoc.skip=true + + - name: Run compression tests + run: | + mvn test \ + -Dtest=TopKCompressorTest,ProbabilisticQuantizationCompressorTest \ + -Dmaven.test.failure.ignore=false \ + -Dmaven.javadoc.skip=true + + - name: Upload test results + if: always() + uses: actions/upload-artifact@v3 + with: + name: test-results + path: target/surefire-reports/ \ No newline at end of file From 43a937b2f92ec065a03c966e91b8ab9709d4dcd6 Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Mon, 29 Jun 2026 13:46:40 +0200 Subject: [PATCH 15/19] Integrate compression into FederationMap broadcast --- .../federated/FederationMap.java | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java index 68dff4785b3..c086e33b2ef 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java @@ -48,6 +48,12 @@ import org.apache.sysds.runtime.controlprogram.federated.FederatedRequest.RequestType; import org.apache.sysds.runtime.instructions.cp.CPOperand; import org.apache.sysds.runtime.instructions.cp.ScalarObject; +import org.apache.sysds.runtime.compress.CompressionConfig; +import org.apache.sysds.runtime.compress.CompressionFactory; +import org.apache.sysds.runtime.compress.CompressionType; +import org.apache.sysds.runtime.compress.CompressedMatrix; +import org.apache.sysds.runtime.compress.MatrixCompressor; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; import org.apache.sysds.runtime.instructions.cp.VariableCPInstruction; import org.apache.sysds.runtime.lineage.LineageItem; import org.apache.sysds.runtime.util.IndexRange; @@ -138,6 +144,27 @@ private FederatedRequest broadcast(CacheableData data, LineageItem lineageIte // is fine, because with broadcast all data on all workers) data.setFedMapping(copyWithNewIDAndRange( cb.getNumRows(), cb.getNumColumns(), id, FType.BROADCAST)); + + // === COMPRESSION INTEGRATION === + // Attempt TopK compression if the block is a MatrixBlock + if(cb instanceof MatrixBlock) { + try { + CompressionConfig config = CompressionConfig.builder() + .enable(true) + .withType(CompressionType.TOPK) + .withSparsity(0.01) + .build(); + MatrixCompressor compressor = CompressionFactory.create(config); + CompressedMatrix compressed = compressor.compress((MatrixBlock) cb); + MatrixBlock decompressed = compressor.decompress(compressed); + return new FederatedRequest(RequestType.PUT_VAR, lineageItem, id, decompressed); + } + catch(Exception ex) { + // Fall back to uncompressed on any error + } + } + // === END COMPRESSION INTEGRATION === + return new FederatedRequest(RequestType.PUT_VAR, lineageItem, id, cb); } From b173446b800da8a7291a90d4245c5a31ba34b5aa Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Mon, 29 Jun 2026 13:53:01 +0200 Subject: [PATCH 16/19] Fix deprecated upload-artifact v3 -> v4 --- .github/workflows/compression-tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/compression-tests.yml b/.github/workflows/compression-tests.yml index 36bca793362..ddfd08db244 100644 --- a/.github/workflows/compression-tests.yml +++ b/.github/workflows/compression-tests.yml @@ -40,7 +40,7 @@ jobs: - name: Upload test results if: always() - uses: actions/upload-artifact@v3 + uses: actions/upload-artifact@v4 with: name: test-results path: target/surefire-reports/ \ No newline at end of file From fcca7c0bb5d4f9f3a3f50cb328174dc0b2337ce8 Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Thu, 2 Jul 2026 20:00:15 +0200 Subject: [PATCH 17/19] Address code review feedback from David (ywcb00) --- .github/workflows/compression-tests.yml | 46 ---- .../java/org/apache/sysds/api/DMLScript.java | 2 + .../federated/FederationMap.java | 13 +- .../compression/CompressedMatrix.java | 81 ++++++ .../compression/CompressionConfig.java | 110 ++++++++ .../compression/CompressionFactory.java | 94 +++++++ .../compression/CompressionType.java | 57 +++++ .../compression/MatrixCompressor.java | 55 ++++ .../ProbabilisticQuantizationCompressor.java | 169 +++++++++++++ .../Quantization/QuantizedData.java | 74 ++++++ .../compression/TopK/TopKCompressor.java | 237 ++++++++++++++++++ .../federated/compression/TopK/TopKData.java | 61 +++++ .../exceptions/CompressionException.java | 36 +++ .../exceptions/DecompressionException.java | 36 +++ ...obabilisticQuantizationCompressorTest.java | 85 ++++--- .../federated/io}/TopKCompressorTest.java | 103 ++++---- 16 files changed, 1114 insertions(+), 145 deletions(-) delete mode 100644 .github/workflows/compression-tests.yml create mode 100644 src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressedMatrix.java create mode 100644 src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionConfig.java create mode 100644 src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionFactory.java create mode 100644 src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionType.java create mode 100644 src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/MatrixCompressor.java create mode 100644 src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/ProbabilisticQuantizationCompressor.java create mode 100644 src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/QuantizedData.java create mode 100644 src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKCompressor.java create mode 100644 src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKData.java create mode 100644 src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/CompressionException.java create mode 100644 src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/DecompressionException.java rename src/test/java/org/apache/sysds/{runtime/compress/Quantization => test/functions/federated/io}/ProbabilisticQuantizationCompressorTest.java (77%) rename src/test/java/org/apache/sysds/{runtime/compress/TopK => test/functions/federated/io}/TopKCompressorTest.java (73%) diff --git a/.github/workflows/compression-tests.yml b/.github/workflows/compression-tests.yml deleted file mode 100644 index ddfd08db244..00000000000 --- a/.github/workflows/compression-tests.yml +++ /dev/null @@ -1,46 +0,0 @@ -name: Compression Tests - -on: - push: - branches: [ feature/compression ] - pull_request: - branches: [ main ] - -jobs: - test: - runs-on: ubuntu-latest - - steps: - - name: Checkout code - uses: actions/checkout@v3 - - - name: Set up Java 17 - uses: actions/setup-java@v3 - with: - java-version: '17' - distribution: 'temurin' - - - name: Cache Maven dependencies - uses: actions/cache@v3 - with: - path: ~/.m2/repository - key: ${{ runner.os }}-maven-${{ hashFiles('**/pom.xml') }} - restore-keys: | - ${{ runner.os }}-maven- - - - name: Build project - run: mvn clean package -Dmaven.test.skip=true -Dmaven.javadoc.skip=true - - - name: Run compression tests - run: | - mvn test \ - -Dtest=TopKCompressorTest,ProbabilisticQuantizationCompressorTest \ - -Dmaven.test.failure.ignore=false \ - -Dmaven.javadoc.skip=true - - - name: Upload test results - if: always() - uses: actions/upload-artifact@v4 - with: - name: test-results - path: target/surefire-reports/ \ No newline at end of file diff --git a/src/main/java/org/apache/sysds/api/DMLScript.java b/src/main/java/org/apache/sysds/api/DMLScript.java index c286b8d3b52..8391634fdb1 100644 --- a/src/main/java/org/apache/sysds/api/DMLScript.java +++ b/src/main/java/org/apache/sysds/api/DMLScript.java @@ -131,6 +131,8 @@ public class DMLScript public static double GPU_MEMORY_UTILIZATION_FACTOR = 0.9; // Set GPU memory allocator to use public static String GPU_MEMORY_ALLOCATOR = "cuda"; + // Enable/disable federated compression + public static boolean FEDERATED_COMPRESSION = false; // Enable/disable lineage tracing public static boolean LINEAGE = DMLOptions.defaultOptions.lineage; // Enable/disable deduplicate lineage items diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java index c086e33b2ef..4d7a06e5614 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java @@ -42,17 +42,18 @@ import org.apache.sysds.hops.fedplanner.FTypes.AlignType; import org.apache.sysds.hops.fedplanner.FTypes.FType; import org.apache.sysds.lops.RightIndex; +import org.apache.sysds.api.DMLScript; import org.apache.sysds.runtime.DMLRuntimeException; import org.apache.sysds.runtime.controlprogram.caching.CacheBlock; import org.apache.sysds.runtime.controlprogram.caching.CacheableData; import org.apache.sysds.runtime.controlprogram.federated.FederatedRequest.RequestType; import org.apache.sysds.runtime.instructions.cp.CPOperand; import org.apache.sysds.runtime.instructions.cp.ScalarObject; -import org.apache.sysds.runtime.compress.CompressionConfig; -import org.apache.sysds.runtime.compress.CompressionFactory; -import org.apache.sysds.runtime.compress.CompressionType; -import org.apache.sysds.runtime.compress.CompressedMatrix; -import org.apache.sysds.runtime.compress.MatrixCompressor; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionConfig; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionFactory; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionType; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix; +import org.apache.sysds.runtime.controlprogram.federated.compression.MatrixCompressor; import org.apache.sysds.runtime.matrix.data.MatrixBlock; import org.apache.sysds.runtime.instructions.cp.VariableCPInstruction; import org.apache.sysds.runtime.lineage.LineageItem; @@ -147,7 +148,7 @@ private FederatedRequest broadcast(CacheableData data, LineageItem lineageIte // === COMPRESSION INTEGRATION === // Attempt TopK compression if the block is a MatrixBlock - if(cb instanceof MatrixBlock) { + if(DMLScript.FEDERATED_COMPRESSION && cb instanceof MatrixBlock) { try { CompressionConfig config = CompressionConfig.builder() .enable(true) diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressedMatrix.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressedMatrix.java new file mode 100644 index 00000000000..6f349b2d7ab --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressedMatrix.java @@ -0,0 +1,81 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression; + +import java.io.Serializable; + +/** + * Generic container for compressed matrix data. + * Stores the compressed representation along with metadata + * needed for decompression and size estimation. + */ +public class CompressedMatrix implements Serializable { + + private static final long serialVersionUID = 1L; + + private final CompressionType type; + private final int numRows; + private final int numCols; + private final Object compressedData; // Technique-specific data + private final double compressionRatio; + private final byte[] metadata; // Optional: scaling factors, etc. + + public CompressedMatrix(CompressionType type, int numRows, int numCols, + Object compressedData, double compressionRatio) { + this(type, numRows, numCols, compressedData, compressionRatio, null); + } + + public CompressedMatrix(CompressionType type, int numRows, int numCols, + Object compressedData, double compressionRatio, + byte[] metadata) { + this.type = type; + this.numRows = numRows; + this.numCols = numCols; + this.compressedData = compressedData; + this.compressionRatio = compressionRatio; + this.metadata = metadata; + } + + public CompressionType getType() { return type; } + public int getNumRows() { return numRows; } + public int getNumCols() { return numCols; } + public Object getCompressedData() { return compressedData; } + public double getCompressionRatio() { return compressionRatio; } + public byte[] getMetadata() { return metadata; } + + /** Estimate original size in bytes (8 bytes per double) */ + public long estimateOriginalSizeBytes() { + return (long) numRows * numCols * 8; + } + + /** Estimate compressed size in bytes */ + public long getCompressedSizeBytes() { + if (compressedData instanceof byte[]) { + return ((byte[]) compressedData).length; + } + return 0; + } + + @Override + public String toString() { + return String.format("CompressedMatrix[%s, %dx%d, ratio=%.2fx]", + type.getId(), numRows, numCols, compressionRatio); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionConfig.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionConfig.java new file mode 100644 index 00000000000..4f3bfa33554 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionConfig.java @@ -0,0 +1,110 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression; + +import java.util.HashMap; +import java.util.Map; + +/** + * Immutable configuration for compression in federated operations. + * Uses the Builder pattern for flexible, readable configuration. + * + * Usage example: + * CompressionConfig config = CompressionConfig.builder() + * .enable(true) + * .withType(CompressionType.TOPK) + * .withSparsity(0.01) + * .build(); + */ +public class CompressionConfig { + + private final boolean enabled; + private final CompressionType type; + private final Map parameters; + + private CompressionConfig(Builder builder) { + this.enabled = builder.enabled; + this.type = builder.enabled ? builder.type : CompressionType.NONE; + this.parameters = new HashMap<>(builder.parameters); + } + + public boolean isEnabled() { return enabled; } + public CompressionType getType() { return type; } + public Map getParameters() { return new HashMap<>(parameters); } + + /** Convenience getter for sparsity parameter (TopK) */ + public double getSparsity() { + return (double) parameters.getOrDefault("sparsity", 0.01); + } + + /** Convenience getter for bits parameter (Quantization) */ + public int getBits() { + return (int) parameters.getOrDefault("bits", 4); + } + + @Override + public String toString() { + return String.format("CompressionConfig[enabled=%s, type=%s, params=%s]", + enabled, type.getId(), parameters); + } + + // ----------------------------------------------------------------------- + // Builder + // ----------------------------------------------------------------------- + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + private boolean enabled = false; + private CompressionType type = CompressionType.NONE; + private final Map parameters = new HashMap<>(); + + public Builder enable(boolean enabled) { + this.enabled = enabled; + return this; + } + + public Builder withType(CompressionType type) { + this.type = type; + return this; + } + + public Builder withParameter(String key, Object value) { + this.parameters.put(key, value); + return this; + } + + /** Shorthand for TopK sparsity ratio */ + public Builder withSparsity(double sparsity) { + return withParameter("sparsity", sparsity); + } + + /** Shorthand for quantization bit width */ + public Builder withBits(int bits) { + return withParameter("bits", bits); + } + + public CompressionConfig build() { + return new CompressionConfig(this); + } + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionFactory.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionFactory.java new file mode 100644 index 00000000000..bb4919415ee --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionFactory.java @@ -0,0 +1,94 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression; + +import org.apache.sysds.runtime.controlprogram.federated.compression.TopK.TopKCompressor; +import org.apache.sysds.runtime.controlprogram.federated.compression.Quantization.ProbabilisticQuantizationCompressor; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.CompressionException; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.DecompressionException; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; + +/** + * Factory for creating compressor instances from configuration. + * Centralizes compressor instantiation and parameter validation. + * + * Usage: + * CompressionConfig config = CompressionConfig.builder() + * .enable(true) + * .withType(CompressionType.TOPK) + * .withSparsity(0.01) + * .build(); + * MatrixCompressor compressor = CompressionFactory.create(config); + */ +public class CompressionFactory { + + private CompressionFactory() { + // Utility class — no instantiation + } + + /** + * Create a compressor from a CompressionConfig. + * @param config The compression configuration + * @return A ready-to-use MatrixCompressor + * @throws IllegalArgumentException if the config is invalid + */ + public static MatrixCompressor create(CompressionConfig config) { + if(config == null || !config.isEnabled()) + return new PassthroughCompressor(); + switch(config.getType()) { + case TOPK: + return new TopKCompressor(config.getSparsity(), true); + case PROBABILISTIC_QUANTIZATION: + return new ProbabilisticQuantizationCompressor(config.getBits()); + case NONE: + default: + return new PassthroughCompressor(); + } + } + + // ----------------------------------------------------------------------- + // Passthrough compressor (no-op) for when compression is disabled + // ----------------------------------------------------------------------- + + /** + * No-op compressor: returns the matrix as-is. + * Used when compression is disabled or type is NONE. + */ + private static class PassthroughCompressor implements MatrixCompressor { + + @Override + public CompressedMatrix compress(MatrixBlock input) + throws CompressionException { + return new CompressedMatrix( + CompressionType.NONE, + input.getNumRows(), + input.getNumColumns(), + input, + 1.0 + ); + } + + @Override + public MatrixBlock decompress(CompressedMatrix compressed) + throws DecompressionException { + return (MatrixBlock) compressed.getCompressedData(); + } + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionType.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionType.java new file mode 100644 index 00000000000..d942fa9e089 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionType.java @@ -0,0 +1,57 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression; + +/** + * Enumeration of supported compression techniques for federated learning. + * Used for configuration, serialization, and technique selection. + */ +public enum CompressionType { + + /** TopK sparsification: keep largest-magnitude elements only */ + TOPK("topk", "Top-K Sparsification"), + + /** Probabilistic quantization: reduce precision with stochastic rounding */ + PROBABILISTIC_QUANTIZATION("prob_quant", "Probabilistic Quantization"), + + /** No compression (passthrough) */ + NONE("none", "No Compression"); + + private final String id; + private final String description; + + CompressionType(String id, String description) { + this.id = id; + this.description = description; + } + + public String getId() { return id; } + public String getDescription() { return description; } + + /** Parse from string identifier (case-insensitive) */ + public static CompressionType fromString(String text) { + for (CompressionType type : CompressionType.values()) { + if (type.id.equalsIgnoreCase(text)) { + return type; + } + } + throw new IllegalArgumentException("Unknown compression type: " + text); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/MatrixCompressor.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/MatrixCompressor.java new file mode 100644 index 00000000000..12e43b3d7a6 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/MatrixCompressor.java @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression; + +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.CompressionException; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.DecompressionException; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; + +/** + * Interface for matrix compression techniques in federated learning. + * All compressors must implement compress/decompress operations. + */ +public interface MatrixCompressor { + + /** + * Compress a matrix block for transmission. + * @param input The source matrix to compress + * @return CompressedMatrix containing compressed data and metadata + * @throws CompressionException if compression fails + */ + CompressedMatrix compress(MatrixBlock input) throws CompressionException; + + /** + * Decompress a compressed matrix back to MatrixBlock. + * @param compressed The compressed data to decompress + * @return Reconstructed MatrixBlock (may be approximate) + * @throws DecompressionException if decompression fails + */ + MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException; + + /** + * Estimate the compression ratio achieved. + * Higher is better (e.g. 10.0 means 10x smaller). + */ + default double estimateCompressionRatio(long originalSize, long compressedSize) { + return compressedSize == 0 ? Double.MAX_VALUE : (double) originalSize / compressedSize; + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/ProbabilisticQuantizationCompressor.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/ProbabilisticQuantizationCompressor.java new file mode 100644 index 00000000000..c46a9ffb874 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/ProbabilisticQuantizationCompressor.java @@ -0,0 +1,169 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression.Quantization; + +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionType; +import org.apache.sysds.runtime.controlprogram.federated.compression.MatrixCompressor; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.CompressionException; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.DecompressionException; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; + +import java.util.Random; + +/** + * Probabilistic Quantization Compressor. + * + * Reduces numerical precision using stochastic rounding to maintain + * an unbiased estimator — meaning E[quantized] = original on average. + * This is critical for federated learning convergence guarantees. + * + * Supports 2, 4, or 8 bits per value: + * 2-bit → 4 levels → 16x compression vs 32-bit float + * 4-bit → 16 levels → 8x compression + * 8-bit → 256 levels → 4x compression + */ +public class ProbabilisticQuantizationCompressor implements MatrixCompressor { + + private final int bitsPerValue; // 2, 4, or 8 + private final Random rng; + + public ProbabilisticQuantizationCompressor(int bitsPerValue) { + if (bitsPerValue != 2 && bitsPerValue != 4 && bitsPerValue != 8) { + throw new IllegalArgumentException("bitsPerValue must be 2, 4, or 8"); + } + this.bitsPerValue = bitsPerValue; + this.rng = new Random(42); // Fixed seed for reproducibility + } + + @Override + public CompressedMatrix compress(MatrixBlock input) throws CompressionException { + try { + int numRows = input.getNumRows(); + int numCols = input.getNumColumns(); + int totalElements = numRows * numCols; + + // Find min and max for normalization + double[] minMax = findMinMax(input, numRows, numCols); + double min = minMax[0]; + double max = minMax[1]; + + int levels = 1 << bitsPerValue; // 2^bits + + // Quantize each element probabilistically + byte[] quantized = new byte[totalElements]; + for (int i = 0; i < numRows; i++) { + for (int j = 0; j < numCols; j++) { + double value = input.get(i, j); + quantized[i * numCols + j] = probabilisticRound(value, min, max, levels); + } + } + + double ratio = 32.0 / bitsPerValue; // vs 32-bit float + + QuantizedData data = new QuantizedData( + quantized, min, max, levels, bitsPerValue, numRows, numCols); + + return new CompressedMatrix( + CompressionType.PROBABILISTIC_QUANTIZATION, + numRows, numCols, data, ratio); + + } catch (Exception e) { + throw new CompressionException( + "Quantization compression failed: " + e.getMessage(), e); + } + } + + @Override + public MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException { + try { + QuantizedData data = (QuantizedData) compressed.getCompressedData(); + MatrixBlock result = new MatrixBlock(data.numRows, data.numCols, false); + result.allocateDenseBlock(); + + for (int i = 0; i < data.numRows; i++) { + for (int j = 0; j < data.numCols; j++) { + byte levelIndex = data.quantizedValues[i * data.numCols + j]; + double value = data.reconstructValue(levelIndex); + result.set(i, j, value); + } + } + + result.examSparsity(); + return result; + + } catch (ClassCastException e) { + throw new DecompressionException( + "Invalid compressed data type for Quantization", e); + } catch (Exception e) { + throw new DecompressionException( + "Quantization decompression failed: " + e.getMessage(), e); + } + } + + // ----------------------------------------------------------------------- + // Private helpers + // ----------------------------------------------------------------------- + + /** + * Stochastic rounding: for value x between levels q_i and q_{i+1}: + * P(round up) = (x - q_i) / (q_{i+1} - q_i) + * P(round down) = 1 - P(round up) + * This gives E[output] = x (unbiased). + */ + private byte probabilisticRound(double value, double min, double max, int levels) { + // Handle constant matrix edge case + if (max - min < 1e-10) return 0; + + // Normalize to [0, 1] + double normalized = (value - min) / (max - min); + + // Find bounding level indices + double scaled = normalized * (levels - 1); + int lowerIdx = (int) scaled; + int upperIdx = lowerIdx + 1; + + if (lowerIdx == upperIdx) { + return (byte) lowerIdx; + } + + // Probabilistic decision + double probUp = scaled - lowerIdx; + return (rng.nextDouble() < probUp) ? (byte) upperIdx : (byte) lowerIdx; + } + + /** Find min and max values across the entire matrix */ + private double[] findMinMax(MatrixBlock input, int numRows, int numCols) { + double min = Double.MAX_VALUE; + double max = -Double.MAX_VALUE; + + for (int i = 0; i < numRows; i++) { + for (int j = 0; j < numCols; j++) { + double val = input.get(i, j); + if (val < min) min = val; + if (val > max) max = val; + } + } + + // Handle all-zero matrix + if (min == Double.MAX_VALUE) { min = 0; max = 0; } + return new double[]{min, max}; + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/QuantizedData.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/QuantizedData.java new file mode 100644 index 00000000000..595ae482556 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/QuantizedData.java @@ -0,0 +1,74 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression.Quantization; + +import java.io.Serializable; + +/** + * Immutable container for probabilistically quantized matrix data. + * Stores quantized byte indices and the scaling parameters needed + * to reconstruct approximate original values on decompression. + */ +public class QuantizedData implements Serializable { + + private static final long serialVersionUID = 1L; + + public final byte[] quantizedValues; // Quantized level indices + public final double min; // Original minimum value + public final double max; // Original maximum value + public final int levels; // Number of quantization levels (2^bits) + public final int bitsPerValue; // Bits used per element + public final int numRows; + public final int numCols; + + public QuantizedData(byte[] quantizedValues, double min, double max, + int levels, int bitsPerValue, int numRows, int numCols) { + this.quantizedValues = quantizedValues.clone(); // Defensive copy + this.min = min; + this.max = max; + this.levels = levels; + this.bitsPerValue = bitsPerValue; + this.numRows = numRows; + this.numCols = numCols; + } + + /** Number of quantized elements */ + public int size() { + return quantizedValues.length; + } + + /** Estimate serialized size in bytes */ + public long estimateSizeBytes() { + return quantizedValues.length + 64; // +64 for scalar fields and headers + } + + /** Reconstruct a double value from a quantized level index */ + public double reconstructValue(byte levelIndex) { + if (max - min < 1e-10) return min; // Constant matrix + int idx = levelIndex & 0xFF; // Treat byte as unsigned + return min + (idx / (double)(levels - 1)) * (max - min); + } + + @Override + public String toString() { + return String.format("QuantizedData[%dx%d, levels=%d, bits=%d, min=%.4f, max=%.4f]", + numRows, numCols, levels, bitsPerValue, min, max); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKCompressor.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKCompressor.java new file mode 100644 index 00000000000..f21834714f2 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKCompressor.java @@ -0,0 +1,237 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression.TopK; + +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionType; +import org.apache.sysds.runtime.controlprogram.federated.compression.MatrixCompressor; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.CompressionException; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.DecompressionException; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.PriorityQueue; + +/** + * TopK Sparsification Compressor. + * + * Keeps only the K largest-magnitude elements in the matrix, + * setting all others to zero. Optimal for gradient sparsification + * in federated learning where most gradient values are near zero. + * + * Compression ratio: approximately 1/sparsityRatio + * e.g. sparsityRatio=0.01 keeps 1% of elements → ~100x compression + */ +public class TopKCompressor implements MatrixCompressor { + + private final double sparsityRatio; // Fraction of elements to keep (0, 1] + private final boolean useHeap; // Use min-heap for O(n log k) selection + + /** + * @param sparsityRatio Fraction of elements to retain e.g. 0.01 = keep top 1% + * @param useHeap If true, use priority queue (faster for large matrices) + */ + public TopKCompressor(double sparsityRatio, boolean useHeap) { + if (sparsityRatio <= 0 || sparsityRatio > 1) { + throw new IllegalArgumentException("sparsityRatio must be in (0, 1]"); + } + this.sparsityRatio = sparsityRatio; + this.useHeap = useHeap; + } + + /** Default constructor: uses heap-based selection */ + public TopKCompressor(double sparsityRatio) { + this(sparsityRatio, true); + } + + @Override + public CompressedMatrix compress(MatrixBlock input) throws CompressionException { + try { + int numRows = input.getNumRows(); + int numCols = input.getNumColumns(); + int totalElements = numRows * numCols; + int k = (int) Math.max(1, Math.ceil(totalElements * sparsityRatio)); + + // If k covers everything, no compression needed + if (k >= totalElements) { + return new CompressedMatrix( + CompressionType.TOPK, numRows, numCols, input, 1.0); + } + + // Extract all non-zero elements with their linear indices + List elements = extractElements(input, numRows, numCols); + + // If fewer non-zeros than k, keep all of them + List topK = (elements.size() <= k) + ? new ArrayList<>(elements) + : selectTopK(elements, k); + + // Pack into TopKData + TopKData data = convertToTopKData(topK, numCols); + + double ratio = calculateCompressionRatio(totalElements, topK.size()); + + return new CompressedMatrix( + CompressionType.TOPK, numRows, numCols, data, ratio); + + } catch (Exception e) { + throw new CompressionException("TopK compression failed: " + e.getMessage(), e); + } + } + + @Override + public MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException { + try { + // Handle passthrough case (no compression was applied) + if (compressed.getCompressedData() instanceof MatrixBlock) { + return (MatrixBlock) compressed.getCompressedData(); + } + + TopKData data = (TopKData) compressed.getCompressedData(); + MatrixBlock result = new MatrixBlock( + compressed.getNumRows(), + compressed.getNumCols(), + true // Start sparse + ); + result.allocateSparseRowsBlock(); + + // Place values back at their original positions + for (int i = 0; i < data.indices.length; i++) { + int linearIdx = data.indices[i]; + int row = linearIdx / data.numCols; + int col = linearIdx % data.numCols; + result.set(row, col, data.values[i]); + } + + result.examSparsity(); + return result; + + } catch (ClassCastException e) { + throw new DecompressionException("Invalid compressed data type for TopK", e); + } catch (Exception e) { + throw new DecompressionException("TopK decompression failed: " + e.getMessage(), e); + } + } + + // ----------------------------------------------------------------------- + // Private helpers + // ----------------------------------------------------------------------- + + /** + * Extract all non-zero elements with their linear indices. + * Handles both dense and sparse MatrixBlock representations. + */ + private List extractElements(MatrixBlock input, int numRows, int numCols) { + List elements = new ArrayList<>(); + + if (input.isInSparseFormat()) { + // Sparse: iterate only over non-zero entries + for (int i = 0; i < numRows; i++) { + if (input.getSparseBlock() == null) continue; + if (input.getSparseBlock().isEmpty(i)) continue; + int[] rowIndices = input.getSparseBlock().indexes(i); + double[] rowValues = input.getSparseBlock().values(i); + int rowSize = input.getSparseBlock().size(i); + for (int j = 0; j < rowSize; j++) { + double val = rowValues[j]; + if (val != 0.0) { + int linearIdx = i * numCols + rowIndices[j]; + elements.add(new Element(linearIdx, val, Math.abs(val))); + } + } + } + } else { + // Dense: iterate all elements, skip zeros + double[] denseBlock = input.getDenseBlockValues(); + if (denseBlock != null) { + for (int i = 0; i < denseBlock.length; i++) { + if (denseBlock[i] != 0.0) { + elements.add(new Element(i, denseBlock[i], Math.abs(denseBlock[i]))); + } + } + } + } + return elements; + } + + /** + * Select top K elements by absolute value. + * Uses min-heap for O(n log k) when useHeap=true, + * or full sort O(n log n) otherwise. + */ + private List selectTopK(List elements, int k) { + if (useHeap) { + PriorityQueue minHeap = new PriorityQueue<>( + k, Comparator.comparingDouble(e -> e.absValue) + ); + for (Element e : elements) { + if (minHeap.size() < k) { + minHeap.offer(e); + } else if (e.absValue > minHeap.peek().absValue) { + minHeap.poll(); + minHeap.offer(e); + } + } + List result = new ArrayList<>(minHeap); + result.sort(Comparator.comparingInt(e -> e.index)); + return result; + } else { + elements.sort((a, b) -> Double.compare(b.absValue, a.absValue)); + return new ArrayList<>(elements.subList(0, k)); + } + } + + private TopKData convertToTopKData(List topK, int numCols) { + int[] indices = new int[topK.size()]; + double[] values = new double[topK.size()]; + for (int i = 0; i < topK.size(); i++) { + indices[i] = topK.get(i).index; + values[i] = topK.get(i).value; + } + return new TopKData(indices, values, numCols); + } + + private double calculateCompressionRatio(int total, int kept) { + if (kept == 0) return Double.MAX_VALUE; + // Original: total * 8 bytes (doubles) + // Compressed: kept * 12 bytes (int index + double value) + long originalBytes = (long) total * 8; + long compressedBytes = (long) kept * 12; + return (double) originalBytes / compressedBytes; + } + + // ----------------------------------------------------------------------- + // Inner class: element tracking during compression + // ----------------------------------------------------------------------- + + private static class Element { + final int index; // Linear index: row * numCols + col + final double value; // Original value + final double absValue; // Absolute value for magnitude comparison + + Element(int index, double value, double absValue) { + this.index = index; + this.value = value; + this.absValue = absValue; + } + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKData.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKData.java new file mode 100644 index 00000000000..eb737868c29 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKData.java @@ -0,0 +1,61 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression.TopK; + +import java.io.Serializable; + +/** + * Immutable container for TopK-compressed matrix data. + * Stores only the K largest-magnitude elements with their positions, + * designed for efficient serialization across federated workers. + */ +public class TopKData implements Serializable { + + private static final long serialVersionUID = 1L; + + public final int[] indices; // Linear indices of kept elements (row*numCols + col) + public final double[] values; // Corresponding original values + public final int numCols; // Needed for index → (row, col) conversion + + public TopKData(int[] indices, double[] values, int numCols) { + if (indices.length != values.length) { + throw new IllegalArgumentException( + "Indices and values arrays must have the same length"); + } + this.indices = indices.clone(); // Defensive copy + this.values = values.clone(); + this.numCols = numCols; + } + + /** Number of kept elements */ + public int size() { + return indices.length; + } + + /** Estimate serialized size in bytes (4 bytes per int + 8 bytes per double) */ + public long estimateSizeBytes() { + return (long) indices.length * 12 + 64; // +64 for object headers + } + + @Override + public String toString() { + return String.format("TopKData[k=%d, numCols=%d]", indices.length, numCols); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/CompressionException.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/CompressionException.java new file mode 100644 index 00000000000..b56d44f23f0 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/CompressionException.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression.exceptions; + +/** + * Exception thrown when matrix compression fails. + */ +public class CompressionException extends Exception { + + private static final long serialVersionUID = 1L; + + public CompressionException(String message) { + super(message); + } + + public CompressionException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/DecompressionException.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/DecompressionException.java new file mode 100644 index 00000000000..7680d80e98e --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/DecompressionException.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression.exceptions; + +/** + * Exception thrown when matrix decompression fails. + */ +public class DecompressionException extends Exception { + + private static final long serialVersionUID = 1L; + + public DecompressionException(String message) { + super(message); + } + + public DecompressionException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/src/test/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressorTest.java b/src/test/java/org/apache/sysds/test/functions/federated/io/ProbabilisticQuantizationCompressorTest.java similarity index 77% rename from src/test/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressorTest.java rename to src/test/java/org/apache/sysds/test/functions/federated/io/ProbabilisticQuantizationCompressorTest.java index a50ae2f3f79..17105385fb6 100644 --- a/src/test/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressorTest.java +++ b/src/test/java/org/apache/sysds/test/functions/federated/io/ProbabilisticQuantizationCompressorTest.java @@ -1,53 +1,63 @@ -package org.apache.sysds.runtime.compress.Quantization; +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ -import org.apache.sysds.runtime.compress.CompressedMatrix; -import org.apache.sysds.runtime.compress.CompressionType; +package org.apache.sysds.test.functions.federated.io; + +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionType; +import org.apache.sysds.runtime.controlprogram.federated.compression.Quantization.ProbabilisticQuantizationCompressor; import org.apache.sysds.runtime.matrix.data.MatrixBlock; +import org.apache.sysds.test.AutomatedTestBase; import org.junit.Test; import static org.junit.Assert.*; -/** - * Unit tests for ProbabilisticQuantizationCompressor. - * Verifies compression ratio, reconstruction accuracy, - * unbiasedness property, and edge case handling. - * - * - */ -public class ProbabilisticQuantizationCompressorTest { +public class ProbabilisticQuantizationCompressorTest extends AutomatedTestBase { + + @Override + public void setUp() {} + + @Override + public void tearDown() {} // ----------------------------------------------------------------------- - // Basic compression / decompression + // Basic properties (merged from testCompressionTypeIsProbabilisticQuantization, + // testDimensionsPreservedAfterDecompression, testReconstructedValuesWithinOriginalRange, + // testCompressionRatio4Bit) // ----------------------------------------------------------------------- @Test - public void testCompressionTypeIsProbabilisticQuantization() throws Exception { - MatrixBlock input = createRandomMatrix(4, 4); + public void testQuantizationBasicProperties() throws Exception { + MatrixBlock input = createRandomMatrix(10, 20); ProbabilisticQuantizationCompressor compressor = new ProbabilisticQuantizationCompressor(4); CompressedMatrix compressed = compressor.compress(input); + assertEquals(CompressionType.PROBABILISTIC_QUANTIZATION, compressed.getType()); - } - @Test - public void testDimensionsPreservedAfterDecompression() throws Exception { - MatrixBlock input = createRandomMatrix(10, 20); - ProbabilisticQuantizationCompressor compressor = - new ProbabilisticQuantizationCompressor(4); - CompressedMatrix compressed = compressor.compress(input); MatrixBlock result = compressor.decompress(compressed); - assertEquals(10, result.getNumRows()); assertEquals(20, result.getNumColumns()); - } - - @Test - public void testReconstructedValuesWithinOriginalRange() throws Exception { - MatrixBlock input = createRandomMatrix(20, 20); // Find original min/max double origMin = Double.MAX_VALUE; double origMax = -Double.MAX_VALUE; - for (int i = 0; i < 20; i++) { + for (int i = 0; i < 10; i++) { for (int j = 0; j < 20; j++) { double v = input.get(i, j); if (v < origMin) origMin = v; @@ -55,18 +65,16 @@ public void testReconstructedValuesWithinOriginalRange() throws Exception { } } - ProbabilisticQuantizationCompressor compressor = - new ProbabilisticQuantizationCompressor(4); - MatrixBlock result = compressor.decompress(compressor.compress(input)); - // All reconstructed values must be within [min, max] - for (int i = 0; i < 20; i++) { + for (int i = 0; i < 10; i++) { for (int j = 0; j < 20; j++) { double v = result.get(i, j); assertTrue("Value below min: " + v, v >= origMin - 1e-9); assertTrue("Value above max: " + v, v <= origMax + 1e-9); } } + + assertEquals(8.0, compressed.getCompressionRatio(), 1e-10); } // ----------------------------------------------------------------------- @@ -82,15 +90,6 @@ public void testCompressionRatio2Bit() throws Exception { assertEquals(16.0, compressed.getCompressionRatio(), 1e-10); } - @Test - public void testCompressionRatio4Bit() throws Exception { - MatrixBlock input = createRandomMatrix(10, 10); - ProbabilisticQuantizationCompressor compressor = - new ProbabilisticQuantizationCompressor(4); - CompressedMatrix compressed = compressor.compress(input); - assertEquals(8.0, compressed.getCompressionRatio(), 1e-10); - } - @Test public void testCompressionRatio8Bit() throws Exception { MatrixBlock input = createRandomMatrix(10, 10); @@ -184,4 +183,4 @@ private MatrixBlock createRandomMatrix(int rows, int cols) { m.examSparsity(); return m; } -} \ No newline at end of file +} diff --git a/src/test/java/org/apache/sysds/runtime/compress/TopK/TopKCompressorTest.java b/src/test/java/org/apache/sysds/test/functions/federated/io/TopKCompressorTest.java similarity index 73% rename from src/test/java/org/apache/sysds/runtime/compress/TopK/TopKCompressorTest.java rename to src/test/java/org/apache/sysds/test/functions/federated/io/TopKCompressorTest.java index 3fd1fd426f2..65b16c9e429 100644 --- a/src/test/java/org/apache/sysds/runtime/compress/TopK/TopKCompressorTest.java +++ b/src/test/java/org/apache/sysds/test/functions/federated/io/TopKCompressorTest.java @@ -1,19 +1,60 @@ -package org.apache.sysds.runtime.compress.TopK; +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.test.functions.federated.io; -import org.apache.sysds.runtime.compress.CompressedMatrix; -import org.apache.sysds.runtime.compress.CompressionType; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionType; +import org.apache.sysds.runtime.controlprogram.federated.compression.TopK.TopKCompressor; import org.apache.sysds.runtime.matrix.data.MatrixBlock; +import org.apache.sysds.test.AutomatedTestBase; import org.junit.Test; import static org.junit.Assert.*; -/** - * Unit tests for TopKCompressor. - * Verifies compression ratio, reconstruction accuracy, - * and correct handling of edge cases. - * - * - */ -public class TopKCompressorTest { +public class TopKCompressorTest extends AutomatedTestBase { + + @Override + public void setUp() {} + + @Override + public void tearDown() {} + + // ----------------------------------------------------------------------- + // Basic properties (merged from testCompressionTypeIsTopK, + // testDimensionsPreservedAfterDecompression, testCompressionRatioIsPositive) + // ----------------------------------------------------------------------- + + @Test + public void testTopKBasicProperties() throws Exception { + MatrixBlock input = createRandomMatrix(10, 20); + TopKCompressor compressor = new TopKCompressor(0.1); + CompressedMatrix compressed = compressor.compress(input); + + assertEquals(CompressionType.TOPK, compressed.getType()); + + MatrixBlock result = compressor.decompress(compressed); + assertEquals(10, result.getNumRows()); + assertEquals(20, result.getNumColumns()); + + assertTrue("Compression ratio must be > 0", + compressed.getCompressionRatio() > 0); + } // ----------------------------------------------------------------------- // Basic compression / decompression @@ -42,38 +83,10 @@ public void testTopKKeepsLargestElements() throws Exception { assertEquals(0.0, result.get(2, 2), 1e-10); } - @Test - public void testCompressionTypeIsTopK() throws Exception { - MatrixBlock input = createDenseMatrix(4, 4, 1.0); - TopKCompressor compressor = new TopKCompressor(0.5); - CompressedMatrix compressed = compressor.compress(input); - assertEquals(CompressionType.TOPK, compressed.getType()); - } - - @Test - public void testDimensionsPreservedAfterDecompression() throws Exception { - MatrixBlock input = createRandomMatrix(10, 20); - TopKCompressor compressor = new TopKCompressor(0.1); - CompressedMatrix compressed = compressor.compress(input); - MatrixBlock result = compressor.decompress(compressed); - - assertEquals(10, result.getNumRows()); - assertEquals(20, result.getNumColumns()); - } - // ----------------------------------------------------------------------- // Compression ratio // ----------------------------------------------------------------------- - @Test - public void testCompressionRatioIsPositive() throws Exception { - MatrixBlock input = createRandomMatrix(50, 50); - TopKCompressor compressor = new TopKCompressor(0.01); - CompressedMatrix compressed = compressor.compress(input); - assertTrue("Compression ratio must be > 0", - compressed.getCompressionRatio() > 0); - } - @Test public void testLowerSparsityGivesHigherRatio() throws Exception { MatrixBlock input = createRandomMatrix(100, 100); @@ -144,14 +157,4 @@ private MatrixBlock createRandomMatrix(int rows, int cols) { m.examSparsity(); return m; } - - private MatrixBlock createDenseMatrix(int rows, int cols, double fillValue) { - MatrixBlock m = new MatrixBlock(rows, cols, false); - m.allocateDenseBlock(); - for (int i = 0; i < rows; i++) - for (int j = 0; j < cols; j++) - m.set(i, j, fillValue); - m.examSparsity(); - return m; - } -} \ No newline at end of file +} From be8f8697315604b5a2f21f7003fbab8d07b448fc Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Fri, 3 Jul 2026 00:15:23 +0200 Subject: [PATCH 18/19] Fix test files: add content, remove AutomatedTestBase, fix tab indentation --- .../compression/CompressedMatrix.java | 86 ++-- .../compression/CompressionConfig.java | 144 +++---- .../compression/CompressionFactory.java | 94 ++--- .../compression/CompressionType.java | 58 +-- .../compression/MatrixCompressor.java | 42 +- .../ProbabilisticQuantizationCompressor.java | 248 ++++++------ .../Quantization/QuantizedData.java | 74 ++-- .../compression/TopK/TopKCompressor.java | 382 +++++++++--------- .../federated/compression/TopK/TopKData.java | 60 +-- .../exceptions/CompressionException.java | 14 +- .../exceptions/DecompressionException.java | 14 +- ...obabilisticQuantizationCompressorTest.java | 276 ++++++------- .../federated/io/TopKCompressorTest.java | 231 +++++------ 13 files changed, 821 insertions(+), 902 deletions(-) diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressedMatrix.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressedMatrix.java index 6f349b2d7ab..4ad1660bfa2 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressedMatrix.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressedMatrix.java @@ -28,54 +28,54 @@ */ public class CompressedMatrix implements Serializable { - private static final long serialVersionUID = 1L; + private static final long serialVersionUID = 1L; - private final CompressionType type; - private final int numRows; - private final int numCols; - private final Object compressedData; // Technique-specific data - private final double compressionRatio; - private final byte[] metadata; // Optional: scaling factors, etc. + private final CompressionType type; + private final int numRows; + private final int numCols; + private final Object compressedData; // Technique-specific data + private final double compressionRatio; + private final byte[] metadata; // Optional: scaling factors, etc. - public CompressedMatrix(CompressionType type, int numRows, int numCols, - Object compressedData, double compressionRatio) { - this(type, numRows, numCols, compressedData, compressionRatio, null); - } + public CompressedMatrix(CompressionType type, int numRows, int numCols, + Object compressedData, double compressionRatio) { + this(type, numRows, numCols, compressedData, compressionRatio, null); + } - public CompressedMatrix(CompressionType type, int numRows, int numCols, - Object compressedData, double compressionRatio, - byte[] metadata) { - this.type = type; - this.numRows = numRows; - this.numCols = numCols; - this.compressedData = compressedData; - this.compressionRatio = compressionRatio; - this.metadata = metadata; - } + public CompressedMatrix(CompressionType type, int numRows, int numCols, + Object compressedData, double compressionRatio, + byte[] metadata) { + this.type = type; + this.numRows = numRows; + this.numCols = numCols; + this.compressedData = compressedData; + this.compressionRatio = compressionRatio; + this.metadata = metadata; + } - public CompressionType getType() { return type; } - public int getNumRows() { return numRows; } - public int getNumCols() { return numCols; } - public Object getCompressedData() { return compressedData; } - public double getCompressionRatio() { return compressionRatio; } - public byte[] getMetadata() { return metadata; } + public CompressionType getType() { return type; } + public int getNumRows() { return numRows; } + public int getNumCols() { return numCols; } + public Object getCompressedData() { return compressedData; } + public double getCompressionRatio() { return compressionRatio; } + public byte[] getMetadata() { return metadata; } - /** Estimate original size in bytes (8 bytes per double) */ - public long estimateOriginalSizeBytes() { - return (long) numRows * numCols * 8; - } + /** Estimate original size in bytes (8 bytes per double) */ + public long estimateOriginalSizeBytes() { + return (long) numRows * numCols * 8; + } - /** Estimate compressed size in bytes */ - public long getCompressedSizeBytes() { - if (compressedData instanceof byte[]) { - return ((byte[]) compressedData).length; - } - return 0; - } + /** Estimate compressed size in bytes */ + public long getCompressedSizeBytes() { + if (compressedData instanceof byte[]) { + return ((byte[]) compressedData).length; + } + return 0; + } - @Override - public String toString() { - return String.format("CompressedMatrix[%s, %dx%d, ratio=%.2fx]", - type.getId(), numRows, numCols, compressionRatio); - } + @Override + public String toString() { + return String.format("CompressedMatrix[%s, %dx%d, ratio=%.2fx]", + type.getId(), numRows, numCols, compressionRatio); + } } diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionConfig.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionConfig.java index 4f3bfa33554..0f825570c16 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionConfig.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionConfig.java @@ -35,76 +35,76 @@ */ public class CompressionConfig { - private final boolean enabled; - private final CompressionType type; - private final Map parameters; - - private CompressionConfig(Builder builder) { - this.enabled = builder.enabled; - this.type = builder.enabled ? builder.type : CompressionType.NONE; - this.parameters = new HashMap<>(builder.parameters); - } - - public boolean isEnabled() { return enabled; } - public CompressionType getType() { return type; } - public Map getParameters() { return new HashMap<>(parameters); } - - /** Convenience getter for sparsity parameter (TopK) */ - public double getSparsity() { - return (double) parameters.getOrDefault("sparsity", 0.01); - } - - /** Convenience getter for bits parameter (Quantization) */ - public int getBits() { - return (int) parameters.getOrDefault("bits", 4); - } - - @Override - public String toString() { - return String.format("CompressionConfig[enabled=%s, type=%s, params=%s]", - enabled, type.getId(), parameters); - } - - // ----------------------------------------------------------------------- - // Builder - // ----------------------------------------------------------------------- - - public static Builder builder() { - return new Builder(); - } - - public static class Builder { - private boolean enabled = false; - private CompressionType type = CompressionType.NONE; - private final Map parameters = new HashMap<>(); - - public Builder enable(boolean enabled) { - this.enabled = enabled; - return this; - } - - public Builder withType(CompressionType type) { - this.type = type; - return this; - } - - public Builder withParameter(String key, Object value) { - this.parameters.put(key, value); - return this; - } - - /** Shorthand for TopK sparsity ratio */ - public Builder withSparsity(double sparsity) { - return withParameter("sparsity", sparsity); - } - - /** Shorthand for quantization bit width */ - public Builder withBits(int bits) { - return withParameter("bits", bits); - } - - public CompressionConfig build() { - return new CompressionConfig(this); - } - } + private final boolean enabled; + private final CompressionType type; + private final Map parameters; + + private CompressionConfig(Builder builder) { + this.enabled = builder.enabled; + this.type = builder.enabled ? builder.type : CompressionType.NONE; + this.parameters = new HashMap<>(builder.parameters); + } + + public boolean isEnabled() { return enabled; } + public CompressionType getType() { return type; } + public Map getParameters() { return new HashMap<>(parameters); } + + /** Convenience getter for sparsity parameter (TopK) */ + public double getSparsity() { + return (double) parameters.getOrDefault("sparsity", 0.01); + } + + /** Convenience getter for bits parameter (Quantization) */ + public int getBits() { + return (int) parameters.getOrDefault("bits", 4); + } + + @Override + public String toString() { + return String.format("CompressionConfig[enabled=%s, type=%s, params=%s]", + enabled, type.getId(), parameters); + } + + // ----------------------------------------------------------------------- + // Builder + // ----------------------------------------------------------------------- + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + private boolean enabled = false; + private CompressionType type = CompressionType.NONE; + private final Map parameters = new HashMap<>(); + + public Builder enable(boolean enabled) { + this.enabled = enabled; + return this; + } + + public Builder withType(CompressionType type) { + this.type = type; + return this; + } + + public Builder withParameter(String key, Object value) { + this.parameters.put(key, value); + return this; + } + + /** Shorthand for TopK sparsity ratio */ + public Builder withSparsity(double sparsity) { + return withParameter("sparsity", sparsity); + } + + /** Shorthand for quantization bit width */ + public Builder withBits(int bits) { + return withParameter("bits", bits); + } + + public CompressionConfig build() { + return new CompressionConfig(this); + } + } } diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionFactory.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionFactory.java index bb4919415ee..0a6d1d0a3b2 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionFactory.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionFactory.java @@ -39,56 +39,56 @@ */ public class CompressionFactory { - private CompressionFactory() { - // Utility class — no instantiation - } + private CompressionFactory() { + // Utility class — no instantiation + } - /** - * Create a compressor from a CompressionConfig. - * @param config The compression configuration - * @return A ready-to-use MatrixCompressor - * @throws IllegalArgumentException if the config is invalid - */ - public static MatrixCompressor create(CompressionConfig config) { - if(config == null || !config.isEnabled()) - return new PassthroughCompressor(); - switch(config.getType()) { - case TOPK: - return new TopKCompressor(config.getSparsity(), true); - case PROBABILISTIC_QUANTIZATION: - return new ProbabilisticQuantizationCompressor(config.getBits()); - case NONE: - default: - return new PassthroughCompressor(); - } - } + /** + * Create a compressor from a CompressionConfig. + * @param config The compression configuration + * @return A ready-to-use MatrixCompressor + * @throws IllegalArgumentException if the config is invalid + */ + public static MatrixCompressor create(CompressionConfig config) { + if(config == null || !config.isEnabled()) + return new PassthroughCompressor(); + switch(config.getType()) { + case TOPK: + return new TopKCompressor(config.getSparsity(), true); + case PROBABILISTIC_QUANTIZATION: + return new ProbabilisticQuantizationCompressor(config.getBits()); + case NONE: + default: + return new PassthroughCompressor(); + } + } - // ----------------------------------------------------------------------- - // Passthrough compressor (no-op) for when compression is disabled - // ----------------------------------------------------------------------- + // ----------------------------------------------------------------------- + // Passthrough compressor (no-op) for when compression is disabled + // ----------------------------------------------------------------------- - /** - * No-op compressor: returns the matrix as-is. - * Used when compression is disabled or type is NONE. - */ - private static class PassthroughCompressor implements MatrixCompressor { + /** + * No-op compressor: returns the matrix as-is. + * Used when compression is disabled or type is NONE. + */ + private static class PassthroughCompressor implements MatrixCompressor { - @Override - public CompressedMatrix compress(MatrixBlock input) - throws CompressionException { - return new CompressedMatrix( - CompressionType.NONE, - input.getNumRows(), - input.getNumColumns(), - input, - 1.0 - ); - } + @Override + public CompressedMatrix compress(MatrixBlock input) + throws CompressionException { + return new CompressedMatrix( + CompressionType.NONE, + input.getNumRows(), + input.getNumColumns(), + input, + 1.0 + ); + } - @Override - public MatrixBlock decompress(CompressedMatrix compressed) - throws DecompressionException { - return (MatrixBlock) compressed.getCompressedData(); - } - } + @Override + public MatrixBlock decompress(CompressedMatrix compressed) + throws DecompressionException { + return (MatrixBlock) compressed.getCompressedData(); + } + } } diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionType.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionType.java index d942fa9e089..ce8359a7f5a 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionType.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionType.java @@ -25,33 +25,33 @@ */ public enum CompressionType { - /** TopK sparsification: keep largest-magnitude elements only */ - TOPK("topk", "Top-K Sparsification"), - - /** Probabilistic quantization: reduce precision with stochastic rounding */ - PROBABILISTIC_QUANTIZATION("prob_quant", "Probabilistic Quantization"), - - /** No compression (passthrough) */ - NONE("none", "No Compression"); - - private final String id; - private final String description; - - CompressionType(String id, String description) { - this.id = id; - this.description = description; - } - - public String getId() { return id; } - public String getDescription() { return description; } - - /** Parse from string identifier (case-insensitive) */ - public static CompressionType fromString(String text) { - for (CompressionType type : CompressionType.values()) { - if (type.id.equalsIgnoreCase(text)) { - return type; - } - } - throw new IllegalArgumentException("Unknown compression type: " + text); - } + /** TopK sparsification: keep largest-magnitude elements only */ + TOPK("topk", "Top-K Sparsification"), + + /** Probabilistic quantization: reduce precision with stochastic rounding */ + PROBABILISTIC_QUANTIZATION("prob_quant", "Probabilistic Quantization"), + + /** No compression (passthrough) */ + NONE("none", "No Compression"); + + private final String id; + private final String description; + + CompressionType(String id, String description) { + this.id = id; + this.description = description; + } + + public String getId() { return id; } + public String getDescription() { return description; } + + /** Parse from string identifier (case-insensitive) */ + public static CompressionType fromString(String text) { + for (CompressionType type : CompressionType.values()) { + if (type.id.equalsIgnoreCase(text)) { + return type; + } + } + throw new IllegalArgumentException("Unknown compression type: " + text); + } } diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/MatrixCompressor.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/MatrixCompressor.java index 12e43b3d7a6..301b245b0e3 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/MatrixCompressor.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/MatrixCompressor.java @@ -29,27 +29,27 @@ */ public interface MatrixCompressor { - /** - * Compress a matrix block for transmission. - * @param input The source matrix to compress - * @return CompressedMatrix containing compressed data and metadata - * @throws CompressionException if compression fails - */ - CompressedMatrix compress(MatrixBlock input) throws CompressionException; + /** + * Compress a matrix block for transmission. + * @param input The source matrix to compress + * @return CompressedMatrix containing compressed data and metadata + * @throws CompressionException if compression fails + */ + CompressedMatrix compress(MatrixBlock input) throws CompressionException; - /** - * Decompress a compressed matrix back to MatrixBlock. - * @param compressed The compressed data to decompress - * @return Reconstructed MatrixBlock (may be approximate) - * @throws DecompressionException if decompression fails - */ - MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException; + /** + * Decompress a compressed matrix back to MatrixBlock. + * @param compressed The compressed data to decompress + * @return Reconstructed MatrixBlock (may be approximate) + * @throws DecompressionException if decompression fails + */ + MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException; - /** - * Estimate the compression ratio achieved. - * Higher is better (e.g. 10.0 means 10x smaller). - */ - default double estimateCompressionRatio(long originalSize, long compressedSize) { - return compressedSize == 0 ? Double.MAX_VALUE : (double) originalSize / compressedSize; - } + /** + * Estimate the compression ratio achieved. + * Higher is better (e.g. 10.0 means 10x smaller). + */ + default double estimateCompressionRatio(long originalSize, long compressedSize) { + return compressedSize == 0 ? Double.MAX_VALUE : (double) originalSize / compressedSize; + } } diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/ProbabilisticQuantizationCompressor.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/ProbabilisticQuantizationCompressor.java index c46a9ffb874..7fcf9b9cf49 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/ProbabilisticQuantizationCompressor.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/ProbabilisticQuantizationCompressor.java @@ -42,128 +42,128 @@ */ public class ProbabilisticQuantizationCompressor implements MatrixCompressor { - private final int bitsPerValue; // 2, 4, or 8 - private final Random rng; - - public ProbabilisticQuantizationCompressor(int bitsPerValue) { - if (bitsPerValue != 2 && bitsPerValue != 4 && bitsPerValue != 8) { - throw new IllegalArgumentException("bitsPerValue must be 2, 4, or 8"); - } - this.bitsPerValue = bitsPerValue; - this.rng = new Random(42); // Fixed seed for reproducibility - } - - @Override - public CompressedMatrix compress(MatrixBlock input) throws CompressionException { - try { - int numRows = input.getNumRows(); - int numCols = input.getNumColumns(); - int totalElements = numRows * numCols; - - // Find min and max for normalization - double[] minMax = findMinMax(input, numRows, numCols); - double min = minMax[0]; - double max = minMax[1]; - - int levels = 1 << bitsPerValue; // 2^bits - - // Quantize each element probabilistically - byte[] quantized = new byte[totalElements]; - for (int i = 0; i < numRows; i++) { - for (int j = 0; j < numCols; j++) { - double value = input.get(i, j); - quantized[i * numCols + j] = probabilisticRound(value, min, max, levels); - } - } - - double ratio = 32.0 / bitsPerValue; // vs 32-bit float - - QuantizedData data = new QuantizedData( - quantized, min, max, levels, bitsPerValue, numRows, numCols); - - return new CompressedMatrix( - CompressionType.PROBABILISTIC_QUANTIZATION, - numRows, numCols, data, ratio); - - } catch (Exception e) { - throw new CompressionException( - "Quantization compression failed: " + e.getMessage(), e); - } - } - - @Override - public MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException { - try { - QuantizedData data = (QuantizedData) compressed.getCompressedData(); - MatrixBlock result = new MatrixBlock(data.numRows, data.numCols, false); - result.allocateDenseBlock(); - - for (int i = 0; i < data.numRows; i++) { - for (int j = 0; j < data.numCols; j++) { - byte levelIndex = data.quantizedValues[i * data.numCols + j]; - double value = data.reconstructValue(levelIndex); - result.set(i, j, value); - } - } - - result.examSparsity(); - return result; - - } catch (ClassCastException e) { - throw new DecompressionException( - "Invalid compressed data type for Quantization", e); - } catch (Exception e) { - throw new DecompressionException( - "Quantization decompression failed: " + e.getMessage(), e); - } - } - - // ----------------------------------------------------------------------- - // Private helpers - // ----------------------------------------------------------------------- - - /** - * Stochastic rounding: for value x between levels q_i and q_{i+1}: - * P(round up) = (x - q_i) / (q_{i+1} - q_i) - * P(round down) = 1 - P(round up) - * This gives E[output] = x (unbiased). - */ - private byte probabilisticRound(double value, double min, double max, int levels) { - // Handle constant matrix edge case - if (max - min < 1e-10) return 0; - - // Normalize to [0, 1] - double normalized = (value - min) / (max - min); - - // Find bounding level indices - double scaled = normalized * (levels - 1); - int lowerIdx = (int) scaled; - int upperIdx = lowerIdx + 1; - - if (lowerIdx == upperIdx) { - return (byte) lowerIdx; - } - - // Probabilistic decision - double probUp = scaled - lowerIdx; - return (rng.nextDouble() < probUp) ? (byte) upperIdx : (byte) lowerIdx; - } - - /** Find min and max values across the entire matrix */ - private double[] findMinMax(MatrixBlock input, int numRows, int numCols) { - double min = Double.MAX_VALUE; - double max = -Double.MAX_VALUE; - - for (int i = 0; i < numRows; i++) { - for (int j = 0; j < numCols; j++) { - double val = input.get(i, j); - if (val < min) min = val; - if (val > max) max = val; - } - } - - // Handle all-zero matrix - if (min == Double.MAX_VALUE) { min = 0; max = 0; } - return new double[]{min, max}; - } + private final int bitsPerValue; // 2, 4, or 8 + private final Random rng; + + public ProbabilisticQuantizationCompressor(int bitsPerValue) { + if (bitsPerValue != 2 && bitsPerValue != 4 && bitsPerValue != 8) { + throw new IllegalArgumentException("bitsPerValue must be 2, 4, or 8"); + } + this.bitsPerValue = bitsPerValue; + this.rng = new Random(42); // Fixed seed for reproducibility + } + + @Override + public CompressedMatrix compress(MatrixBlock input) throws CompressionException { + try { + int numRows = input.getNumRows(); + int numCols = input.getNumColumns(); + int totalElements = numRows * numCols; + + // Find min and max for normalization + double[] minMax = findMinMax(input, numRows, numCols); + double min = minMax[0]; + double max = minMax[1]; + + int levels = 1 << bitsPerValue; // 2^bits + + // Quantize each element probabilistically + byte[] quantized = new byte[totalElements]; + for (int i = 0; i < numRows; i++) { + for (int j = 0; j < numCols; j++) { + double value = input.get(i, j); + quantized[i * numCols + j] = probabilisticRound(value, min, max, levels); + } + } + + double ratio = 32.0 / bitsPerValue; // vs 32-bit float + + QuantizedData data = new QuantizedData( + quantized, min, max, levels, bitsPerValue, numRows, numCols); + + return new CompressedMatrix( + CompressionType.PROBABILISTIC_QUANTIZATION, + numRows, numCols, data, ratio); + + } catch (Exception e) { + throw new CompressionException( + "Quantization compression failed: " + e.getMessage(), e); + } + } + + @Override + public MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException { + try { + QuantizedData data = (QuantizedData) compressed.getCompressedData(); + MatrixBlock result = new MatrixBlock(data.numRows, data.numCols, false); + result.allocateDenseBlock(); + + for (int i = 0; i < data.numRows; i++) { + for (int j = 0; j < data.numCols; j++) { + byte levelIndex = data.quantizedValues[i * data.numCols + j]; + double value = data.reconstructValue(levelIndex); + result.set(i, j, value); + } + } + + result.examSparsity(); + return result; + + } catch (ClassCastException e) { + throw new DecompressionException( + "Invalid compressed data type for Quantization", e); + } catch (Exception e) { + throw new DecompressionException( + "Quantization decompression failed: " + e.getMessage(), e); + } + } + + // ----------------------------------------------------------------------- + // Private helpers + // ----------------------------------------------------------------------- + + /** + * Stochastic rounding: for value x between levels q_i and q_{i+1}: + * P(round up) = (x - q_i) / (q_{i+1} - q_i) + * P(round down) = 1 - P(round up) + * This gives E[output] = x (unbiased). + */ + private byte probabilisticRound(double value, double min, double max, int levels) { + // Handle constant matrix edge case + if (max - min < 1e-10) return 0; + + // Normalize to [0, 1] + double normalized = (value - min) / (max - min); + + // Find bounding level indices + double scaled = normalized * (levels - 1); + int lowerIdx = (int) scaled; + int upperIdx = lowerIdx + 1; + + if (lowerIdx == upperIdx) { + return (byte) lowerIdx; + } + + // Probabilistic decision + double probUp = scaled - lowerIdx; + return (rng.nextDouble() < probUp) ? (byte) upperIdx : (byte) lowerIdx; + } + + /** Find min and max values across the entire matrix */ + private double[] findMinMax(MatrixBlock input, int numRows, int numCols) { + double min = Double.MAX_VALUE; + double max = -Double.MAX_VALUE; + + for (int i = 0; i < numRows; i++) { + for (int j = 0; j < numCols; j++) { + double val = input.get(i, j); + if (val < min) min = val; + if (val > max) max = val; + } + } + + // Handle all-zero matrix + if (min == Double.MAX_VALUE) { min = 0; max = 0; } + return new double[]{min, max}; + } } diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/QuantizedData.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/QuantizedData.java index 595ae482556..f1b2c2cbbe0 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/QuantizedData.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/QuantizedData.java @@ -28,47 +28,47 @@ */ public class QuantizedData implements Serializable { - private static final long serialVersionUID = 1L; + private static final long serialVersionUID = 1L; - public final byte[] quantizedValues; // Quantized level indices - public final double min; // Original minimum value - public final double max; // Original maximum value - public final int levels; // Number of quantization levels (2^bits) - public final int bitsPerValue; // Bits used per element - public final int numRows; - public final int numCols; + public final byte[] quantizedValues; // Quantized level indices + public final double min; // Original minimum value + public final double max; // Original maximum value + public final int levels; // Number of quantization levels (2^bits) + public final int bitsPerValue; // Bits used per element + public final int numRows; + public final int numCols; - public QuantizedData(byte[] quantizedValues, double min, double max, - int levels, int bitsPerValue, int numRows, int numCols) { - this.quantizedValues = quantizedValues.clone(); // Defensive copy - this.min = min; - this.max = max; - this.levels = levels; - this.bitsPerValue = bitsPerValue; - this.numRows = numRows; - this.numCols = numCols; - } + public QuantizedData(byte[] quantizedValues, double min, double max, + int levels, int bitsPerValue, int numRows, int numCols) { + this.quantizedValues = quantizedValues.clone(); // Defensive copy + this.min = min; + this.max = max; + this.levels = levels; + this.bitsPerValue = bitsPerValue; + this.numRows = numRows; + this.numCols = numCols; + } - /** Number of quantized elements */ - public int size() { - return quantizedValues.length; - } + /** Number of quantized elements */ + public int size() { + return quantizedValues.length; + } - /** Estimate serialized size in bytes */ - public long estimateSizeBytes() { - return quantizedValues.length + 64; // +64 for scalar fields and headers - } + /** Estimate serialized size in bytes */ + public long estimateSizeBytes() { + return quantizedValues.length + 64; // +64 for scalar fields and headers + } - /** Reconstruct a double value from a quantized level index */ - public double reconstructValue(byte levelIndex) { - if (max - min < 1e-10) return min; // Constant matrix - int idx = levelIndex & 0xFF; // Treat byte as unsigned - return min + (idx / (double)(levels - 1)) * (max - min); - } + /** Reconstruct a double value from a quantized level index */ + public double reconstructValue(byte levelIndex) { + if (max - min < 1e-10) return min; // Constant matrix + int idx = levelIndex & 0xFF; // Treat byte as unsigned + return min + (idx / (double)(levels - 1)) * (max - min); + } - @Override - public String toString() { - return String.format("QuantizedData[%dx%d, levels=%d, bits=%d, min=%.4f, max=%.4f]", - numRows, numCols, levels, bitsPerValue, min, max); - } + @Override + public String toString() { + return String.format("QuantizedData[%dx%d, levels=%d, bits=%d, min=%.4f, max=%.4f]", + numRows, numCols, levels, bitsPerValue, min, max); + } } diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKCompressor.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKCompressor.java index f21834714f2..df3afc0b7f1 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKCompressor.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKCompressor.java @@ -43,195 +43,195 @@ */ public class TopKCompressor implements MatrixCompressor { - private final double sparsityRatio; // Fraction of elements to keep (0, 1] - private final boolean useHeap; // Use min-heap for O(n log k) selection - - /** - * @param sparsityRatio Fraction of elements to retain e.g. 0.01 = keep top 1% - * @param useHeap If true, use priority queue (faster for large matrices) - */ - public TopKCompressor(double sparsityRatio, boolean useHeap) { - if (sparsityRatio <= 0 || sparsityRatio > 1) { - throw new IllegalArgumentException("sparsityRatio must be in (0, 1]"); - } - this.sparsityRatio = sparsityRatio; - this.useHeap = useHeap; - } - - /** Default constructor: uses heap-based selection */ - public TopKCompressor(double sparsityRatio) { - this(sparsityRatio, true); - } - - @Override - public CompressedMatrix compress(MatrixBlock input) throws CompressionException { - try { - int numRows = input.getNumRows(); - int numCols = input.getNumColumns(); - int totalElements = numRows * numCols; - int k = (int) Math.max(1, Math.ceil(totalElements * sparsityRatio)); - - // If k covers everything, no compression needed - if (k >= totalElements) { - return new CompressedMatrix( - CompressionType.TOPK, numRows, numCols, input, 1.0); - } - - // Extract all non-zero elements with their linear indices - List elements = extractElements(input, numRows, numCols); - - // If fewer non-zeros than k, keep all of them - List topK = (elements.size() <= k) - ? new ArrayList<>(elements) - : selectTopK(elements, k); - - // Pack into TopKData - TopKData data = convertToTopKData(topK, numCols); - - double ratio = calculateCompressionRatio(totalElements, topK.size()); - - return new CompressedMatrix( - CompressionType.TOPK, numRows, numCols, data, ratio); - - } catch (Exception e) { - throw new CompressionException("TopK compression failed: " + e.getMessage(), e); - } - } - - @Override - public MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException { - try { - // Handle passthrough case (no compression was applied) - if (compressed.getCompressedData() instanceof MatrixBlock) { - return (MatrixBlock) compressed.getCompressedData(); - } - - TopKData data = (TopKData) compressed.getCompressedData(); - MatrixBlock result = new MatrixBlock( - compressed.getNumRows(), - compressed.getNumCols(), - true // Start sparse - ); - result.allocateSparseRowsBlock(); - - // Place values back at their original positions - for (int i = 0; i < data.indices.length; i++) { - int linearIdx = data.indices[i]; - int row = linearIdx / data.numCols; - int col = linearIdx % data.numCols; - result.set(row, col, data.values[i]); - } - - result.examSparsity(); - return result; - - } catch (ClassCastException e) { - throw new DecompressionException("Invalid compressed data type for TopK", e); - } catch (Exception e) { - throw new DecompressionException("TopK decompression failed: " + e.getMessage(), e); - } - } - - // ----------------------------------------------------------------------- - // Private helpers - // ----------------------------------------------------------------------- - - /** - * Extract all non-zero elements with their linear indices. - * Handles both dense and sparse MatrixBlock representations. - */ - private List extractElements(MatrixBlock input, int numRows, int numCols) { - List elements = new ArrayList<>(); - - if (input.isInSparseFormat()) { - // Sparse: iterate only over non-zero entries - for (int i = 0; i < numRows; i++) { - if (input.getSparseBlock() == null) continue; - if (input.getSparseBlock().isEmpty(i)) continue; - int[] rowIndices = input.getSparseBlock().indexes(i); - double[] rowValues = input.getSparseBlock().values(i); - int rowSize = input.getSparseBlock().size(i); - for (int j = 0; j < rowSize; j++) { - double val = rowValues[j]; - if (val != 0.0) { - int linearIdx = i * numCols + rowIndices[j]; - elements.add(new Element(linearIdx, val, Math.abs(val))); - } - } - } - } else { - // Dense: iterate all elements, skip zeros - double[] denseBlock = input.getDenseBlockValues(); - if (denseBlock != null) { - for (int i = 0; i < denseBlock.length; i++) { - if (denseBlock[i] != 0.0) { - elements.add(new Element(i, denseBlock[i], Math.abs(denseBlock[i]))); - } - } - } - } - return elements; - } - - /** - * Select top K elements by absolute value. - * Uses min-heap for O(n log k) when useHeap=true, - * or full sort O(n log n) otherwise. - */ - private List selectTopK(List elements, int k) { - if (useHeap) { - PriorityQueue minHeap = new PriorityQueue<>( - k, Comparator.comparingDouble(e -> e.absValue) - ); - for (Element e : elements) { - if (minHeap.size() < k) { - minHeap.offer(e); - } else if (e.absValue > minHeap.peek().absValue) { - minHeap.poll(); - minHeap.offer(e); - } - } - List result = new ArrayList<>(minHeap); - result.sort(Comparator.comparingInt(e -> e.index)); - return result; - } else { - elements.sort((a, b) -> Double.compare(b.absValue, a.absValue)); - return new ArrayList<>(elements.subList(0, k)); - } - } - - private TopKData convertToTopKData(List topK, int numCols) { - int[] indices = new int[topK.size()]; - double[] values = new double[topK.size()]; - for (int i = 0; i < topK.size(); i++) { - indices[i] = topK.get(i).index; - values[i] = topK.get(i).value; - } - return new TopKData(indices, values, numCols); - } - - private double calculateCompressionRatio(int total, int kept) { - if (kept == 0) return Double.MAX_VALUE; - // Original: total * 8 bytes (doubles) - // Compressed: kept * 12 bytes (int index + double value) - long originalBytes = (long) total * 8; - long compressedBytes = (long) kept * 12; - return (double) originalBytes / compressedBytes; - } - - // ----------------------------------------------------------------------- - // Inner class: element tracking during compression - // ----------------------------------------------------------------------- - - private static class Element { - final int index; // Linear index: row * numCols + col - final double value; // Original value - final double absValue; // Absolute value for magnitude comparison - - Element(int index, double value, double absValue) { - this.index = index; - this.value = value; - this.absValue = absValue; - } - } + private final double sparsityRatio; // Fraction of elements to keep (0, 1] + private final boolean useHeap; // Use min-heap for O(n log k) selection + + /** + * @param sparsityRatio Fraction of elements to retain e.g. 0.01 = keep top 1% + * @param useHeap If true, use priority queue (faster for large matrices) + */ + public TopKCompressor(double sparsityRatio, boolean useHeap) { + if (sparsityRatio <= 0 || sparsityRatio > 1) { + throw new IllegalArgumentException("sparsityRatio must be in (0, 1]"); + } + this.sparsityRatio = sparsityRatio; + this.useHeap = useHeap; + } + + /** Default constructor: uses heap-based selection */ + public TopKCompressor(double sparsityRatio) { + this(sparsityRatio, true); + } + + @Override + public CompressedMatrix compress(MatrixBlock input) throws CompressionException { + try { + int numRows = input.getNumRows(); + int numCols = input.getNumColumns(); + int totalElements = numRows * numCols; + int k = (int) Math.max(1, Math.ceil(totalElements * sparsityRatio)); + + // If k covers everything, no compression needed + if (k >= totalElements) { + return new CompressedMatrix( + CompressionType.TOPK, numRows, numCols, input, 1.0); + } + + // Extract all non-zero elements with their linear indices + List elements = extractElements(input, numRows, numCols); + + // If fewer non-zeros than k, keep all of them + List topK = (elements.size() <= k) + ? new ArrayList<>(elements) + : selectTopK(elements, k); + + // Pack into TopKData + TopKData data = convertToTopKData(topK, numCols); + + double ratio = calculateCompressionRatio(totalElements, topK.size()); + + return new CompressedMatrix( + CompressionType.TOPK, numRows, numCols, data, ratio); + + } catch (Exception e) { + throw new CompressionException("TopK compression failed: " + e.getMessage(), e); + } + } + + @Override + public MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException { + try { + // Handle passthrough case (no compression was applied) + if (compressed.getCompressedData() instanceof MatrixBlock) { + return (MatrixBlock) compressed.getCompressedData(); + } + + TopKData data = (TopKData) compressed.getCompressedData(); + MatrixBlock result = new MatrixBlock( + compressed.getNumRows(), + compressed.getNumCols(), + true // Start sparse + ); + result.allocateSparseRowsBlock(); + + // Place values back at their original positions + for (int i = 0; i < data.indices.length; i++) { + int linearIdx = data.indices[i]; + int row = linearIdx / data.numCols; + int col = linearIdx % data.numCols; + result.set(row, col, data.values[i]); + } + + result.examSparsity(); + return result; + + } catch (ClassCastException e) { + throw new DecompressionException("Invalid compressed data type for TopK", e); + } catch (Exception e) { + throw new DecompressionException("TopK decompression failed: " + e.getMessage(), e); + } + } + + // ----------------------------------------------------------------------- + // Private helpers + // ----------------------------------------------------------------------- + + /** + * Extract all non-zero elements with their linear indices. + * Handles both dense and sparse MatrixBlock representations. + */ + private List extractElements(MatrixBlock input, int numRows, int numCols) { + List elements = new ArrayList<>(); + + if (input.isInSparseFormat()) { + // Sparse: iterate only over non-zero entries + for (int i = 0; i < numRows; i++) { + if (input.getSparseBlock() == null) continue; + if (input.getSparseBlock().isEmpty(i)) continue; + int[] rowIndices = input.getSparseBlock().indexes(i); + double[] rowValues = input.getSparseBlock().values(i); + int rowSize = input.getSparseBlock().size(i); + for (int j = 0; j < rowSize; j++) { + double val = rowValues[j]; + if (val != 0.0) { + int linearIdx = i * numCols + rowIndices[j]; + elements.add(new Element(linearIdx, val, Math.abs(val))); + } + } + } + } else { + // Dense: iterate all elements, skip zeros + double[] denseBlock = input.getDenseBlockValues(); + if (denseBlock != null) { + for (int i = 0; i < denseBlock.length; i++) { + if (denseBlock[i] != 0.0) { + elements.add(new Element(i, denseBlock[i], Math.abs(denseBlock[i]))); + } + } + } + } + return elements; + } + + /** + * Select top K elements by absolute value. + * Uses min-heap for O(n log k) when useHeap=true, + * or full sort O(n log n) otherwise. + */ + private List selectTopK(List elements, int k) { + if (useHeap) { + PriorityQueue minHeap = new PriorityQueue<>( + k, Comparator.comparingDouble(e -> e.absValue) + ); + for (Element e : elements) { + if (minHeap.size() < k) { + minHeap.offer(e); + } else if (e.absValue > minHeap.peek().absValue) { + minHeap.poll(); + minHeap.offer(e); + } + } + List result = new ArrayList<>(minHeap); + result.sort(Comparator.comparingInt(e -> e.index)); + return result; + } else { + elements.sort((a, b) -> Double.compare(b.absValue, a.absValue)); + return new ArrayList<>(elements.subList(0, k)); + } + } + + private TopKData convertToTopKData(List topK, int numCols) { + int[] indices = new int[topK.size()]; + double[] values = new double[topK.size()]; + for (int i = 0; i < topK.size(); i++) { + indices[i] = topK.get(i).index; + values[i] = topK.get(i).value; + } + return new TopKData(indices, values, numCols); + } + + private double calculateCompressionRatio(int total, int kept) { + if (kept == 0) return Double.MAX_VALUE; + // Original: total * 8 bytes (doubles) + // Compressed: kept * 12 bytes (int index + double value) + long originalBytes = (long) total * 8; + long compressedBytes = (long) kept * 12; + return (double) originalBytes / compressedBytes; + } + + // ----------------------------------------------------------------------- + // Inner class: element tracking during compression + // ----------------------------------------------------------------------- + + private static class Element { + final int index; // Linear index: row * numCols + col + final double value; // Original value + final double absValue; // Absolute value for magnitude comparison + + Element(int index, double value, double absValue) { + this.index = index; + this.value = value; + this.absValue = absValue; + } + } } diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKData.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKData.java index eb737868c29..5d7c6f7fa8f 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKData.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKData.java @@ -28,34 +28,34 @@ */ public class TopKData implements Serializable { - private static final long serialVersionUID = 1L; - - public final int[] indices; // Linear indices of kept elements (row*numCols + col) - public final double[] values; // Corresponding original values - public final int numCols; // Needed for index → (row, col) conversion - - public TopKData(int[] indices, double[] values, int numCols) { - if (indices.length != values.length) { - throw new IllegalArgumentException( - "Indices and values arrays must have the same length"); - } - this.indices = indices.clone(); // Defensive copy - this.values = values.clone(); - this.numCols = numCols; - } - - /** Number of kept elements */ - public int size() { - return indices.length; - } - - /** Estimate serialized size in bytes (4 bytes per int + 8 bytes per double) */ - public long estimateSizeBytes() { - return (long) indices.length * 12 + 64; // +64 for object headers - } - - @Override - public String toString() { - return String.format("TopKData[k=%d, numCols=%d]", indices.length, numCols); - } + private static final long serialVersionUID = 1L; + + public final int[] indices; // Linear indices of kept elements (row*numCols + col) + public final double[] values; // Corresponding original values + public final int numCols; // Needed for index → (row, col) conversion + + public TopKData(int[] indices, double[] values, int numCols) { + if (indices.length != values.length) { + throw new IllegalArgumentException( + "Indices and values arrays must have the same length"); + } + this.indices = indices.clone(); // Defensive copy + this.values = values.clone(); + this.numCols = numCols; + } + + /** Number of kept elements */ + public int size() { + return indices.length; + } + + /** Estimate serialized size in bytes (4 bytes per int + 8 bytes per double) */ + public long estimateSizeBytes() { + return (long) indices.length * 12 + 64; // +64 for object headers + } + + @Override + public String toString() { + return String.format("TopKData[k=%d, numCols=%d]", indices.length, numCols); + } } diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/CompressionException.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/CompressionException.java index b56d44f23f0..24574d7fb4d 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/CompressionException.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/CompressionException.java @@ -24,13 +24,13 @@ */ public class CompressionException extends Exception { - private static final long serialVersionUID = 1L; + private static final long serialVersionUID = 1L; - public CompressionException(String message) { - super(message); - } + public CompressionException(String message) { + super(message); + } - public CompressionException(String message, Throwable cause) { - super(message, cause); - } + public CompressionException(String message, Throwable cause) { + super(message, cause); + } } diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/DecompressionException.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/DecompressionException.java index 7680d80e98e..7ba3601f107 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/DecompressionException.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/DecompressionException.java @@ -24,13 +24,13 @@ */ public class DecompressionException extends Exception { - private static final long serialVersionUID = 1L; + private static final long serialVersionUID = 1L; - public DecompressionException(String message) { - super(message); - } + public DecompressionException(String message) { + super(message); + } - public DecompressionException(String message, Throwable cause) { - super(message, cause); - } + public DecompressionException(String message, Throwable cause) { + super(message, cause); + } } diff --git a/src/test/java/org/apache/sysds/test/functions/federated/io/ProbabilisticQuantizationCompressorTest.java b/src/test/java/org/apache/sysds/test/functions/federated/io/ProbabilisticQuantizationCompressorTest.java index 17105385fb6..20f9c433c2d 100644 --- a/src/test/java/org/apache/sysds/test/functions/federated/io/ProbabilisticQuantizationCompressorTest.java +++ b/src/test/java/org/apache/sysds/test/functions/federated/io/ProbabilisticQuantizationCompressorTest.java @@ -16,171 +16,127 @@ * specific language governing permissions and limitations * under the License. */ - package org.apache.sysds.test.functions.federated.io; import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix; import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionType; import org.apache.sysds.runtime.controlprogram.federated.compression.Quantization.ProbabilisticQuantizationCompressor; import org.apache.sysds.runtime.matrix.data.MatrixBlock; -import org.apache.sysds.test.AutomatedTestBase; import org.junit.Test; -import static org.junit.Assert.*; - -public class ProbabilisticQuantizationCompressorTest extends AutomatedTestBase { - - @Override - public void setUp() {} - - @Override - public void tearDown() {} - - // ----------------------------------------------------------------------- - // Basic properties (merged from testCompressionTypeIsProbabilisticQuantization, - // testDimensionsPreservedAfterDecompression, testReconstructedValuesWithinOriginalRange, - // testCompressionRatio4Bit) - // ----------------------------------------------------------------------- - - @Test - public void testQuantizationBasicProperties() throws Exception { - MatrixBlock input = createRandomMatrix(10, 20); - ProbabilisticQuantizationCompressor compressor = - new ProbabilisticQuantizationCompressor(4); - CompressedMatrix compressed = compressor.compress(input); - - assertEquals(CompressionType.PROBABILISTIC_QUANTIZATION, compressed.getType()); - - MatrixBlock result = compressor.decompress(compressed); - assertEquals(10, result.getNumRows()); - assertEquals(20, result.getNumColumns()); - - // Find original min/max - double origMin = Double.MAX_VALUE; - double origMax = -Double.MAX_VALUE; - for (int i = 0; i < 10; i++) { - for (int j = 0; j < 20; j++) { - double v = input.get(i, j); - if (v < origMin) origMin = v; - if (v > origMax) origMax = v; - } - } - - // All reconstructed values must be within [min, max] - for (int i = 0; i < 10; i++) { - for (int j = 0; j < 20; j++) { - double v = result.get(i, j); - assertTrue("Value below min: " + v, v >= origMin - 1e-9); - assertTrue("Value above max: " + v, v <= origMax + 1e-9); - } - } - - assertEquals(8.0, compressed.getCompressionRatio(), 1e-10); - } - - // ----------------------------------------------------------------------- - // Compression ratio - // ----------------------------------------------------------------------- - - @Test - public void testCompressionRatio2Bit() throws Exception { - MatrixBlock input = createRandomMatrix(10, 10); - ProbabilisticQuantizationCompressor compressor = - new ProbabilisticQuantizationCompressor(2); - CompressedMatrix compressed = compressor.compress(input); - assertEquals(16.0, compressed.getCompressionRatio(), 1e-10); - } - - @Test - public void testCompressionRatio8Bit() throws Exception { - MatrixBlock input = createRandomMatrix(10, 10); - ProbabilisticQuantizationCompressor compressor = - new ProbabilisticQuantizationCompressor(8); - CompressedMatrix compressed = compressor.compress(input); - assertEquals(4.0, compressed.getCompressionRatio(), 1e-10); - } - - @Test - public void testFewerBitsGivesHigherRatio() throws Exception { - MatrixBlock input = createRandomMatrix(20, 20); - - double ratio2bit = new ProbabilisticQuantizationCompressor(2) - .compress(input).getCompressionRatio(); - double ratio8bit = new ProbabilisticQuantizationCompressor(8) - .compress(input).getCompressionRatio(); - - assertTrue("2-bit should compress more than 8-bit", ratio2bit > ratio8bit); - } - - // ----------------------------------------------------------------------- - // Unbiasedness: E[quantized] ≈ original over many runs - // ----------------------------------------------------------------------- - - @Test - public void testUnbiasednessOverManyRuns() throws Exception { - // Single element matrix with value 0.5 (midpoint) - MatrixBlock input = new MatrixBlock(1, 1, false); - input.allocateDenseBlock(); - input.set(0, 0, 5.0); - input.examSparsity(); - - // Run quantization 1000 times and average the results - double sum = 0.0; - int runs = 1000; - for (int r = 0; r < runs; r++) { - ProbabilisticQuantizationCompressor compressor = - new ProbabilisticQuantizationCompressor(4); - MatrixBlock result = compressor.decompress(compressor.compress(input)); - sum += result.get(0, 0); - } - double average = sum / runs; - - // Average should be close to original value (unbiased estimator) - assertEquals("Quantization should be unbiased", - 5.0, average, 0.5); // Allow 0.5 tolerance - } - - // ----------------------------------------------------------------------- - // Edge cases - // ----------------------------------------------------------------------- - - @Test - public void testConstantMatrix() throws Exception { - // All same value — min == max edge case - MatrixBlock input = new MatrixBlock(3, 3, false); - input.allocateDenseBlock(); - for (int i = 0; i < 3; i++) - for (int j = 0; j < 3; j++) - input.set(i, j, 7.0); - input.examSparsity(); - - ProbabilisticQuantizationCompressor compressor = - new ProbabilisticQuantizationCompressor(4); - CompressedMatrix compressed = compressor.compress(input); - MatrixBlock result = compressor.decompress(compressed); - - // Should not throw; all values should reconstruct to min (7.0 or close) - assertNotNull(result); - assertEquals(3, result.getNumRows()); - assertEquals(3, result.getNumColumns()); - } - - @Test(expected = IllegalArgumentException.class) - public void testInvalidBitsThrowsException() { - new ProbabilisticQuantizationCompressor(3); // Only 2, 4, 8 allowed - } - - // ----------------------------------------------------------------------- - // Helpers - // ----------------------------------------------------------------------- - - private MatrixBlock createRandomMatrix(int rows, int cols) { - MatrixBlock m = new MatrixBlock(rows, cols, false); - m.allocateDenseBlock(); - java.util.Random rng = new java.util.Random(42); - for (int i = 0; i < rows; i++) - for (int j = 0; j < cols; j++) - m.set(i, j, rng.nextGaussian() * 10); - m.examSparsity(); - return m; - } + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +public class ProbabilisticQuantizationCompressorTest { + + @Test + public void testQuantizationBasicProperties() throws Exception { + MatrixBlock input = createRandomMatrix(10, 20); + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(4); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + assertEquals(CompressionType.PROBABILISTIC_QUANTIZATION, compressed.getType()); + assertEquals(10, result.getNumRows()); + assertEquals(20, result.getNumColumns()); + assertEquals(8.0, compressed.getCompressionRatio(), 1e-10); + + double origMin = Double.MAX_VALUE; + double origMax = -Double.MAX_VALUE; + for(int i = 0; i < 10; i++) { + for(int j = 0; j < 20; j++) { + double v = input.get(i, j); + if(v < origMin) origMin = v; + if(v > origMax) origMax = v; + } + } + for(int i = 0; i < 10; i++) + for(int j = 0; j < 20; j++) { + double v = result.get(i, j); + assertTrue(v >= origMin - 1e-9); + assertTrue(v <= origMax + 1e-9); + } + } + + @Test + public void testFewerBitsGivesHigherRatio() throws Exception { + MatrixBlock input = createRandomMatrix(20, 20); + + double ratio2bit = new ProbabilisticQuantizationCompressor(2) + .compress(input).getCompressionRatio(); + double ratio8bit = new ProbabilisticQuantizationCompressor(8) + .compress(input).getCompressionRatio(); + + assertTrue(ratio2bit > ratio8bit); + } + + @Test + public void testUnbiasednessOverManyRuns() throws Exception { + MatrixBlock input = new MatrixBlock(1, 1, false); + input.allocateDenseBlock(); + input.set(0, 0, 5.0); + input.examSparsity(); + + double sum = 0.0; + int runs = 1000; + for(int r = 0; r < runs; r++) { + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(4); + MatrixBlock result = compressor.decompress(compressor.compress(input)); + sum += result.get(0, 0); + } + assertEquals(5.0, sum / runs, 0.5); + } + + @Test + public void testConstantMatrix() throws Exception { + MatrixBlock input = new MatrixBlock(3, 3, false); + input.allocateDenseBlock(); + for(int i = 0; i < 3; i++) + for(int j = 0; j < 3; j++) + input.set(i, j, 7.0); + input.examSparsity(); + + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(4); + MatrixBlock result = compressor.decompress(compressor.compress(input)); + + assertNotNull(result); + assertEquals(3, result.getNumRows()); + assertEquals(3, result.getNumColumns()); + } + + @Test(expected = IllegalArgumentException.class) + public void testInvalidBitsThrowsException() { + new ProbabilisticQuantizationCompressor(3); + } + + @Test + public void testCompressionRatio2Bit() throws Exception { + MatrixBlock input = createRandomMatrix(10, 10); + assertEquals(16.0, + new ProbabilisticQuantizationCompressor(2).compress(input).getCompressionRatio(), + 1e-10); + } + + @Test + public void testCompressionRatio8Bit() throws Exception { + MatrixBlock input = createRandomMatrix(10, 10); + assertEquals(4.0, + new ProbabilisticQuantizationCompressor(8).compress(input).getCompressionRatio(), + 1e-10); + } + + private MatrixBlock createRandomMatrix(int rows, int cols) { + MatrixBlock m = new MatrixBlock(rows, cols, false); + m.allocateDenseBlock(); + java.util.Random rng = new java.util.Random(42); + for(int i = 0; i < rows; i++) + for(int j = 0; j < cols; j++) + m.set(i, j, rng.nextGaussian() * 10); + m.examSparsity(); + return m; + } } diff --git a/src/test/java/org/apache/sysds/test/functions/federated/io/TopKCompressorTest.java b/src/test/java/org/apache/sysds/test/functions/federated/io/TopKCompressorTest.java index 65b16c9e429..3e07d802c05 100644 --- a/src/test/java/org/apache/sysds/test/functions/federated/io/TopKCompressorTest.java +++ b/src/test/java/org/apache/sysds/test/functions/federated/io/TopKCompressorTest.java @@ -16,145 +16,108 @@ * specific language governing permissions and limitations * under the License. */ - package org.apache.sysds.test.functions.federated.io; import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix; import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionType; import org.apache.sysds.runtime.controlprogram.federated.compression.TopK.TopKCompressor; import org.apache.sysds.runtime.matrix.data.MatrixBlock; -import org.apache.sysds.test.AutomatedTestBase; import org.junit.Test; -import static org.junit.Assert.*; - -public class TopKCompressorTest extends AutomatedTestBase { - - @Override - public void setUp() {} - - @Override - public void tearDown() {} - - // ----------------------------------------------------------------------- - // Basic properties (merged from testCompressionTypeIsTopK, - // testDimensionsPreservedAfterDecompression, testCompressionRatioIsPositive) - // ----------------------------------------------------------------------- - - @Test - public void testTopKBasicProperties() throws Exception { - MatrixBlock input = createRandomMatrix(10, 20); - TopKCompressor compressor = new TopKCompressor(0.1); - CompressedMatrix compressed = compressor.compress(input); - - assertEquals(CompressionType.TOPK, compressed.getType()); - - MatrixBlock result = compressor.decompress(compressed); - assertEquals(10, result.getNumRows()); - assertEquals(20, result.getNumColumns()); - - assertTrue("Compression ratio must be > 0", - compressed.getCompressionRatio() > 0); - } - - // ----------------------------------------------------------------------- - // Basic compression / decompression - // ----------------------------------------------------------------------- - - @Test - public void testTopKKeepsLargestElements() throws Exception { - // 3x3 matrix with three distinct magnitudes - MatrixBlock input = new MatrixBlock(3, 3, false); - input.allocateDenseBlock(); - input.set(0, 0, 10.0); // Largest - input.set(1, 1, 5.0); // Medium - input.set(2, 2, 1.0); // Smallest - input.examSparsity(); - - // Keep top 2 of 9 elements (~22% sparsity) - TopKCompressor compressor = new TopKCompressor(0.22); - CompressedMatrix compressed = compressor.compress(input); - MatrixBlock result = compressor.decompress(compressed); - - // Largest two values must be preserved exactly - assertEquals(10.0, result.get(0, 0), 1e-10); - assertEquals(5.0, result.get(1, 1), 1e-10); - - // Smallest should be zeroed out - assertEquals(0.0, result.get(2, 2), 1e-10); - } - - // ----------------------------------------------------------------------- - // Compression ratio - // ----------------------------------------------------------------------- - - @Test - public void testLowerSparsityGivesHigherRatio() throws Exception { - MatrixBlock input = createRandomMatrix(100, 100); - - TopKCompressor c1 = new TopKCompressor(0.1); - TopKCompressor c2 = new TopKCompressor(0.01); - - double ratio1 = c1.compress(input).getCompressionRatio(); - double ratio2 = c2.compress(input).getCompressionRatio(); - - assertTrue("1% sparsity should compress more than 10%", ratio2 > ratio1); - } - - // ----------------------------------------------------------------------- - // Edge cases - // ----------------------------------------------------------------------- - - @Test - public void testAllZeroMatrix() throws Exception { - MatrixBlock input = new MatrixBlock(5, 5, false); - input.allocateDenseBlock(); - input.examSparsity(); - - TopKCompressor compressor = new TopKCompressor(0.1); - CompressedMatrix compressed = compressor.compress(input); - MatrixBlock result = compressor.decompress(compressed); - - // All zeros in → all zeros out - for (int i = 0; i < 5; i++) - for (int j = 0; j < 5; j++) - assertEquals(0.0, result.get(i, j), 1e-10); - } - - @Test - public void testSparsityOfOneKeepsEverything() throws Exception { - MatrixBlock input = createRandomMatrix(5, 5); - TopKCompressor compressor = new TopKCompressor(1.0); - CompressedMatrix compressed = compressor.compress(input); - MatrixBlock result = compressor.decompress(compressed); - - // With sparsity=1.0, all values should be preserved - for (int i = 0; i < 5; i++) - for (int j = 0; j < 5; j++) - assertEquals(input.get(i, j), result.get(i, j), 1e-10); - } - - @Test(expected = IllegalArgumentException.class) - public void testInvalidSparsityThrowsException() { - new TopKCompressor(0.0); // Must be > 0 - } - - @Test(expected = IllegalArgumentException.class) - public void testSparsityAboveOneThrowsException() { - new TopKCompressor(1.5); // Must be <= 1 - } - - // ----------------------------------------------------------------------- - // Helpers - // ----------------------------------------------------------------------- - - private MatrixBlock createRandomMatrix(int rows, int cols) { - MatrixBlock m = new MatrixBlock(rows, cols, false); - m.allocateDenseBlock(); - java.util.Random rng = new java.util.Random(42); - for (int i = 0; i < rows; i++) - for (int j = 0; j < cols; j++) - m.set(i, j, rng.nextGaussian() * 10); - m.examSparsity(); - return m; - } + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +public class TopKCompressorTest { + + @Test + public void testTopKBasicProperties() throws Exception { + MatrixBlock input = createRandomMatrix(10, 20); + TopKCompressor compressor = new TopKCompressor(0.1); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + assertEquals(CompressionType.TOPK, compressed.getType()); + assertEquals(10, result.getNumRows()); + assertEquals(20, result.getNumColumns()); + assertTrue(compressed.getCompressionRatio() > 0); + } + + @Test + public void testTopKKeepsLargestElements() throws Exception { + MatrixBlock input = new MatrixBlock(3, 3, false); + input.allocateDenseBlock(); + input.set(0, 0, 10.0); + input.set(1, 1, 5.0); + input.set(2, 2, 1.0); + input.examSparsity(); + + TopKCompressor compressor = new TopKCompressor(0.22); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + assertEquals(10.0, result.get(0, 0), 1e-10); + assertEquals(5.0, result.get(1, 1), 1e-10); + assertEquals(0.0, result.get(2, 2), 1e-10); + } + + @Test + public void testLowerSparsityGivesHigherRatio() throws Exception { + MatrixBlock input = createRandomMatrix(100, 100); + + TopKCompressor c1 = new TopKCompressor(0.1); + TopKCompressor c2 = new TopKCompressor(0.01); + + double ratio1 = c1.compress(input).getCompressionRatio(); + double ratio2 = c2.compress(input).getCompressionRatio(); + + assertTrue(ratio2 > ratio1); + } + + @Test + public void testAllZeroMatrix() throws Exception { + MatrixBlock input = new MatrixBlock(5, 5, false); + input.allocateDenseBlock(); + input.examSparsity(); + + TopKCompressor compressor = new TopKCompressor(0.1); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + for(int i = 0; i < 5; i++) + for(int j = 0; j < 5; j++) + assertEquals(0.0, result.get(i, j), 1e-10); + } + + @Test + public void testSparsityOfOneKeepsEverything() throws Exception { + MatrixBlock input = createRandomMatrix(5, 5); + TopKCompressor compressor = new TopKCompressor(1.0); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + for(int i = 0; i < 5; i++) + for(int j = 0; j < 5; j++) + assertEquals(input.get(i, j), result.get(i, j), 1e-10); + } + + @Test(expected = IllegalArgumentException.class) + public void testInvalidSparsityThrowsException() { + new TopKCompressor(0.0); + } + + @Test(expected = IllegalArgumentException.class) + public void testSparsityAboveOneThrowsException() { + new TopKCompressor(1.5); + } + + private MatrixBlock createRandomMatrix(int rows, int cols) { + MatrixBlock m = new MatrixBlock(rows, cols, false); + m.allocateDenseBlock(); + java.util.Random rng = new java.util.Random(42); + for(int i = 0; i < rows; i++) + for(int j = 0; j < cols; j++) + m.set(i, j, rng.nextGaussian() * 10); + m.examSparsity(); + return m; + } } From aafc6023adcfaabef2d3406e6dd80b66005bfbb1 Mon Sep 17 00:00:00 2001 From: Charansurya Udaysingh Jhurree Date: Fri, 3 Jul 2026 10:35:28 +0200 Subject: [PATCH 19/19] Remove old compress/ package files (moved to federated/compression/) --- .../runtime/compress/CompressedMatrix.java | 64 ----- .../runtime/compress/CompressionConfig.java | 93 ------- .../runtime/compress/CompressionFactory.java | 95 -------- .../runtime/compress/CompressionType.java | 43 ---- .../runtime/compress/MatrixCompressor.java | 44 ---- .../ProbabilisticQuantizationCompressor.java | 158 ------------ .../compress/Quantization/QuantizedData.java | 57 ----- .../runtime/compress/TopK/TopKCompressor.java | 226 ------------------ .../sysds/runtime/compress/TopK/TopKData.java | 44 ---- .../exceptions/CompressionException.java | 19 -- .../exceptions/DecompressionException.java | 19 -- 11 files changed, 862 deletions(-) delete mode 100644 src/main/java/org/apache/sysds/runtime/compress/CompressedMatrix.java delete mode 100644 src/main/java/org/apache/sysds/runtime/compress/CompressionConfig.java delete mode 100644 src/main/java/org/apache/sysds/runtime/compress/CompressionFactory.java delete mode 100644 src/main/java/org/apache/sysds/runtime/compress/CompressionType.java delete mode 100644 src/main/java/org/apache/sysds/runtime/compress/MatrixCompressor.java delete mode 100644 src/main/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressor.java delete mode 100644 src/main/java/org/apache/sysds/runtime/compress/Quantization/QuantizedData.java delete mode 100644 src/main/java/org/apache/sysds/runtime/compress/TopK/TopKCompressor.java delete mode 100644 src/main/java/org/apache/sysds/runtime/compress/TopK/TopKData.java delete mode 100644 src/main/java/org/apache/sysds/runtime/compress/exceptions/CompressionException.java delete mode 100644 src/main/java/org/apache/sysds/runtime/compress/exceptions/DecompressionException.java diff --git a/src/main/java/org/apache/sysds/runtime/compress/CompressedMatrix.java b/src/main/java/org/apache/sysds/runtime/compress/CompressedMatrix.java deleted file mode 100644 index a24b5255d1c..00000000000 --- a/src/main/java/org/apache/sysds/runtime/compress/CompressedMatrix.java +++ /dev/null @@ -1,64 +0,0 @@ -package org.apache.sysds.runtime.compress; - -import java.io.Serializable; - -/** - * Generic container for compressed matrix data. - * Stores the compressed representation along with metadata - * needed for decompression and size estimation. - * - * @author Nirvan C. Udaysingh Jhurree - */ -public class CompressedMatrix implements Serializable { - - private static final long serialVersionUID = 1L; - - private final CompressionType type; - private final int numRows; - private final int numCols; - private final Object compressedData; // Technique-specific data - private final double compressionRatio; - private final byte[] metadata; // Optional: scaling factors, etc. - - public CompressedMatrix(CompressionType type, int numRows, int numCols, - Object compressedData, double compressionRatio) { - this(type, numRows, numCols, compressedData, compressionRatio, null); - } - - public CompressedMatrix(CompressionType type, int numRows, int numCols, - Object compressedData, double compressionRatio, - byte[] metadata) { - this.type = type; - this.numRows = numRows; - this.numCols = numCols; - this.compressedData = compressedData; - this.compressionRatio = compressionRatio; - this.metadata = metadata; - } - - public CompressionType getType() { return type; } - public int getNumRows() { return numRows; } - public int getNumCols() { return numCols; } - public Object getCompressedData() { return compressedData; } - public double getCompressionRatio() { return compressionRatio; } - public byte[] getMetadata() { return metadata; } - - /** Estimate original size in bytes (8 bytes per double) */ - public long estimateOriginalSizeBytes() { - return (long) numRows * numCols * 8; - } - - /** Estimate compressed size in bytes */ - public long getCompressedSizeBytes() { - if (compressedData instanceof byte[]) { - return ((byte[]) compressedData).length; - } - return 0; - } - - @Override - public String toString() { - return String.format("CompressedMatrix[%s, %dx%d, ratio=%.2fx]", - type.getId(), numRows, numCols, compressionRatio); - } -} \ No newline at end of file diff --git a/src/main/java/org/apache/sysds/runtime/compress/CompressionConfig.java b/src/main/java/org/apache/sysds/runtime/compress/CompressionConfig.java deleted file mode 100644 index d2634bda528..00000000000 --- a/src/main/java/org/apache/sysds/runtime/compress/CompressionConfig.java +++ /dev/null @@ -1,93 +0,0 @@ -package org.apache.sysds.runtime.compress; - -import java.util.HashMap; -import java.util.Map; - -/** - * Immutable configuration for compression in federated operations. - * Uses the Builder pattern for flexible, readable configuration. - * - * Usage example: - * CompressionConfig config = CompressionConfig.builder() - * .enable(true) - * .withType(CompressionType.TOPK) - * .withSparsity(0.01) - * .build(); - * - * - */ -public class CompressionConfig { - - private final boolean enabled; - private final CompressionType type; - private final Map parameters; - - private CompressionConfig(Builder builder) { - this.enabled = builder.enabled; - this.type = builder.enabled ? builder.type : CompressionType.NONE; - this.parameters = new HashMap<>(builder.parameters); - } - - public boolean isEnabled() { return enabled; } - public CompressionType getType() { return type; } - public Map getParameters() { return new HashMap<>(parameters); } - - /** Convenience getter for sparsity parameter (TopK) */ - public double getSparsity() { - return (double) parameters.getOrDefault("sparsity", 0.01); - } - - /** Convenience getter for bits parameter (Quantization) */ - public int getBits() { - return (int) parameters.getOrDefault("bits", 4); - } - - @Override - public String toString() { - return String.format("CompressionConfig[enabled=%s, type=%s, params=%s]", - enabled, type.getId(), parameters); - } - - // ----------------------------------------------------------------------- - // Builder - // ----------------------------------------------------------------------- - - public static Builder builder() { - return new Builder(); - } - - public static class Builder { - private boolean enabled = false; - private CompressionType type = CompressionType.NONE; - private final Map parameters = new HashMap<>(); - - public Builder enable(boolean enabled) { - this.enabled = enabled; - return this; - } - - public Builder withType(CompressionType type) { - this.type = type; - return this; - } - - public Builder withParameter(String key, Object value) { - this.parameters.put(key, value); - return this; - } - - /** Shorthand for TopK sparsity ratio */ - public Builder withSparsity(double sparsity) { - return withParameter("sparsity", sparsity); - } - - /** Shorthand for quantization bit width */ - public Builder withBits(int bits) { - return withParameter("bits", bits); - } - - public CompressionConfig build() { - return new CompressionConfig(this); - } - } -} \ No newline at end of file diff --git a/src/main/java/org/apache/sysds/runtime/compress/CompressionFactory.java b/src/main/java/org/apache/sysds/runtime/compress/CompressionFactory.java deleted file mode 100644 index 04c7a982fc2..00000000000 --- a/src/main/java/org/apache/sysds/runtime/compress/CompressionFactory.java +++ /dev/null @@ -1,95 +0,0 @@ -package org.apache.sysds.runtime.compress; - -import org.apache.sysds.runtime.compress.TopK.TopKCompressor; -import org.apache.sysds.runtime.compress.Quantization.ProbabilisticQuantizationCompressor; - -/** - * Factory for creating compressor instances from configuration. - * Centralizes compressor instantiation and parameter validation. - * - * Usage: - * CompressionConfig config = CompressionConfig.builder() - * .enable(true) - * .withType(CompressionType.TOPK) - * .withSparsity(0.01) - * .build(); - * MatrixCompressor compressor = CompressionFactory.create(config); - * - * - */ -public class CompressionFactory { - - private CompressionFactory() { - // Utility class — no instantiation - } - - /** - * Create a compressor from a CompressionConfig. - * @param config The compression configuration - * @return A ready-to-use MatrixCompressor - * @throws IllegalArgumentException if the config is invalid - */ - public static MatrixCompressor create(CompressionConfig config) { - if (config == null || !config.isEnabled()) { - return new PassthroughCompressor(); - } - return create(config.getType(), config); - } - - /** - * Create a compressor for a specific type with given config. - */ - public static MatrixCompressor create(CompressionType type, CompressionConfig config) { - switch (type) { - case TOPK: - double sparsity = config.getSparsity(); - return new TopKCompressor(sparsity, true); - - case PROBABILISTIC_QUANTIZATION: - int bits = config.getBits(); - return new ProbabilisticQuantizationCompressor(bits); - - case ONE_BIT_CS: - throw new UnsupportedOperationException( - "1-Bit Compressed Sensing not yet implemented"); - - case NONE: - default: - return new PassthroughCompressor(); - } - } - - // ----------------------------------------------------------------------- - // Passthrough compressor (no-op) for when compression is disabled - // ----------------------------------------------------------------------- - - /** - * No-op compressor: returns the matrix as-is. - * Used when compression is disabled or type is NONE. - */ - private static class PassthroughCompressor implements MatrixCompressor { - - @Override - public CompressedMatrix compress(org.apache.sysds.runtime.matrix.data.MatrixBlock input) - throws org.apache.sysds.runtime.compress.exceptions.CompressionException { - return new CompressedMatrix( - CompressionType.NONE, - input.getNumRows(), - input.getNumColumns(), - input, - 1.0 - ); - } - - @Override - public org.apache.sysds.runtime.matrix.data.MatrixBlock decompress(CompressedMatrix compressed) - throws org.apache.sysds.runtime.compress.exceptions.DecompressionException { - return (org.apache.sysds.runtime.matrix.data.MatrixBlock) compressed.getCompressedData(); - } - - @Override - public CompressionType getCompressionType() { - return CompressionType.NONE; - } - } -} \ No newline at end of file diff --git a/src/main/java/org/apache/sysds/runtime/compress/CompressionType.java b/src/main/java/org/apache/sysds/runtime/compress/CompressionType.java deleted file mode 100644 index ba6a7b9851d..00000000000 --- a/src/main/java/org/apache/sysds/runtime/compress/CompressionType.java +++ /dev/null @@ -1,43 +0,0 @@ -package org.apache.sysds.runtime.compress; - -/** - * Enumeration of supported compression techniques for federated learning. - * Used for configuration, serialization, and technique selection. - * - * @author Nirvan C. Udaysingh Jhurree - */ -public enum CompressionType { - - /** TopK sparsification: keep largest-magnitude elements only */ - TOPK("topk", "Top-K Sparsification"), - - /** Probabilistic quantization: reduce precision with stochastic rounding */ - PROBABILISTIC_QUANTIZATION("prob_quant", "Probabilistic Quantization"), - - /** 1-bit compressed sensing: sign-only transmission + iterative reconstruction */ - ONE_BIT_CS("1bit_cs", "1-Bit Compressed Sensing"), - - /** No compression (passthrough) */ - NONE("none", "No Compression"); - - private final String id; - private final String description; - - CompressionType(String id, String description) { - this.id = id; - this.description = description; - } - - public String getId() { return id; } - public String getDescription() { return description; } - - /** Parse from string identifier (case-insensitive) */ - public static CompressionType fromString(String text) { - for (CompressionType type : CompressionType.values()) { - if (type.id.equalsIgnoreCase(text)) { - return type; - } - } - throw new IllegalArgumentException("Unknown compression type: " + text); - } -} \ No newline at end of file diff --git a/src/main/java/org/apache/sysds/runtime/compress/MatrixCompressor.java b/src/main/java/org/apache/sysds/runtime/compress/MatrixCompressor.java deleted file mode 100644 index 88d235e9d2f..00000000000 --- a/src/main/java/org/apache/sysds/runtime/compress/MatrixCompressor.java +++ /dev/null @@ -1,44 +0,0 @@ -package org.apache.sysds.runtime.compress; - -import org.apache.sysds.runtime.compress.exceptions.CompressionException; -import org.apache.sysds.runtime.compress.exceptions.DecompressionException; -import org.apache.sysds.runtime.matrix.data.MatrixBlock; - -/** - * Interface for matrix compression techniques in federated learning. - * All compressors must implement compress/decompress operations. - * - * @author Nirvan C. Udaysingh Jhurree - */ -public interface MatrixCompressor { - - /** - * Compress a matrix block for transmission. - * @param input The source matrix to compress - * @return CompressedMatrix containing compressed data and metadata - * @throws CompressionException if compression fails - */ - CompressedMatrix compress(MatrixBlock input) throws CompressionException; - - /** - * Decompress a compressed matrix back to MatrixBlock. - * @param compressed The compressed data to decompress - * @return Reconstructed MatrixBlock (may be approximate) - * @throws DecompressionException if decompression fails - */ - MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException; - - /** - * Get the compression technique identifier. - * @return CompressionType enum value - */ - CompressionType getCompressionType(); - - /** - * Estimate the compression ratio achieved. - * Higher is better (e.g. 10.0 means 10x smaller). - */ - default double estimateCompressionRatio(long originalSize, long compressedSize) { - return compressedSize == 0 ? Double.MAX_VALUE : (double) originalSize / compressedSize; - } -} \ No newline at end of file diff --git a/src/main/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressor.java b/src/main/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressor.java deleted file mode 100644 index 95ce2562b09..00000000000 --- a/src/main/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressor.java +++ /dev/null @@ -1,158 +0,0 @@ -package org.apache.sysds.runtime.compress.Quantization; - -import org.apache.sysds.runtime.compress.CompressedMatrix; -import org.apache.sysds.runtime.compress.CompressionType; -import org.apache.sysds.runtime.compress.MatrixCompressor; -import org.apache.sysds.runtime.compress.exceptions.CompressionException; -import org.apache.sysds.runtime.compress.exceptions.DecompressionException; -import org.apache.sysds.runtime.matrix.data.MatrixBlock; - -import java.util.Random; - -/** - * Probabilistic Quantization Compressor. - * - * Reduces numerical precision using stochastic rounding to maintain - * an unbiased estimator — meaning E[quantized] = original on average. - * This is critical for federated learning convergence guarantees. - * - * Supports 2, 4, or 8 bits per value: - * 2-bit → 4 levels → 16x compression vs 32-bit float - * 4-bit → 16 levels → 8x compression - * 8-bit → 256 levels → 4x compression - * - * - */ -public class ProbabilisticQuantizationCompressor implements MatrixCompressor { - - private final int bitsPerValue; // 2, 4, or 8 - private final Random rng; - - public ProbabilisticQuantizationCompressor(int bitsPerValue) { - if (bitsPerValue != 2 && bitsPerValue != 4 && bitsPerValue != 8) { - throw new IllegalArgumentException("bitsPerValue must be 2, 4, or 8"); - } - this.bitsPerValue = bitsPerValue; - this.rng = new Random(42); // Fixed seed for reproducibility - } - - @Override - public CompressedMatrix compress(MatrixBlock input) throws CompressionException { - try { - int numRows = input.getNumRows(); - int numCols = input.getNumColumns(); - int totalElements = numRows * numCols; - - // Find min and max for normalization - double[] minMax = findMinMax(input, numRows, numCols); - double min = minMax[0]; - double max = minMax[1]; - - int levels = 1 << bitsPerValue; // 2^bits - - // Quantize each element probabilistically - byte[] quantized = new byte[totalElements]; - for (int i = 0; i < numRows; i++) { - for (int j = 0; j < numCols; j++) { - double value = input.get(i, j); - quantized[i * numCols + j] = probabilisticRound(value, min, max, levels); - } - } - - double ratio = 32.0 / bitsPerValue; // vs 32-bit float - - QuantizedData data = new QuantizedData( - quantized, min, max, levels, bitsPerValue, numRows, numCols); - - return new CompressedMatrix( - CompressionType.PROBABILISTIC_QUANTIZATION, - numRows, numCols, data, ratio); - - } catch (Exception e) { - throw new CompressionException( - "Quantization compression failed: " + e.getMessage(), e); - } - } - - @Override - public MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException { - try { - QuantizedData data = (QuantizedData) compressed.getCompressedData(); - MatrixBlock result = new MatrixBlock(data.numRows, data.numCols, false); - result.allocateDenseBlock(); - - for (int i = 0; i < data.numRows; i++) { - for (int j = 0; j < data.numCols; j++) { - byte levelIndex = data.quantizedValues[i * data.numCols + j]; - double value = data.reconstructValue(levelIndex); - result.set(i, j, value); - } - } - - result.examSparsity(); - return result; - - } catch (ClassCastException e) { - throw new DecompressionException( - "Invalid compressed data type for Quantization", e); - } catch (Exception e) { - throw new DecompressionException( - "Quantization decompression failed: " + e.getMessage(), e); - } - } - - @Override - public CompressionType getCompressionType() { - return CompressionType.PROBABILISTIC_QUANTIZATION; - } - - // ----------------------------------------------------------------------- - // Private helpers - // ----------------------------------------------------------------------- - - /** - * Stochastic rounding: for value x between levels q_i and q_{i+1}: - * P(round up) = (x - q_i) / (q_{i+1} - q_i) - * P(round down) = 1 - P(round up) - * This gives E[output] = x (unbiased). - */ - private byte probabilisticRound(double value, double min, double max, int levels) { - // Handle constant matrix edge case - if (max - min < 1e-10) return 0; - - // Normalize to [0, 1] - double normalized = (value - min) / (max - min); - normalized = Math.max(0.0, Math.min(1.0, normalized)); // Clamp - - // Find bounding level indices - double scaled = normalized * (levels - 1); - int lowerIdx = (int) scaled; - int upperIdx = Math.min(lowerIdx + 1, levels - 1); - - if (lowerIdx == upperIdx) { - return (byte) lowerIdx; - } - - // Probabilistic decision - double probUp = scaled - lowerIdx; - return (rng.nextDouble() < probUp) ? (byte) upperIdx : (byte) lowerIdx; - } - - /** Find min and max values across the entire matrix */ - private double[] findMinMax(MatrixBlock input, int numRows, int numCols) { - double min = Double.MAX_VALUE; - double max = -Double.MAX_VALUE; - - for (int i = 0; i < numRows; i++) { - for (int j = 0; j < numCols; j++) { - double val = input.get(i, j); - if (val < min) min = val; - if (val > max) max = val; - } - } - - // Handle all-zero matrix - if (min == Double.MAX_VALUE) { min = 0; max = 0; } - return new double[]{min, max}; - } -} \ No newline at end of file diff --git a/src/main/java/org/apache/sysds/runtime/compress/Quantization/QuantizedData.java b/src/main/java/org/apache/sysds/runtime/compress/Quantization/QuantizedData.java deleted file mode 100644 index 632066e9854..00000000000 --- a/src/main/java/org/apache/sysds/runtime/compress/Quantization/QuantizedData.java +++ /dev/null @@ -1,57 +0,0 @@ -package org.apache.sysds.runtime.compress.Quantization; - -import java.io.Serializable; - -/** - * Immutable container for probabilistically quantized matrix data. - * Stores quantized byte indices and the scaling parameters needed - * to reconstruct approximate original values on decompression. - * - * - */ -public class QuantizedData implements Serializable { - - private static final long serialVersionUID = 1L; - - public final byte[] quantizedValues; // Quantized level indices - public final double min; // Original minimum value - public final double max; // Original maximum value - public final int levels; // Number of quantization levels (2^bits) - public final int bitsPerValue; // Bits used per element - public final int numRows; - public final int numCols; - - public QuantizedData(byte[] quantizedValues, double min, double max, - int levels, int bitsPerValue, int numRows, int numCols) { - this.quantizedValues = quantizedValues.clone(); // Defensive copy - this.min = min; - this.max = max; - this.levels = levels; - this.bitsPerValue = bitsPerValue; - this.numRows = numRows; - this.numCols = numCols; - } - - /** Number of quantized elements */ - public int size() { - return quantizedValues.length; - } - - /** Estimate serialized size in bytes */ - public long estimateSizeBytes() { - return quantizedValues.length + 64; // +64 for scalar fields and headers - } - - /** Reconstruct a double value from a quantized level index */ - public double reconstructValue(byte levelIndex) { - if (max - min < 1e-10) return min; // Constant matrix - int idx = levelIndex & 0xFF; // Treat byte as unsigned - return min + (idx / (double)(levels - 1)) * (max - min); - } - - @Override - public String toString() { - return String.format("QuantizedData[%dx%d, levels=%d, bits=%d, min=%.4f, max=%.4f]", - numRows, numCols, levels, bitsPerValue, min, max); - } -} \ No newline at end of file diff --git a/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKCompressor.java b/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKCompressor.java deleted file mode 100644 index a8de26db3a6..00000000000 --- a/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKCompressor.java +++ /dev/null @@ -1,226 +0,0 @@ -package org.apache.sysds.runtime.compress.TopK; - -import org.apache.sysds.runtime.compress.CompressedMatrix; -import org.apache.sysds.runtime.compress.CompressionType; -import org.apache.sysds.runtime.compress.MatrixCompressor; -import org.apache.sysds.runtime.compress.exceptions.CompressionException; -import org.apache.sysds.runtime.compress.exceptions.DecompressionException; -import org.apache.sysds.runtime.matrix.data.MatrixBlock; - -import java.util.ArrayList; -import java.util.Comparator; -import java.util.List; -import java.util.PriorityQueue; - -/** - * TopK Sparsification Compressor. - * - * Keeps only the K largest-magnitude elements in the matrix, - * setting all others to zero. Optimal for gradient sparsification - * in federated learning where most gradient values are near zero. - * - * Compression ratio: approximately 1/sparsityRatio - * e.g. sparsityRatio=0.01 keeps 1% of elements → ~100x compression - * - * @author Nirvan Jhurree - */ -public class TopKCompressor implements MatrixCompressor { - - private final double sparsityRatio; // Fraction of elements to keep (0, 1] - private final boolean useHeap; // Use min-heap for O(n log k) selection - - /** - * @param sparsityRatio Fraction of elements to retain e.g. 0.01 = keep top 1% - * @param useHeap If true, use priority queue (faster for large matrices) - */ - public TopKCompressor(double sparsityRatio, boolean useHeap) { - if (sparsityRatio <= 0 || sparsityRatio > 1) { - throw new IllegalArgumentException("sparsityRatio must be in (0, 1]"); - } - this.sparsityRatio = sparsityRatio; - this.useHeap = useHeap; - } - - /** Default constructor: uses heap-based selection */ - public TopKCompressor(double sparsityRatio) { - this(sparsityRatio, true); - } - - @Override - public CompressedMatrix compress(MatrixBlock input) throws CompressionException { - try { - int numRows = input.getNumRows(); - int numCols = input.getNumColumns(); - int totalElements = numRows * numCols; - int k = (int) Math.max(1, Math.ceil(totalElements * sparsityRatio)); - - // If k covers everything, no compression needed - if (k >= totalElements) { - return new CompressedMatrix( - CompressionType.TOPK, numRows, numCols, input, 1.0); - } - - // Extract all non-zero elements with their linear indices - List elements = extractElements(input, numRows, numCols); - - // If fewer non-zeros than k, keep all of them - List topK = (elements.size() <= k) - ? new ArrayList<>(elements) - : selectTopK(elements, k); - - // Pack into TopKData - TopKData data = convertToTopKData(topK, numCols); - - double ratio = calculateCompressionRatio(totalElements, topK.size()); - - return new CompressedMatrix( - CompressionType.TOPK, numRows, numCols, data, ratio); - - } catch (Exception e) { - throw new CompressionException("TopK compression failed: " + e.getMessage(), e); - } - } - - @Override - public MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException { - try { - // Handle passthrough case (no compression was applied) - if (compressed.getCompressedData() instanceof MatrixBlock) { - return (MatrixBlock) compressed.getCompressedData(); - } - - TopKData data = (TopKData) compressed.getCompressedData(); - MatrixBlock result = new MatrixBlock( - compressed.getNumRows(), - compressed.getNumCols(), - true // Start sparse - ); - result.allocateSparseRowsBlock(); - - // Place values back at their original positions - for (int i = 0; i < data.indices.length; i++) { - int linearIdx = data.indices[i]; - int row = linearIdx / data.numCols; - int col = linearIdx % data.numCols; - //result.setValue(row, col, data.values[i]); - result.set(row, col, data.values[i]); - } - - result.examSparsity(); - return result; - - } catch (ClassCastException e) { - throw new DecompressionException("Invalid compressed data type for TopK", e); - } catch (Exception e) { - throw new DecompressionException("TopK decompression failed: " + e.getMessage(), e); - } - } - - @Override - public CompressionType getCompressionType() { - return CompressionType.TOPK; - } - - // ----------------------------------------------------------------------- - // Private helpers - // ----------------------------------------------------------------------- - - /** - * Extract all non-zero elements with their linear indices. - * Handles both dense and sparse MatrixBlock representations. - */ - private List extractElements(MatrixBlock input, int numRows, int numCols) { - List elements = new ArrayList<>(); - - if (input.isInSparseFormat()) { - // Sparse: iterate only over non-zero entries - for (int i = 0; i < numRows; i++) { - if (input.getSparseBlock() == null) continue; - if (input.getSparseBlock().isEmpty(i)) continue; - int[] rowIndices = input.getSparseBlock().indexes(i); - double[] rowValues = input.getSparseBlock().values(i); - int rowSize = input.getSparseBlock().size(i); - for (int j = 0; j < rowSize; j++) { - double val = rowValues[j]; - if (val != 0.0) { - int linearIdx = i * numCols + rowIndices[j]; - elements.add(new Element(linearIdx, val, Math.abs(val))); - } - } - } - } else { - // Dense: iterate all elements, skip zeros - double[] denseBlock = input.getDenseBlockValues(); - if (denseBlock != null) { - for (int i = 0; i < denseBlock.length; i++) { - if (denseBlock[i] != 0.0) { - elements.add(new Element(i, denseBlock[i], Math.abs(denseBlock[i]))); - } - } - } - } - return elements; - } - - /** - * Select top K elements by absolute value. - * Uses min-heap for O(n log k) when useHeap=true, - * or full sort O(n log n) otherwise. - */ - private List selectTopK(List elements, int k) { - if (useHeap) { - PriorityQueue minHeap = new PriorityQueue<>( - k, Comparator.comparingDouble(e -> e.absValue) - ); - for (Element e : elements) { - if (minHeap.size() < k) { - minHeap.offer(e); - } else if (e.absValue > minHeap.peek().absValue) { - minHeap.poll(); - minHeap.offer(e); - } - } - List result = new ArrayList<>(minHeap); - result.sort(Comparator.comparingInt(e -> e.index)); - return result; - } else { - elements.sort((a, b) -> Double.compare(b.absValue, a.absValue)); - return new ArrayList<>(elements.subList(0, k)); - } - } - - private TopKData convertToTopKData(List topK, int numCols) { - int[] indices = new int[topK.size()]; - double[] values = new double[topK.size()]; - for (int i = 0; i < topK.size(); i++) { - indices[i] = topK.get(i).index; - values[i] = topK.get(i).value; - } - return new TopKData(indices, values, numCols); - } - - private double calculateCompressionRatio(int total, int kept) { - if (kept == 0) return Double.MAX_VALUE; - // Original: total * 8 bytes (doubles) - // Compressed: kept * 12 bytes (int index + double value) - long originalBytes = (long) total * 8; - long compressedBytes = (long) kept * 12; - return (double) originalBytes / compressedBytes; - } - - // ----------------------------------------------------------------------- - // Inner class: element tracking during compression - // ----------------------------------------------------------------------- - - private static class Element { - final int index; // Linear index: row * numCols + col - final double value; // Original value - final double absValue; // Absolute value for magnitude comparison - - Element(int index, double value, double absValue) { - this.index = index; - this.value = value; - this.absValue = absValue; - } - } -} \ No newline at end of file diff --git a/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKData.java b/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKData.java deleted file mode 100644 index 509653f7cd0..00000000000 --- a/src/main/java/org/apache/sysds/runtime/compress/TopK/TopKData.java +++ /dev/null @@ -1,44 +0,0 @@ -package org.apache.sysds.runtime.compress.TopK; - -import java.io.Serializable; - -/** - * Immutable container for TopK-compressed matrix data. - * Stores only the K largest-magnitude elements with their positions, - * designed for efficient serialization across federated workers. - * - * @author Nirvan C. Udaysingh Jhurree - */ -public class TopKData implements Serializable { - - private static final long serialVersionUID = 1L; - - public final int[] indices; // Linear indices of kept elements (row*numCols + col) - public final double[] values; // Corresponding original values - public final int numCols; // Needed for index → (row, col) conversion - - public TopKData(int[] indices, double[] values, int numCols) { - if (indices.length != values.length) { - throw new IllegalArgumentException( - "Indices and values arrays must have the same length"); - } - this.indices = indices.clone(); // Defensive copy - this.values = values.clone(); - this.numCols = numCols; - } - - /** Number of kept elements */ - public int size() { - return indices.length; - } - - /** Estimate serialized size in bytes (4 bytes per int + 8 bytes per double) */ - public long estimateSizeBytes() { - return (long) indices.length * 12 + 64; // +64 for object headers - } - - @Override - public String toString() { - return String.format("TopKData[k=%d, numCols=%d]", indices.length, numCols); - } -} \ No newline at end of file diff --git a/src/main/java/org/apache/sysds/runtime/compress/exceptions/CompressionException.java b/src/main/java/org/apache/sysds/runtime/compress/exceptions/CompressionException.java deleted file mode 100644 index da62873aa7d..00000000000 --- a/src/main/java/org/apache/sysds/runtime/compress/exceptions/CompressionException.java +++ /dev/null @@ -1,19 +0,0 @@ -package org.apache.sysds.runtime.compress.exceptions; - -/** - * Exception thrown when matrix compression fails. - * - * @author Nirvan C. UdaysinghJhurree - */ -public class CompressionException extends Exception { - - private static final long serialVersionUID = 1L; - - public CompressionException(String message) { - super(message); - } - - public CompressionException(String message, Throwable cause) { - super(message, cause); - } -} \ No newline at end of file diff --git a/src/main/java/org/apache/sysds/runtime/compress/exceptions/DecompressionException.java b/src/main/java/org/apache/sysds/runtime/compress/exceptions/DecompressionException.java deleted file mode 100644 index ff060588200..00000000000 --- a/src/main/java/org/apache/sysds/runtime/compress/exceptions/DecompressionException.java +++ /dev/null @@ -1,19 +0,0 @@ -package org.apache.sysds.runtime.compress.exceptions; - -/** - * Exception thrown when matrix decompression fails. - * - * @author Nirvan C. Udaysingh Jhurree - */ -public class DecompressionException extends Exception { - - private static final long serialVersionUID = 1L; - - public DecompressionException(String message) { - super(message); - } - - public DecompressionException(String message, Throwable cause) { - super(message, cause); - } -} \ No newline at end of file