Matrix compression for federated learning broadcasts

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

src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java

@@ -42,12 +42,19 @@
 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.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;
 import org.apache.sysds.runtime.util.IndexRange;
@@ -138,6 +145,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(DMLScript.FEDERATED_COMPRESSION && 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);
 	}
 

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);
+	}
+}

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<String, Object> 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<String, Object> 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<String, Object> 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);
+		}
+	}
+}

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();
+		}
+	}
+}