[SYSTEMDS-3730] Multithreaded roll operation and improved tests

This patch introduces multi-threading support for the roll operation to
improve performance. The RollTest.java has been updated to cover
both single and multithreaded execution modes.

Furthermore, this update adds comprehensive consistency checks to ensure
mathematical correctness. New tests were created to validate both dense
and sparse matrix inputs. Additionally, cross-verification tests were
added to confirm that sparse and dense rolling for single and
multithreaded executions produce identical results.

Closes #2376.

Author: Biranavan Parameswaran

src/main/java/org/apache/sysds/hops/ReorgOp.java

@@ -173,7 +173,9 @@ _op, getDataType(), getValueType(), et,
 				for (int i = 0; i < 2; i++)
 					linputs[i] = getInput().get(i).constructLops();
 
-				Transform transform1 = new Transform(linputs, _op, getDataType(), getValueType(), et, 1);
+				Transform transform1 = new Transform(
+						linputs, _op, getDataType(), getValueType(), et,
+						OptimizerUtils.getConstrainedNumThreads(_maxNumThreads));
 
 				setOutputDimensions(transform1);
 				setLineNumbers(transform1);

src/main/java/org/apache/sysds/lops/Transform.java

@@ -180,7 +180,7 @@ private String getInstructions(String input1, int numInputs, String output) {
 		sb.append( this.prepOutputOperand(output));
 
 		if( (getExecType()==ExecType.CP || getExecType()==ExecType.FED || getExecType()==ExecType.OOC)
-			&& (_operation == ReOrgOp.TRANS || _operation == ReOrgOp.REV || _operation == ReOrgOp.SORT) ) {
+			&& (_operation == ReOrgOp.TRANS || _operation == ReOrgOp.REV || _operation == ReOrgOp.SORT || _operation == ReOrgOp.ROLL) ) {
 			sb.append( OPERAND_DELIMITOR );
 			sb.append( _numThreads );
 			if ( getExecType()==ExecType.FED ) {

src/main/java/org/apache/sysds/runtime/instructions/cp/ReorgCPInstruction.java

@@ -118,11 +118,13 @@ else if ( opcode.equalsIgnoreCase(Opcodes.REV.toString()) ) {
 			return new ReorgCPInstruction(new ReorgOperator(RevIndex.getRevIndexFnObject(), k), in, out, opcode, str);
 		}
 		else if (opcode.equalsIgnoreCase(Opcodes.ROLL.toString())) {
-			InstructionUtils.checkNumFields(str, 3);
+			InstructionUtils.checkNumFields(str, 3, 4);
 			in.split(parts[1]);
 			out.split(parts[3]);
 			CPOperand shift = new CPOperand(parts[2]);
-			return new ReorgCPInstruction(new ReorgOperator(new RollIndex(0)), in, out, shift, opcode, str);
+			int k = (parts.length > 4) ? Integer.parseInt(parts[4]) : 1;
+
+			return new ReorgCPInstruction(new ReorgOperator(new RollIndex(0), k), in, out, shift, opcode, str);
 		}
 		else if ( opcode.equalsIgnoreCase(Opcodes.DIAG.toString()) ) {
 			parseUnaryInstruction(str, in, out); //max 2 operands

src/main/java/org/apache/sysds/runtime/matrix/data/LibMatrixReorg.java

@@ -134,6 +134,8 @@ public static MatrixBlock reorg( MatrixBlock in, MatrixBlock out, ReorgOperator
 					return rev(in, out);
 			case ROLL:
 				RollIndex rix = (RollIndex) op.fn;
+				if(op.getNumThreads() > 1)
+					return roll(in, out, rix.getShift(), op.getNumThreads());
 				return roll(in, out, rix.getShift());
 			case DIAG:
 				return diag(in, out);
@@ -514,6 +516,124 @@ public static MatrixBlock roll(MatrixBlock in, MatrixBlock out, int shift) {
 		return out;
 	}
 
+	public static MatrixBlock roll(MatrixBlock input, MatrixBlock output, int shift, int numThreads) {
+
+		final int numRows = input.rlen;
+		final int numCols = input.clen;
+		final boolean isSparse = input.sparse;
+
+		// sparse-safe operation
+		if(input.isEmptyBlock(false))
+			return output;
+
+		// special case: row vector
+		if(numRows == 1) {
+			output.copy(input);
+			return output;
+		}
+
+		if(numThreads <= 1 || input.getLength() < PAR_NUMCELL_THRESHOLD) {
+			return roll(input, output, shift); // fallback to single-threaded
+		}
+
+		final int normalizedShift = getNormalizedShiftForRoll(shift, numRows);
+
+		output.reset(numRows, numCols, isSparse);
+		output.nonZeros = input.nonZeros;
+
+		if(isSparse) {
+			output.allocateSparseRowsBlock(false);
+		}
+		else {
+			output.allocateDenseBlock(false);
+		}
+
+		//TODO experiment with more tasks per thread for better load balance
+		//TODO call common kernel from both single- and multi-threaded execution
+		
+		ExecutorService threadPool = CommonThreadPool.get(numThreads);
+		try {
+			final int rowsPerThread = (int) Math.ceil((double) numRows / numThreads);
+			List<Future<?>> tasks = new ArrayList<>();
+
+			for(int threadIndex = 0; threadIndex < numThreads; threadIndex++) {
+
+				final int startRow = threadIndex * rowsPerThread;
+				final int endRow = Math.min((threadIndex + 1) * rowsPerThread, numRows);
+
+				tasks.add(threadPool.submit(() -> {
+					if(isSparse)
+						rollSparseBlock(input, output, normalizedShift, startRow, endRow);
+					else
+						rollDenseBlock(input, output, normalizedShift, startRow, endRow);
+				}));
+			}
+
+			for(Future<?> task : tasks)
+				task.get();
+
+		}
+		catch(Exception ex) {
+			throw new DMLRuntimeException(ex);
+		}
+		finally {
+			threadPool.shutdown();
+		}
+
+		return output;
+	}
+
+	private static int getNormalizedShiftForRoll(int shift, int numRows) {
+		shift = shift % numRows;
+		if(shift < 0)
+			shift += numRows;
+
+		return shift;
+	}
+
+	private static void rollDenseBlock(MatrixBlock input, MatrixBlock output, 
+			int shift, int startRow, int endRow) 
+	{
+		DenseBlock inputBlock = input.getDenseBlock();
+		DenseBlock outputBlock = output.getDenseBlock();
+		final int numRows = input.rlen;
+		final int numCols = input.clen;
+
+		for(int targetRow = startRow; targetRow < endRow; targetRow++) {
+			int sourceRow = targetRow - shift;
+			if(sourceRow < 0)
+				sourceRow += numRows;
+
+			System.arraycopy(inputBlock.values(sourceRow), inputBlock.pos(sourceRow), outputBlock.values(targetRow),
+				outputBlock.pos(targetRow), numCols);
+		}
+	}
+
+	private static void rollSparseBlock(MatrixBlock input, MatrixBlock output, 
+		int shift, int startRow, int endRow) 
+	{
+		SparseBlock inputBlock = input.getSparseBlock();
+		SparseBlock outputBlock = output.getSparseBlock();
+		final int numRows = input.rlen;
+
+		for(int targetRow = startRow; targetRow < endRow; targetRow++) {
+			int sourceRow = targetRow - shift;
+			if(sourceRow < 0)
+				sourceRow += numRows;
+
+			if(!inputBlock.isEmpty(sourceRow)) {
+				int rowStart = inputBlock.pos(sourceRow);
+				int rowEnd = rowStart + inputBlock.size(sourceRow);
+				int[] colIndexes = inputBlock.indexes(sourceRow);
+				double[] values = inputBlock.values(sourceRow);
+
+				for(int k = rowStart; k < rowEnd; k++) {
+					outputBlock.set(targetRow, colIndexes[k], values[k]);
+				}
+			}
+		}
+	}
+
 	public static void roll(IndexedMatrixValue in, long rlen, int blen, int shift, ArrayList<IndexedMatrixValue> out) {
 		MatrixIndexes inMtxIdx = in.getIndexes();
 		MatrixBlock inMtxBlk = (MatrixBlock) in.getValue();
@@ -2554,15 +2674,15 @@ private static void reverseSparse(MatrixBlock in, MatrixBlock out, int rl, int r
 
 	private static void rollDense(MatrixBlock in, MatrixBlock out, int shift) {
 		final int m = in.rlen;
-		shift %= (m != 0 ? m : 1); // roll matrix with axis=none
+		shift = getNormalizedShiftForRoll(shift, m); // roll matrix with axis=none
 
 		copyDenseMtx(in, out, 0, shift, m - shift, false, true);
 		copyDenseMtx(in, out, m - shift, 0, shift, true, true);
 	}
 
 	private static void rollSparse(MatrixBlock in, MatrixBlock out, int shift) {
 		final int m = in.rlen;
-		shift %= (m != 0 ? m : 1); // roll matrix with axis=0
+		shift = getNormalizedShiftForRoll(shift, m); // roll matrix with axis=0
 
 		copySparseMtx(in, out, 0, shift, m - shift, false, true);
 		copySparseMtx(in, out, m-shift, 0, shift, false, true);

src/test/java/org/apache/sysds/performance/matrix/MatrixRollPerf.java

@@ -0,0 +1,127 @@
+/*
+ * 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.performance.matrix;
+
+import org.apache.sysds.performance.compression.APerfTest;
+import org.apache.sysds.performance.generators.ConstMatrix;
+import org.apache.sysds.performance.generators.IGenerate;
+import org.apache.sysds.runtime.functionobjects.IndexFunction;
+import org.apache.sysds.runtime.functionobjects.RollIndex;
+import org.apache.sysds.runtime.matrix.data.MatrixBlock;
+import org.apache.sysds.runtime.matrix.operators.ReorgOperator;
+import org.apache.sysds.test.TestUtils;
+import org.apache.sysds.utils.stats.InfrastructureAnalyzer;
+
+import java.util.Random;
+
+public class MatrixRollPerf extends APerfTest<Object, MatrixBlock> {
+
+	private final int rows;
+	private final int cols;
+	private final int shift;
+	private final int k;
+
+	private final ReorgOperator reorg;
+	private MatrixBlock out;
+
+	public MatrixRollPerf(int N, int W, IGenerate<MatrixBlock> gen, int rows, int cols, int shift, int k) {
+		super(N, W, gen);
+		this.rows = rows;
+		this.cols = cols;
+		this.shift = shift;
+		this.k = k;
+
+		IndexFunction op = new RollIndex(shift);
+		this.reorg = new ReorgOperator(op, k);
+	}
+
+	public void run() throws Exception {
+		MatrixBlock mb = gen.take();
+		logInfos(rows, cols, shift, mb.getSparsity(), k);
+
+
+		String info = String.format("rows: %5d cols: %5d sp: %.4f shift: %4d k: %2d",
+			rows, cols, mb.getSparsity(), shift, k);
+
+
+		warmup(this::rollOnce, W);
+
+		execute(this::rollOnce, info);
+	}
+
+	private void logInfos(int rows, int cols, int shift, double sparsity, int k) {
+		String matrixType = sparsity == 1 ? "Dense" : "Sparse";
+		if (k == 1) {
+			System.out.println("---------------------------------------------------------------------------------------------------------");
+			System.out.printf("%s Experiment for rows %d columns %d and shift %d \n", matrixType, rows, cols, shift);
+			System.out.println("---------------------------------------------------------------------------------------------------------");
+		}
+	}
+
+	private void rollOnce() {
+		MatrixBlock in = gen.take();
+
+		if (out == null)
+			out = new MatrixBlock(rows, cols, in.isInSparseFormat());
+
+		out.reset(rows, cols, in.isInSparseFormat());
+
+		in.reorgOperations(reorg, out, 0, 0, 0);
+
+		ret.add(null);
+	}
+
+	@Override
+	protected String makeResString() {
+		return "";
+	}
+
+	public static void main(String[] args) throws Exception {
+		int kMulti = InfrastructureAnalyzer.getLocalParallelism();
+		int reps = 2000;
+		int warmup = 200;
+
+		int minRows = 2017;
+		int minCols = 1001;
+		double spSparse = 0.01;
+		int minShift = -50;
+		int maxShift = 1022;
+		int iterations = 10;
+
+		Random rand = new Random(42);
+
+		for (int i = 0; i < iterations; i++) {
+			int rows = 10_000_000;
+			int cols = 10;
+			int shift = rand.nextInt((maxShift - minShift) + 1) + minShift;
+
+			MatrixBlock denseIn = TestUtils.generateTestMatrixBlock(rows, cols, -100, 100, 1.0, 42);
+			MatrixBlock sparseIn = TestUtils.generateTestMatrixBlock(rows, cols, -100, 100, spSparse, 42);
+
+			// Run Dense Case (Single vs Multi-threaded)
+			new MatrixRollPerf(reps, warmup, new ConstMatrix(denseIn, -1), rows, cols, shift, 1).run();
+			new MatrixRollPerf(reps, warmup, new ConstMatrix(denseIn, -1), rows, cols, shift, kMulti).run();
+
+			// Run Sparse Case (Single vs Multi-threaded)
+			new MatrixRollPerf(reps, warmup, new ConstMatrix(sparseIn, -1), rows, cols, shift, 1).run();
+			new MatrixRollPerf(reps, warmup, new ConstMatrix(sparseIn, -1), rows, cols, shift, kMulti).run();
+		}
+	}
+}