[SYSTEMDS-3891] Add Broadcast and Grouped Reduction Primitives

src/main/java/org/apache/sysds/runtime/instructions/ooc/MMultOOCInstruction.java

@@ -33,6 +33,9 @@
 import org.apache.sysds.runtime.matrix.operators.AggregateOperator;
 import org.apache.sysds.runtime.matrix.operators.BinaryOperator;
 import org.apache.sysds.runtime.matrix.operators.Operator;
+import org.apache.sysds.runtime.meta.DataCharacteristics;
+import org.apache.sysds.runtime.ooc.store.CountingLiveness;
+import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils;
 
 public class MMultOOCInstruction extends ComputationOOCInstruction {
 
@@ -60,6 +63,33 @@ public void processInstruction( ExecutionContext ec ) {
 		// 1. Identify the inputs
 		MatrixObject min = ec.getMatrixObject(input1); // big matrix
 		MatrixObject vin = ec.getMatrixObject(input2); // streamed vector
+		DataCharacteristics mdc = min.getDataCharacteristics();
+		DataCharacteristics vdc = vin.getDataCharacteristics();
+
+		if(min != vin && mdc.getRows() > 0 && mdc.getCols() > 0 && vdc.getCols() > 0 &&
+			mdc.getCols() == vdc.getRows() && vdc.getNumColBlocks() == 1) {
+			OOCStream<IndexedMatrixValue> partials = createWritableStream();
+			OOCStream<IndexedMatrixValue> out = createWritableStream();
+			partials.setData(min);
+			ec.getMatrixObject(output).setStreamHandle(out);
+			OOCInstructionUtils.indexedBroadcastMap(min.getStreamable(), vin.getStreamable(), partials,
+				left -> Math.toIntExact(left.getIndexes().getColumnIndex() - 1),
+				() -> new CountingLiveness(Math.toIntExact(vin.getDataCharacteristics().getNumRowBlocks()),
+					Math.toIntExact(min.getDataCharacteristics().getNumRowBlocks())),
+				(left, right) -> {
+					MatrixBlock leftBlock = (MatrixBlock) left.getValue();
+					MatrixBlock rightBlock = (MatrixBlock) right.getValue();
+					MatrixBlock partial = leftBlock.aggregateBinaryOperations(leftBlock, rightBlock, new MatrixBlock(),
+						(AggregateBinaryOperator) _optr);
+					MatrixIndexes indexes = left.getIndexes();
+					return new IndexedMatrixValue(new MatrixIndexes(indexes.getRowIndex(), indexes.getColumnIndex()),
+						partial);
+				}, getContext());
+			BinaryOperator plus = InstructionUtils.parseBinaryOperator(Opcodes.PLUS.toString());
+			OOCInstructionUtils.rowGroupedReduce(partials, out,
+				(left, right) -> left.binaryOperations(plus, right, new MatrixBlock()), getContext());
+			return;
+		}
 
 		int emitLeftThreshold = (int)vin.getDataCharacteristics().getNumColBlocks();
 		int emitRightThreshold = (int)min.getDataCharacteristics().getNumRowBlocks();

src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCache.java

@@ -24,6 +24,14 @@
 import java.util.function.LongUnaryOperator;
 
 public interface OOCCache {
+	/**
+	 * Maximum bytes charged given the logical bytes of the requested entry. Use this method for reservation budget
+	 * planning as logical byte size and pinned entry bytes may differ.
+	 */
+	default long maxPhysicalPinBytes(long logicalBytes) {
+		return logicalBytes;
+	}
+
 	/**
 	 * Pins an item backed by an allowance. A successful pin transfers memory ownership from the cache to the owner of
 	 * the allowance and guarantees data availability. While pinned, the bytes of the entry are not counted as

src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java

@@ -129,6 +129,14 @@ public OOCPackedCache(OOCCacheImpl physical, long packThresholdBytes, long packT
 		});
 	}
 
+	@Override
+	public long maxPhysicalPinBytes(long logicalBytes) {
+		if(logicalBytes >= _packThresholdBytes)
+			return logicalBytes;
+		return _packTargetBytes > Long.MAX_VALUE - _packThresholdBytes ? Long.MAX_VALUE : _packTargetBytes +
+			_packThresholdBytes;
+	}
+
 	@Override
 	public BlockEntry putPinned(long sId, long tId, Object data, long size, MemoryAllowance allowance) {
 		if(size >= _packThresholdBytes)

src/main/java/org/apache/sysds/runtime/ooc/memory/ReservationBudget.java

@@ -26,6 +26,7 @@ public final class ReservationBudget implements MemoryAllowance, AutoCloseable {
 	private long _outstanding;
 	private long _available;
 	private boolean _closed;
+	private boolean _reusable;
 
 	public ReservationBudget(MemoryAllowance parent, long bytes) {
 		if(parent == null)
@@ -37,6 +38,13 @@ public ReservationBudget(MemoryAllowance parent, long bytes) {
 		_available = bytes;
 	}
 
+	public synchronized ReservationBudget enableReuse() {
+		if(_closed || getUsedMemory() != 0)
+			throw new IllegalStateException("Budget reuse must be enabled before reserving memory.");
+		_reusable = true;
+		return this;
+	}
+
 	@Override
 	public synchronized boolean tryReserve(long bytes) {
 		checkNonNegative(bytes);
@@ -64,13 +72,19 @@ public void release(long bytes) {
 		checkNonNegative(bytes);
 		if(bytes == 0)
 			return;
+		boolean releaseParent;
 		synchronized(this) {
 			long used = _outstanding - _available;
 			if(bytes > used)
 				throw new IllegalStateException("Cannot release " + bytes + " bytes from a budget using " + used);
-			_outstanding -= bytes;
+			releaseParent = _closed || !_reusable;
+			if(releaseParent)
+				_outstanding -= bytes;
+			else
+				_available += bytes;
 		}
-		_parent.release(bytes);
+		if(releaseParent)
+			_parent.release(bytes);
 	}
 
 	@Override