[BWARE] Track more compressed-friendly ops in FederatedWorkloadAnalyzer (#2481)

* Track more compressed-friendly ops in FederatedWorkloadAnalyzer

Extends the federated workload counter so that compression decisions
account for additional instruction shapes beyond AggregateBinary.

- Pass the right-hand column count to incOverlappingDecompressions
  so the cost model reflects the actual decompression size rather
  than counting a single column
- Count MMChainCPInstruction as one LMM and one RMM contribution per
  invocation
- Count AggregateUnaryCPInstruction: when reducing columns with a
  sum/mean operator, treat it as a dict-op (compression-friendly);
  otherwise count it as a decompression
- Minor formatting cleanup in compressRun

* Add unit tests for FederatedWorkloadAnalyzer workload tracking

Cover the instruction-shape branches in incrementWorkload that drive
federated compression decisions, which previously had no direct tests:

- AggregateBinary: RMM/LMM counting, overlapping-decompress sizing by the
  right-hand column count, and the validSize row/column guards
- MMChain: one LMM and one RMM contribution per invocation
- AggregateUnary: dict-op vs decompression classification across
  ReduceAll/ReduceRow/ReduceCol with sum, mean, product, and max operators
- Instance-level dispatch and compressRun threshold behavior, asserting
  async compression materializes when the cost model would compress

Author: Baunsgaard

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

@@ -27,9 +27,18 @@
 import org.apache.sysds.runtime.compress.cost.InstructionTypeCounter;
 import org.apache.sysds.runtime.controlprogram.caching.MatrixObject;
 import org.apache.sysds.runtime.controlprogram.context.ExecutionContext;
+import org.apache.sysds.runtime.functionobjects.IndexFunction;
+import org.apache.sysds.runtime.functionobjects.KahanPlus;
+import org.apache.sysds.runtime.functionobjects.Mean;
+import org.apache.sysds.runtime.functionobjects.Plus;
+import org.apache.sysds.runtime.functionobjects.ReduceCol;
 import org.apache.sysds.runtime.instructions.Instruction;
 import org.apache.sysds.runtime.instructions.cp.AggregateBinaryCPInstruction;
+import org.apache.sysds.runtime.instructions.cp.AggregateUnaryCPInstruction;
 import org.apache.sysds.runtime.instructions.cp.ComputationCPInstruction;
+import org.apache.sysds.runtime.instructions.cp.MMChainCPInstruction;
+import org.apache.sysds.runtime.matrix.operators.AggregateUnaryOperator;
+import org.apache.sysds.runtime.matrix.operators.Operator;
 
 public class FederatedWorkloadAnalyzer {
 	protected static final Log LOG = LogFactory.getLog(FederatedWorkloadAnalyzer.class.getName());
@@ -55,7 +64,7 @@ public void incrementWorkload(ExecutionContext ec, long tid, Instruction ins) {
 	}
 
 	public void compressRun(ExecutionContext ec, long tid) {
-		if(counter >= compressRunFrequency ){
+		if(counter >= compressRunFrequency) {
 			counter = 0;
 			get(tid).forEach((K, V) -> CompressedMatrixBlockFactory.compressAsync(ec, Long.toString(K), V));
 		}
@@ -68,6 +77,7 @@ private void incrementWorkload(ExecutionContext ec, long tid, ComputationCPInstr
 	public void incrementWorkload(ExecutionContext ec, ConcurrentHashMap<Long, InstructionTypeCounter> mm,
 		ComputationCPInstruction cpIns) {
 		// TODO: Count transitive closure via lineage
+		// TODO: add more operations
 		if(cpIns instanceof AggregateBinaryCPInstruction) {
 			final String n1 = cpIns.input1.getName();
 			MatrixObject d1 = (MatrixObject) ec.getCacheableData(n1);
@@ -81,15 +91,45 @@ public void incrementWorkload(ExecutionContext ec, ConcurrentHashMap<Long, Instr
 			if(validSize(r1, c1)) {
 				getOrMakeCounter(mm, Long.parseLong(n1)).incRMM(c2);
 				// safety add overlapping decompress for RMM
-				getOrMakeCounter(mm, Long.parseLong(n1)).incOverlappingDecompressions();
+				getOrMakeCounter(mm, Long.parseLong(n1)).incOverlappingDecompressions(c2);
 				counter++;
 			}
 			if(validSize(r2, c2)) {
 				getOrMakeCounter(mm, Long.parseLong(n2)).incLMM(r1);
 				counter++;
 			}
-			
 		}
+		else if(cpIns instanceof MMChainCPInstruction) {
+			final String n1 = cpIns.input1.getName();
+			getOrMakeCounter(mm, Long.parseLong(n1)).incRMM(1);
+			getOrMakeCounter(mm, Long.parseLong(n1)).incLMM(1);
+			counter++;
+		}
+		else if(cpIns instanceof AggregateUnaryCPInstruction) {
+			Operator op = cpIns.getOperator();
+			final String n1 = cpIns.input1.getName();
+			long id = Long.parseLong(n1);
+			if(op instanceof AggregateUnaryOperator) {
+				AggregateUnaryOperator aop = (AggregateUnaryOperator) op;
+				IndexFunction idxF = aop.indexFn;
+				getOrMakeCounter(mm, id).incDictOps();
+				if(idxF instanceof ReduceCol) {
+					if((aop.aggOp.increOp.fn instanceof KahanPlus //
+						|| aop.aggOp.increOp.fn instanceof Plus //
+						|| aop.aggOp.increOp.fn instanceof Mean)) {
+						getOrMakeCounter(mm, id).incDictOps();
+					}
+					else {
+						// increment decompression if row reduce.
+						getOrMakeCounter(mm, id).incDecompressions();
+					}
+				}
+				else {
+					getOrMakeCounter(mm, id).incDictOps();
+				}
+			}
+		}
+
 	}
 
 	private static InstructionTypeCounter getOrMakeCounter(ConcurrentHashMap<Long, InstructionTypeCounter> mm, long id) {
@@ -117,8 +157,8 @@ private static boolean validSize(int nRow, int nCol) {
 		return nRow > 90 && nRow >= nCol;
 	}
 
-	@Override 
-	public String toString(){
+	@Override
+	public String toString() {
 		StringBuilder sb = new StringBuilder();
 		sb.append(this.getClass().getSimpleName());
 		sb.append("  Counter: ");