ywcb00 commented on code in PR #2524:
URL: https://github.com/apache/systemds/pull/2524#discussion_r3496748071
##########
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
Review Comment:
Remove author information.
##########
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
Review Comment:
Author information is already stored by git. No need to add it explicitly to
the function header.
##########
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.
+ *
+ *
Review Comment:
Remove empty lines.
##########
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
+ *
+ *
Review Comment:
Remove empty lines.
##########
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) {
Review Comment:
You can merge these two methods into a single one, as only the above method
is called from outside this class.
##########
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();
+ *
+ *
Review Comment:
Remove empty lines.
##########
src/main/java/org/apache/sysds/runtime/compress/exceptions/CompressionException.java:
##########
@@ -0,0 +1,19 @@
+package org.apache.sysds.runtime.compress.exceptions;
Review Comment:
Please move everything that you created in this package to the package
`org.apache.sysds.runtime.controlprogram.federated.compression` and the
corresponding folder.
##########
src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java:
##########
@@ -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) {
Review Comment:
Introduce an execution option to control whether federated compression is
enabled or disabled, similar to `DMLScript.LINEAGE` for enabling lineage
functionality.
##########
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;
Review Comment:
Remove if not needed.
##########
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;
+ }
Review Comment:
Remove if not needed.
##########
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();
Review Comment:
Remove if not needed.
##########
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 {
Review Comment:
Please try to merge some of these tests. For example, you can check the
compression type, the dimensions, the range, and the ratio within the same unit
test.
[This applies to both test classes that are added with this PR]
##########
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);
+ }
+}
Review Comment:
Please add an empty line to the very end of every file.
##########
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 {
Review Comment:
You can extend from the `AutomatedTestBase` class and leverage the utility
functions in `org.apache.sysds.test.TestUtils`.
[This applies to both test classes that are added with this PR]
##########
.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@v4
+ with:
+ name: test-results
+ path: target/surefire-reports/
Review Comment:
No need to create a separate workflow for this test. You can move your test
to class to one of the test folders so that it gets executed automatically, or
you can add the path to your test to the `javaTests.yml` workflow.
##########
src/test/java/org/apache/sysds/runtime/compress/Quantization/ProbabilisticQuantizationCompressorTest.java:
##########
@@ -0,0 +1,187 @@
+package org.apache.sysds.runtime.compress.Quantization;
Review Comment:
Move both test classes to `org.apache.sysds.test.functions.federated.io`.
Then, the test will run with the general java test workflow.
[This applies to both test classes that are added with this PR]
##########
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
Review Comment:
Why do we need this line? Can the value be smaller than zero or larger than
1 after normalizing it to [0, 1]?
##########
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");
Review Comment:
Are you still planning on implementing compressed sensing?
If yes, here is a reference that might be useful:
[https://doi.org/10.1109/CISS.2008.4558487](https://doi.org/10.1109/CISS.2008.4558487)
Otherwise, please remove this case.
##########
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);
Review Comment:
Do we need the Math.min operation here?
##########
src/main/java/org/apache/sysds/runtime/compress/TopK/TopKCompressor.java:
##########
@@ -0,0 +1,226 @@
+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<Element> elements = extractElements(input, numRows, numCols);
+
+ // If fewer non-zeros than k, keep all of them
+ List<Element> 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;
+ }
Review Comment:
Remove if not needed.
##########
src/main/java/org/apache/sysds/runtime/compress/exceptions/CompressionException.java:
##########
@@ -0,0 +1,19 @@
+package org.apache.sysds.runtime.compress.exceptions;
Review Comment:
Please add the header of the Apache v2.0 license to every file. You can find
this header in every other file of the project.
##########
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();
Review Comment:
Import the MatrixBlock and use it as `MatrixBlock` instead of writing the
full path.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
