<commit_before>/**
 * IBM Confidential
 * OCO Source Materials
 * (C) Copyright IBM Corp. 2010, 2015
 * The source code for this program is not published or otherwise divested of its trade secrets, irrespective of what has been deposited with the U.S. Copyright Office.
 */


package com.ibm.bi.dml.runtime.matrix.data;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import com.ibm.bi.dml.lops.MapMultChain.ChainType;
import com.ibm.bi.dml.runtime.DMLRuntimeException;
import com.ibm.bi.dml.runtime.DMLUnsupportedOperationException;
import com.ibm.bi.dml.runtime.functionobjects.SwapIndex;
import com.ibm.bi.dml.runtime.matrix.operators.ReorgOperator;
import com.ibm.bi.dml.runtime.util.UtilFunctions;

/**
 * MB:
 * Library for matrix multiplications including MM, MV, VV for all
 * combinations of dense and sparse representations and special
 * operations such as transpose-self matrix multiplication.
 * 
 * In general all implementations use internally dense outputs
 * for direct access, but change the final result to sparse if necessary.  
 * 
 * NOTES on BLAS:
 * * Experiments in 04/2013 showed that even on dense-dense this implementation 
 *   is 3x faster than f2j-BLAS-DGEMM, 2x faster than f2c-BLAS-DGEMM, and
 *   level (+10% after JIT) with a native C implementation. 
 * * Calling native BLAS would loose platform independence and would require 
 *   JNI calls incl data transfer. Furthermore, BLAS does not support sparse 
 *   matrices (except Sparse BLAS, with dedicated function calls and matrix formats) 
 *   and would be an external dependency. 
 * * Experiments in 02/2014 showed that on dense-dense this implementation now achieves
 *   almost 30% peak FP performance. Compared to Intel MKL 11.1 (dgemm, N=1000) it is
 *   just 3.2x (sparsity=1.0) and 1.9x (sparsity=0.5) slower, respectively.
 *   
 *  
 */
public class LibMatrixMult 
{
	@SuppressWarnings("unused")
	private static final String _COPYRIGHT = "Licensed Materials - Property of IBM\n(C) Copyright IBM Corp. 2010, 2015\n" +
                                             "US Government Users Restricted Rights - Use, duplication  disclosure restricted by GSA ADP Schedule Contract with IBM Corp.";
	
	public static final boolean LOW_LEVEL_OPTIMIZATION = true;

	private LibMatrixMult() {
		//prevent instantiation via private constructor
	}
	
	////////////////////////////////
	// public matrix mult interface
	////////////////////////////////
	
	/**
	 * Performs a matrix multiplication and stores the result in the output matrix.
	 * 
	 * All variants use a IKJ access pattern, and internally use dense output. After the
	 * actual computation, we recompute nnz and check for sparse/dense representation.
	 *  
	 * 
	 * @param m1 first matrix
	 * @param m2 second matrix
	 * @param ret result matrix
	 * @throws DMLRuntimeException 
	 */
	public static void matrixMult(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret) 
		throws DMLRuntimeException
	{	
		//check inputs / outputs
		if( m1.isEmptyBlock(false) || m2.isEmptyBlock(false) )
			return;
		
		//Timing time = new Timing(true);
		
		//pre-processing
		ret.sparse = (m1.isUltraSparse() || m2.isUltraSparse());
		if( !ret.sparse )
			ret.allocateDenseBlock();
		
		//core matrix mult computation
		if( m1.isUltraSparse() || m2.isUltraSparse() )
			matrixMultUltraSparse(m1, m2, ret, 0, m1.rlen);
		else if(!m1.sparse && !m2.sparse)
			matrixMultDenseDense(m1, m2, ret, 0, m1.rlen);
		else if(m1.sparse && m2.sparse)
			matrixMultSparseSparse(m1, m2, ret, 0, m1.rlen);
		else if(m1.sparse)
			matrixMultSparseDense(m1, m2, ret, 0, m1.rlen);
		else
			matrixMultDenseSparse(m1, m2, ret, 0, m1.rlen);
		
		//post-processing
		if( !ret.sparse )
			ret.recomputeNonZeros();
		ret.examSparsity();
		
		//System.out.println("MM ("+m1.isInSparseFormat()+","+m1.getNumRows()+","+m1.getNumColumns()+","+m1.getNonZeros()+")x" +
		//		              "("+m2.isInSparseFormat()+","+m2.getNumRows()+","+m2.getNumColumns()+","+m2.getNonZeros()+") in "+time.stop());
	}
	
	/**
	 * Performs a multi-threaded matrix multiplication and stores the result in the output matrix.
	 * The parameter k (k>=1) determines the max parallelism k' with k'=min(k, vcores, m1.rlen).
	 * 
	 * @param m1
	 * @param m2
	 * @param ret
	 * @param k
	 * @throws DMLRuntimeException 
	 */
	public static void matrixMult(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret, int k) 
		throws DMLRuntimeException
	{	
		//check inputs / outputs
		if( m1.isEmptyBlock(false) || m2.isEmptyBlock(false) )
			return;
		//check no parallelization benefit (fallback to sequential)
		if( m1.rlen == 1 ) {
			matrixMult(m1, m2, ret);
			return;
		}
		
		//Timing time = new Timing(true);
		
		//pre-processing
		ret.sparse = (m1.isUltraSparse() || m2.isUltraSparse());
		if( !ret.sparse )
			ret.allocateDenseBlock();
		
		//core multi-threaded matrix mult computation
		//(currently: always parallelization over number of rows)
		try {
			ExecutorService pool = Executors.newFixedThreadPool( k );
			ArrayList<MatrixMultTask> tasks = new ArrayList<MatrixMultTask>();
			int blklen = (int)(Math.ceil((double)m1.rlen/k));
			for( int i=0; i<k & i*blklen<m1.rlen; i++ )
				tasks.add(new MatrixMultTask(m1, m2, ret, i*blklen, Math.min((i+1)*blklen, m1.rlen)));
			pool.invokeAll(tasks);	
			pool.shutdown();
		}
		catch(Exception ex) {
			throw new DMLRuntimeException(ex);
		}
		
		//post-processing
		if( !ret.sparse )
			ret.recomputeNonZeros();
		ret.examSparsity();
		
		//System.out.println("MM ("+m1.isInSparseFormat()+","+m1.getNumRows()+","+m1.getNumColumns()+","+m1.getNonZeros()+")x" +
		//		              "("+m2.isInSparseFormat()+","+m2.getNumRows()+","+m2.getNumColumns()+","+m2.getNonZeros()+") in "+time.stop());
	}
	
	/**
	 * Performs a matrix multiplication chain operation of type t(X)%*%(X%*%v) or t(X)%*%(w*(X%*%v)).
	 * 
	 * All variants use a IKJ access pattern, and internally use dense output. After the
	 * actual computation, we recompute nnz and check for sparse/dense representation.
	 * 
	 * @param m1
	 * @param m2
	 * @param w
	 * @param ret
	 * @param ct
	 * @throws DMLRuntimeException
	 * @throws DMLUnsupportedOperationException
	 */
	public static void matrixMultChain(MatrixBlock mX, MatrixBlock mV, MatrixBlock mW, MatrixBlock ret, ChainType ct) 
		throws DMLRuntimeException, DMLUnsupportedOperationException
	{		
		//check inputs / outputs (after that mV and mW guaranteed to be dense)
		if( mX.isEmptyBlock(false) || mV.isEmptyBlock(false) || (mW !=null && mW.isEmptyBlock(false)) )
			return;

		//Timing time = new Timing(true);
				
		//pre-processing
		ret.sparse = false;
		ret.allocateDenseBlock();
		
		//core matrix mult chain computation
		if( mX.sparse )
			matrixMultChainSparse(mX, mV, mW, ret, ct, 0, mX.rlen);
		else
			matrixMultChainDense(mX, mV, mW, ret, ct, 0, mX.rlen);
		
		//post-processing
		ret.recomputeNonZeros();
		ret.examSparsity();
		
		//System.out.println("MMChain "+ct.toString()+" ("+mX.isInSparseFormat()+","+mX.getNumRows()+","+mX.getNumColumns()+","+mX.getNonZeros()+")x" +
		//		              "("+mV.isInSparseFormat()+","+mV.getNumRows()+","+mV.getNumColumns()+","+mV.getNonZeros()+") in "+time.stop());
	}

	/**
	 * Performs a parallel matrix multiplication chain operation of type t(X)%*%(X%*%v) or t(X)%*%(w*(X%*%v)).
	 * The parameter k (k>=1) determines the max parallelism k' with k'=min(k, vcores, m1.rlen).
	 * 
	 * NOTE: This multi-threaded mmchain operation has additional memory requirements of k*ncol(X)*8bytes 
	 * for partial aggregation. Current max memory: 256KB; otherwise redirectly to sequential execution.
	 * 
	 * @param mX
	 * @param mV
	 * @param mW
	 * @param ret
	 * @param ct
	 * @param k
	 * @throws DMLRuntimeException
	 * @throws DMLUnsupportedOperationException
	 */
	public static void matrixMultChain(MatrixBlock mX, MatrixBlock mV, MatrixBlock mW, MatrixBlock ret, ChainType ct, int k) 
		throws DMLRuntimeException, DMLUnsupportedOperationException
	{		
		//check inputs / outputs (after that mV and mW guaranteed to be dense)
		if( mX.isEmptyBlock(false) || mV.isEmptyBlock(false) || (mW !=null && mW.isEmptyBlock(false)) )
			return;

		//check too high additional memory requirements (fallback to sequential)
		if( mV.rlen * 8 * k > 256*1024 ) { //256KB
			matrixMultChain(mX, mV, mW, ret, ct);
			return;
		}
		
		//Timing time = new Timing(true);
				
		//pre-processing
		ret.sparse = false;
		ret.allocateDenseBlock();
		
		//core matrix mult chain computation
		//(currently: always parallelization over number of rows)
		try {
			ExecutorService pool = Executors.newFixedThreadPool( k );
			ArrayList<MatrixMultChainTask> tasks = new ArrayList<MatrixMultChainTask>();
			int blklen = (int)(Math.ceil((double)mX.rlen/k));
			blklen += (blklen%24 != 0)?24-blklen%24:0;
			for( int i=0; i<k & i*blklen<mX.rlen; i++ )
				tasks.add(new MatrixMultChainTask(mX, mV, mW, ret, ct, i*blklen, Math.min((i+1)*blklen, mX.rlen)));
			pool.invokeAll(tasks);	
			pool.shutdown();
			//aggregate partial results
			for( MatrixMultChainTask task : tasks )
				vectAdd(task.getResult().denseBlock, ret.denseBlock, 0, 0, mX.clen);
		}
		catch(Exception ex) {
			throw new DMLRuntimeException(ex);
		}
		
		//post-processing
		ret.recomputeNonZeros();
		ret.examSparsity();
		
		//System.out.println("MMChain "+ct.toString()+" ("+mX.isInSparseFormat()+","+mX.getNumRows()+","+mX.getNumColumns()+","+mX.getNonZeros()+")x" +
		//		              "("+mV.isInSparseFormat()+","+mV.getNumRows()+","+mV.getNumColumns()+","+mV.getNonZeros()+") in "+time.stop());
	}

	/**
	 * 
	 * @param m1
	 * @param ret
	 * @param leftTranspose
	 * @throws DMLRuntimeException 
	 * @throws DMLUnsupportedOperationException 
	 */
	public static void matrixMultTransposeSelf( MatrixBlock m1, MatrixBlock ret, boolean leftTranspose )
		throws DMLUnsupportedOperationException, DMLRuntimeException
	{
		//Timing time = new Timing(true);
		
		if( m1.sparse )
			matrixMultTransposeSelfSparse(m1, ret, leftTranspose);
		else 
			matrixMultTransposeSelfDense(m1, ret, leftTranspose);

		//System.out.println("TSMM ("+m1.isInSparseFormat()+","+m1.getNumRows()+","+m1.getNumColumns()+","+m1.getNonZeros()+","+leftTranspose+") in "+time.stop());
	}
	
	/**
	 * 
	 * @param m1
	 * @param m2
	 * @param ret1
	 * @param ret2
	 * @throws DMLUnsupportedOperationException
	 * @throws DMLRuntimeException
	 */
	public static void matrixMultPermute( MatrixBlock pm1, MatrixBlock m2, MatrixBlock ret1, MatrixBlock ret2 )
		throws DMLUnsupportedOperationException, DMLRuntimeException
	{
		//Timing time = new Timing(true);
		
		if( m2.sparse )
			matrixMultPermuteSparse(pm1, m2, ret1, ret2);
		else if( ret1.sparse )
			matrixMultPermuteDenseSparse(pm1, m2, ret1, ret2);
		else 
			matrixMultPermuteDense(pm1, m2, ret1, ret2);
		
		//System.out.println("PMM ("+m1.isInSparseFormat()+","+m1.getNumRows()+","+m1.getNumColumns()+","+m1.getNonZeros()+")x" +
		//		              "("+m2.isInSparseFormat()+","+m2.getNumRows()+","+m2.getNumColumns()+","+m2.getNonZeros()+") in "+time.stop());
	}	
	
	
	//////////////////////////////////////////
	// optimized matrix mult implementation //
	//////////////////////////////////////////
	
	/**
	 * 
	 * @param m1
	 * @param m2
	 * @param ret
	 * @throws DMLRuntimeException 
	 */
	private static void matrixMultDenseDense(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret, int rl, int ru) 
		throws DMLRuntimeException
	{			
		double[] a = m1.denseBlock;
		double[] b = m2.denseBlock;
		double[] c = ret.denseBlock;
		final int m = m1.rlen;
		final int n = m2.clen;
		final int cd = m1.clen;

		if( LOW_LEVEL_OPTIMIZATION )
		{
			if( m==1 && n==1 ) 		   //DOT PRODUCT
			{
				c[0] = dotProduct(a, b, cd);
			}
			else if( n>1 && cd == 1 )  //OUTER PRODUCT
			{
				for( int i=rl, cix=rl*n; i < ru; i++, cix+=n) {
					if( a[i] == 1 )
						System.arraycopy(b, 0, c, cix, n);
				    else if( a[i] != 0 )
						vectMultiplyWrite(a[i], b, c, 0, cix, n);
					else
						Arrays.fill(c, cix, cix+n, 0);
				}
			}
			else if( n==1 && cd == 1 ) //VECTOR-SCALAR
			{
				vectMultiplyWrite(b[0], a, c, rl, rl, ru-rl);
			}
			else if( n==1 )            //MATRIX-VECTOR
			{
				for( int i=rl, aix=rl*cd; i < ru; i++, aix+=cd) 
					c[ i ] = dotProduct(a, b, aix, 0, cd);	
			}
			else //MATRIX-MATRIX
			{	
				//1) Unrolled inner loop (for better instruction-level parallelism)
				//2) Blocked execution (for less cache trashing in parallel exec) 	
				//3) Asymmetric block sizes (for less misses in inner loop, yet blocks in L1/L2)
				
				final int blocksizeI = 32; //64//256KB c block (typical L2 size per core), 32KB a block 
				final int blocksizeK = 24; //64//256KB b block (typical L2 size per core), used while read 512B of a / read/write 4KB of c 
				final int blocksizeJ = 1024; //512//4KB (typical main-memory page size), for scan 

				//temporary arrays (nnz a, b index)
				double[] ta = new double[ blocksizeK ];
				int[]  tbi  = new int[ blocksizeK ];
				
				//blocked execution
				for( int bi = rl; bi < ru; bi+=blocksizeI )
					for( int bk = 0, bimin = Math.min(ru, bi+blocksizeI); bk < cd; bk+=blocksizeK ) 
						for( int bj = 0, bkmin = Math.min(cd, bk+blocksizeK); bj < n; bj+=blocksizeJ ) 
						{
							int bklen = bkmin-bk;
							int bjlen = Math.min(n, bj+blocksizeJ)-bj;
							
							//core sub block matrix multiplication
				    		for( int i = bi; i < bimin; i++) 
				    		{
				    			int aixi = i * cd + bk; //start index on a
				    			int cixj = i * n + bj; //scan index on c
				    			
				    			//determine nnz of a (for sparsity-aware skipping of rows)
				    			int knnz = copyNonZeroElements(a, aixi, bk, bj, n, ta, tbi, bklen);
				    			//if( knnz > 0 ) //for skipping empty rows
				    			
			    				//rest not aligned to blocks of 4 rows
				    			final int bn = knnz % 4;
				    			switch( bn ){
					    			case 1: vectMultiplyAdd(ta[0], b, c, tbi[0], cixj, bjlen); break;
					    	    	case 2: vectMultiplyAdd2(ta[0],ta[1], b, c, tbi[0], tbi[1], cixj, bjlen); break;
					    			case 3: vectMultiplyAdd3(ta[0],ta[1],ta[2], b, c, tbi[0], tbi[1],tbi[2], cixj, bjlen); break;
				    			}
				    			
				    			//compute blocks of 4 rows (core inner loop)
				    			for( int k = bn; k<knnz; k+=4 ){
				    				vectMultiplyAdd4( ta[k], ta[k+1], ta[k+2], ta[k+3], b, c, 
				    						          tbi[k], tbi[k+1], tbi[k+2], tbi[k+3], cixj, bjlen );
				    			}
				    		}
						}
			}
		}
		else
		{
			//init empty result
			Arrays.fill(c, 0, c.length, 0);
			
			double val;
			for( int i = rl, aix=rl*cd, cix=rl*n; i < ru; i++, cix+=n) 
				for( int k = 0, bix=0; k < cd; k++, aix++, bix+=n)
				{			
					val = a[ aix ];
					if( val != 0 )
						for( int j = 0; j < n; j++) 
							c[ cix+j ] += val * b[ bix+j ];
				}	
		}
		
	}
	
	/**
	 * 
	 * @param m1
	 * @param m2
	 * @param ret
	 * @throws DMLRuntimeException 
	 */
	private static void matrixMultDenseSparse(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret, int rl, int ru) 
		throws DMLRuntimeException 
	{	
		double[] a = m1.denseBlock;
		double[] c = ret.denseBlock;
		int n = m2.clen;
		int cd = m1.clen;
	
		// MATRIX-MATRIX (VV, MV not applicable here because V always dense)
		if( LOW_LEVEL_OPTIMIZATION )
		{
			final int blocksizeI = 32; //256KB c block (typical L2 size per core), 32KB a block 
			final int blocksizeK = 32; 
			//note: in contrast to dense-dense, no blocking over j (would require maintaining blocksizeK indexes, counter-productive on skew)
			
			SparseRow[] b = m2.sparseRows;
			
			//blocked execution
			for( int bi = rl; bi < ru; bi+=blocksizeI )
				for( int bk = 0, bimin = Math.min(ru, bi+blocksizeI); bk < cd; bk+=blocksizeK ) 
				{
					int bklen = Math.min(cd, bk+blocksizeK)-bk;
					
					//core sub block matrix multiplication
		    		for( int i = bi; i < bimin; i++) 
		    		{
		    			int aixi = i * cd + bk; //start index on a
		    			int cixj = i * n + 0; //scan index on c
		    			
		    			for( int k = 0; k < bklen; k++ )
						{
							double val = a[aixi+k];
							SparseRow brow = b[ bk+k ];
							if( val != 0 && brow != null && !brow.isEmpty() ) {
								int blen = brow.size();
								int[] bix = brow.getIndexContainer();
								double[] bvals = brow.getValueContainer();								
								vectMultiplyAdd(val, bvals, c, bix, cixj, blen);
							}
						}
		    		}
				}	
		}
		else
		{
			for( int i=rl, aix=rl*cd, cix=rl*n; i < ru; i++, cix+=n ) 
				for(int k = 0; k < cd; k++, aix++ ) 
				{
					double val = a[aix];
					if( val!=0 )
					{
						SparseRow brow = m2.sparseRows[ k ];
						if( brow != null && !brow.isEmpty() ) 
						{
							int blen = brow.size();
							int[] bix = brow.getIndexContainer();
							double[] bvals = brow.getValueContainer();	
							for(int j = 0; j < blen; j++)
								c[cix+bix[j]] += val * bvals[j];								
						}
					}
				}		
		}
	}
	
	/**
	 * 
	 * @param m1
	 * @param m2
	 * @param ret
	 * @throws DMLRuntimeException 
	 */
	private static void matrixMultSparseDense(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret, int rl, int ru) 
		throws DMLRuntimeException
	{	
		double[] b = m2.denseBlock;
		double[] c = ret.denseBlock;
		final int m = m1.rlen;
		final int n = m2.clen;

		if( LOW_LEVEL_OPTIMIZATION )
		{
		
			if( m==1 && n==1 ) //DOT PRODUCT
			{
				SparseRow arow = m1.sparseRows[0];
				if( arow != null && !arow.isEmpty() )
				{
					int alen = arow.size();
					int[] aix = arow.getIndexContainer();
					double[] avals = arow.getValueContainer();
					
					c[0] = dotProduct(avals, b, aix, 0, alen);
				}
			}
			else if( n==1 ) //MATRIX-VECTOR
			{
				for( int i=rl; i<Math.min(ru, m1.sparseRows.length); i++ )
				{
					SparseRow arow = m1.sparseRows[i];
					if( arow != null && !arow.isEmpty() ) 
					{
						int alen = arow.size();
						int[] aix = arow.getIndexContainer();
						double[] avals = arow.getValueContainer();					
					
						c[i] = dotProduct(avals, b, aix, 0, alen);							
					}
				}
			}
			else //MATRIX-MATRIX
			{
				for( int i=rl, cix=rl*n; i<Math.min(ru, m1.sparseRows.length); i++, cix+=n )
				{
					SparseRow arow = m1.sparseRows[i];
					if( arow != null && !arow.isEmpty() ) 
					{
						int alen = arow.size();
						int[] aix = arow.getIndexContainer();
						double[] avals = arow.getValueContainer();					
						
						if( alen==1 && avals[0]==1 ) //ROW SELECTION 
						{
							//plain memcopy for permutation matrices
							System.arraycopy(b, aix[0]*n, c, cix, n);
						}
						else //GENERAL CASE
						{
							//rest not aligned to blocks of 4 rows
			    			final int bn = alen % 4;
			    			switch( bn ){
				    			case 1: vectMultiplyAdd(avals[0], b, c, aix[0]*n, cix, n); break;
				    	    	case 2: vectMultiplyAdd2(avals[0],avals[1], b, c, aix[0]*n, aix[1]*n, cix, n); break;
				    			case 3: vectMultiplyAdd3(avals[0],avals[1],avals[2], b, c, aix[0]*n, aix[1]*n, aix[2]*n, cix, n); break;
			    			}
			    			
			    			//compute blocks of 4 rows (core inner loop)
			    			for( int k = bn; k<alen; k+=4 ) {
			    				vectMultiplyAdd4( avals[k], avals[k+1], avals[k+2], avals[k+3], b, c, 
			    						          aix[k]*n, aix[k+1]*n, aix[k+2]*n, aix[k+3]*n, cix, n );
			    			}
						}
					}
				}					
			}
		}
		else
		{
			//init empty result
			Arrays.fill(c, 0, c.length, 0);
			
			for( int i=rl, cix=rl*n; i<Math.min(ru, m1.sparseRows.length); i++, cix+=n )
			{
				SparseRow arow = m1.sparseRows[i];
				if( arow != null && !arow.isEmpty() ) 
				{
					int alen = arow.size();
					int[] aix = arow.getIndexContainer();
					double[] avals = arow.getValueContainer();					
					
					for(int k = 0; k < alen; k++) 
					{
						double val = avals[k];
						for(int j = 0, bix=aix[k]*n; j < n; j++)
							c[cix+j] += val * b[bix+j];								
					}						
				}
			}
		}
	}
	
	/**
	 * 
	 * @param m1
	 * @param m2
	 * @param ret
	 * @throws DMLRuntimeException 
	 */
	private static void matrixMultSparseSparse(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret, int rl, int ru) 
		throws DMLRuntimeException
	{	
		double[] c = ret.denseBlock;
		int n = m2.clen;
		
		// MATRIX-MATRIX (VV, MV not applicable here because V always dense)
		if(LOW_LEVEL_OPTIMIZATION)
		{
			for( int i=rl, cix=rl*n; i<Math.min(ru, m1.sparseRows.length); i++, cix+=n )
			{
				SparseRow arow = m1.sparseRows[i];
				if( arow != null && !arow.isEmpty() ) 
				{
					int alen = arow.size();
					int[] aix = arow.getIndexContainer();
					double[] avals = arow.getValueContainer();					
					
					for(int k = 0; k < alen; k++) 
					{
						double val = avals[k];
						SparseRow brow = m2.sparseRows[ aix[k] ];
						if( brow != null && !brow.isEmpty() ) 
						{
							int blen = brow.size();
							int[] bix = brow.getIndexContainer();
							double[] bvals = brow.getValueContainer();	
							
							vectMultiplyAdd(val, bvals, c, bix, cix, blen);
						}
					}						
				}
			}
		}
		else
		{
			for( int i=rl, cix=rl*n; i<Math.min(ru, m1.sparseRows.length); i++, cix+=n )
			{
				SparseRow arow = m1.sparseRows[i];
				if( arow != null && !arow.isEmpty() ) 
				{
					int alen = arow.size();
					int[] aix = arow.getIndexContainer();
					double[] avals = arow.getValueContainer();					
					
					for(int k = 0; k < alen; k++) 
					{
						double val = avals[k];
						SparseRow brow = m2.sparseRows[ aix[k] ];
						if( brow != null && !brow.isEmpty() ) 
						{
							int blen = brow.size();
							int[] bix = brow.getIndexContainer();
							double[] bvals = brow.getValueContainer();	
							for(int j = 0; j < blen; j++)
								c[cix+bix[j]] += val * bvals[j];								
						}
					}						
				}
			}
		}
	}

	/**
	 * This implementation applies to any combination of dense/sparse if at least one
	 * input is ultrasparse (sparse and very few nnz). In that case, most importantly,
	 * we want to create a sparse output and only iterate over the few nnz as the major
	 * dimension. Low-level optimization have less importance in that case and having
	 * this generic implementation helps to reduce the implementations from (2+1)^2
	 * to 2^2+1.
	 * 
	 * @param m1
	 * @param m2
	 * @param ret
	 * @throws DMLRuntimeException
	 */
	private static void matrixMultUltraSparse(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret, int rl, int ru) 
		throws DMLRuntimeException 
	{
		boolean leftUS = m1.isUltraSparse();
		final int m  = m1.rlen;
		final int cd = m1.clen;
		final int n  = m2.clen;
		
		if( leftUS ) //left is ultra-sparse (IKJ)
		{
			boolean rightSparse = m2.sparse;
			
			for( int i=rl; i<ru; i++ )
			{
				SparseRow arow = m1.sparseRows[ i ];
				if( arow != null && !arow.isEmpty() ) 
				{
					int alen = arow.size();
					int[] aixs = arow.getIndexContainer();
					double[] avals = arow.getValueContainer();	
					
					if( alen==1 && avals[0]==1 ) //ROW SELECTION (no aggregation)
					{
						int aix = aixs[0];
						if( rightSparse ) { //sparse right matrix (full row copy)
							if( m2.sparseRows!=null && m2.sparseRows[aix]!=null ) {
								ret.rlen=m;
								ret.allocateSparseRowsBlock(false); //allocation on demand
								ret.sparseRows[i] = new SparseRow(m2.sparseRows[aix]); 
								ret.nonZeros += ret.sparseRows[i].size();
							}
						}
						else { //dense right matrix (append all values)
							for( int j=0; j<n; j++ )
								ret.appendValue(i, j, m2.quickGetValue(aix, j));
						}
					}
					else //GENERAL CASE
					{
						for( int k=0; k<alen; k++ )
						{
							double aval = avals[k];
							int aix = aixs[k];
							for( int j=0; j<n; j++ )
							{
								double cval = ret.quickGetValue(i, j);
								double cvald = aval*m2.quickGetValue(aix, j);
								if( cvald != 0 )
									ret.quickSetValue(i, j, cval+cvald);
							}
						}
					}
				}
			}
		}
		else //right is ultra-sparse (KJI)
		{
			for(int k = 0; k < cd; k++ ) 
			{			
				SparseRow brow = m2.sparseRows[ k ];
				if( brow != null && !brow.isEmpty() ) 
				{
					int blen = brow.size();
					int[] bixs = brow.getIndexContainer();
					double[] bvals = brow.getValueContainer();								
					for( int j=0; j<blen; j++ )
					{
						double bval = bvals[j];
						int bix = bixs[j];
						for( int i=rl; i<ru; i++ )
						{
							double cvald = bval*m1.quickGetValue(i, k);
							if( cvald != 0 ){
								double cval = ret.quickGetValue(i, bix);
								ret.quickSetValue(i, bix, cval+cvald);
							}
						}
					}
				}
			}
		}
		//no need to recompute nonzeros because maintained internally
	}

	/**
	 * 
	 * @param mX
	 * @param mV
	 * @param mW
	 * @param ret
	 * @param ct
	 * @throws DMLRuntimeException
	 * @throws DMLUnsupportedOperationException
	 */
	private static void matrixMultChainDense(MatrixBlock mX, MatrixBlock mV, MatrixBlock mW, MatrixBlock ret, ChainType ct, int rl, int ru) 
	{
		double[] a = mX.denseBlock;
		double[] b = mV.denseBlock;
		double[] w = (mW!=null) ? mW.denseBlock : null;
		double[] c = ret.denseBlock;
		final int cd = mX.clen; //features in X
		boolean weights = (ct == ChainType.XtwXv);

		//temporary array for cache blocking
		//(blocksize chosen to fit b+v in L2 (256KB) for default 1k blocks)
		final int blocksize = 24; // constraint: factor of 4
		double[] tmp = new double[blocksize];
			
		//blockwise mmchain computation
		final int bn = ru - ru % blocksize; //rl blocksize aligned
		for( int bi=rl; bi < bn; bi+=blocksize ) 
		{
			//compute 1st matrix-vector for row block
			for( int j=0, aix=bi*cd; j < blocksize; j++, aix+=cd)
				tmp[j] = dotProduct(a, b, aix, 0, cd);
			
			//multiply weights (in-place), if required
			if( weights ) 
				vectMultiply(w, tmp, bi, 0, blocksize);	
			
			//compute 2nd matrix vector for row block and aggregate
			for (int j=0, aix=bi*cd; j < blocksize; j+=4, aix+=4*cd)
				vectMultiplyAdd4(tmp[j], tmp[j+1], tmp[j+2], tmp[j+3], a, c, aix, aix+cd, aix+2*cd, aix+3*cd, 0, cd);
		}
		
		//compute rest (not aligned to blocksize)
		for( int i=bn, aix=i*cd; i < ru; i++, aix+=cd ) {
			double val = dotProduct(a, b, aix, 0, cd);
			val *= (weights) ? w[i] : 1; 
			vectMultiplyAdd(val, a, c, aix, 0, cd);				
		}
	}
	<commit_msg>94449: Scale-up CP Runtime - Fix multi-threaded matrix multiply (pre-allocation/nnz recompute sparse outputs)<commit_after>/**
 * IBM Confidential
 * OCO Source Materials
 * (C) Copyright IBM Corp. 2010, 2015
 * The source code for this program is not published or otherwise divested of its trade secrets, irrespective of what has been deposited with the U.S. Copyright Office.
 */


package com.ibm.bi.dml.runtime.matrix.data;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import com.ibm.bi.dml.lops.MapMultChain.ChainType;
import com.ibm.bi.dml.runtime.DMLRuntimeException;
import com.ibm.bi.dml.runtime.DMLUnsupportedOperationException;
import com.ibm.bi.dml.runtime.functionobjects.SwapIndex;
import com.ibm.bi.dml.runtime.matrix.operators.ReorgOperator;
import com.ibm.bi.dml.runtime.util.UtilFunctions;

/**
 * MB:
 * Library for matrix multiplications including MM, MV, VV for all
 * combinations of dense and sparse representations and special
 * operations such as transpose-self matrix multiplication.
 * 
 * In general all implementations use internally dense outputs
 * for direct access, but change the final result to sparse if necessary.  
 * 
 * NOTES on BLAS:
 * * Experiments in 04/2013 showed that even on dense-dense this implementation 
 *   is 3x faster than f2j-BLAS-DGEMM, 2x faster than f2c-BLAS-DGEMM, and
 *   level (+10% after JIT) with a native C implementation. 
 * * Calling native BLAS would loose platform independence and would require 
 *   JNI calls incl data transfer. Furthermore, BLAS does not support sparse 
 *   matrices (except Sparse BLAS, with dedicated function calls and matrix formats) 
 *   and would be an external dependency. 
 * * Experiments in 02/2014 showed that on dense-dense this implementation now achieves
 *   almost 30% peak FP performance. Compared to Intel MKL 11.1 (dgemm, N=1000) it is
 *   just 3.2x (sparsity=1.0) and 1.9x (sparsity=0.5) slower, respectively.
 *   
 *  
 */
public class LibMatrixMult 
{
	@SuppressWarnings("unused")
	private static final String _COPYRIGHT = "Licensed Materials - Property of IBM\n(C) Copyright IBM Corp. 2010, 2015\n" +
                                             "US Government Users Restricted Rights - Use, duplication  disclosure restricted by GSA ADP Schedule Contract with IBM Corp.";
	
	public static final boolean LOW_LEVEL_OPTIMIZATION = true;

	private LibMatrixMult() {
		//prevent instantiation via private constructor
	}
	
	////////////////////////////////
	// public matrix mult interface
	////////////////////////////////
	
	/**
	 * Performs a matrix multiplication and stores the result in the output matrix.
	 * 
	 * All variants use a IKJ access pattern, and internally use dense output. After the
	 * actual computation, we recompute nnz and check for sparse/dense representation.
	 *  
	 * 
	 * @param m1 first matrix
	 * @param m2 second matrix
	 * @param ret result matrix
	 * @throws DMLRuntimeException 
	 */
	public static void matrixMult(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret) 
		throws DMLRuntimeException
	{	
		//check inputs / outputs
		if( m1.isEmptyBlock(false) || m2.isEmptyBlock(false) )
			return;
		
		//Timing time = new Timing(true);
		
		//pre-processing: output allocation
		ret.sparse = (m1.isUltraSparse() || m2.isUltraSparse());
		if( !ret.sparse )
			ret.allocateDenseBlock();
		
		//core matrix mult computation
		if( m1.isUltraSparse() || m2.isUltraSparse() )
			matrixMultUltraSparse(m1, m2, ret, 0, m1.rlen);
		else if(!m1.sparse && !m2.sparse)
			matrixMultDenseDense(m1, m2, ret, 0, m1.rlen);
		else if(m1.sparse && m2.sparse)
			matrixMultSparseSparse(m1, m2, ret, 0, m1.rlen);
		else if(m1.sparse)
			matrixMultSparseDense(m1, m2, ret, 0, m1.rlen);
		else
			matrixMultDenseSparse(m1, m2, ret, 0, m1.rlen);
		
		//post-processing: nnz/representation
		if( !ret.sparse )
			ret.recomputeNonZeros();
		ret.examSparsity();
		
		//System.out.println("MM ("+m1.isInSparseFormat()+","+m1.getNumRows()+","+m1.getNumColumns()+","+m1.getNonZeros()+")x" +
		//		              "("+m2.isInSparseFormat()+","+m2.getNumRows()+","+m2.getNumColumns()+","+m2.getNonZeros()+") in "+time.stop());
	}
	
	/**
	 * Performs a multi-threaded matrix multiplication and stores the result in the output matrix.
	 * The parameter k (k>=1) determines the max parallelism k' with k'=min(k, vcores, m1.rlen).
	 * 
	 * @param m1
	 * @param m2
	 * @param ret
	 * @param k
	 * @throws DMLRuntimeException 
	 */
	public static void matrixMult(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret, int k) 
		throws DMLRuntimeException
	{	
		//check inputs / outputs
		if( m1.isEmptyBlock(false) || m2.isEmptyBlock(false) )
			return;
		//check no parallelization benefit (fallback to sequential)
		if( m1.rlen == 1 ) {
			matrixMult(m1, m2, ret);
			return;
		}
		
		//Timing time = new Timing(true);
		
		//pre-processing: output allocation (in contrast to single-threaded,
		//we need to allocate sparse as well in order to prevent synchronization)
		ret.sparse = (m1.isUltraSparse() || m2.isUltraSparse());
		if( !ret.sparse )
			ret.allocateDenseBlock();
		else
			ret.allocateSparseRowsBlock();
		
		//core multi-threaded matrix mult computation
		//(currently: always parallelization over number of rows)
		try {
			ExecutorService pool = Executors.newFixedThreadPool( k );
			ArrayList<MatrixMultTask> tasks = new ArrayList<MatrixMultTask>();
			int blklen = (int)(Math.ceil((double)m1.rlen/k));
			for( int i=0; i<k & i*blklen<m1.rlen; i++ )
				tasks.add(new MatrixMultTask(m1, m2, ret, i*blklen, Math.min((i+1)*blklen, m1.rlen)));
			pool.invokeAll(tasks);	
			pool.shutdown();
		}
		catch(Exception ex) {
			throw new DMLRuntimeException(ex);
		}
		
		//post-processing: nnz/representation (in contrast to single-threaded,
		//we need to recompute nnz for sparse as well in order to prevent synchronization)
		ret.recomputeNonZeros();
		ret.examSparsity();
		
		//System.out.println("MM k="+k+" ("+m1.isInSparseFormat()+","+m1.getNumRows()+","+m1.getNumColumns()+","+m1.getNonZeros()+")x" +
		//		              "("+m2.isInSparseFormat()+","+m2.getNumRows()+","+m2.getNumColumns()+","+m2.getNonZeros()+") in "+time.stop());
	}
	
	/**
	 * Performs a matrix multiplication chain operation of type t(X)%*%(X%*%v) or t(X)%*%(w*(X%*%v)).
	 * 
	 * All variants use a IKJ access pattern, and internally use dense output. After the
	 * actual computation, we recompute nnz and check for sparse/dense representation.
	 * 
	 * @param m1
	 * @param m2
	 * @param w
	 * @param ret
	 * @param ct
	 * @throws DMLRuntimeException
	 * @throws DMLUnsupportedOperationException
	 */
	public static void matrixMultChain(MatrixBlock mX, MatrixBlock mV, MatrixBlock mW, MatrixBlock ret, ChainType ct) 
		throws DMLRuntimeException, DMLUnsupportedOperationException
	{		
		//check inputs / outputs (after that mV and mW guaranteed to be dense)
		if( mX.isEmptyBlock(false) || mV.isEmptyBlock(false) || (mW !=null && mW.isEmptyBlock(false)) )
			return;

		//Timing time = new Timing(true);
				
		//pre-processing
		ret.sparse = false;
		ret.allocateDenseBlock();
		
		//core matrix mult chain computation
		if( mX.sparse )
			matrixMultChainSparse(mX, mV, mW, ret, ct, 0, mX.rlen);
		else
			matrixMultChainDense(mX, mV, mW, ret, ct, 0, mX.rlen);
		
		//post-processing
		ret.recomputeNonZeros();
		ret.examSparsity();
		
		//System.out.println("MMChain "+ct.toString()+" ("+mX.isInSparseFormat()+","+mX.getNumRows()+","+mX.getNumColumns()+","+mX.getNonZeros()+")x" +
		//		              "("+mV.isInSparseFormat()+","+mV.getNumRows()+","+mV.getNumColumns()+","+mV.getNonZeros()+") in "+time.stop());
	}

	/**
	 * Performs a parallel matrix multiplication chain operation of type t(X)%*%(X%*%v) or t(X)%*%(w*(X%*%v)).
	 * The parameter k (k>=1) determines the max parallelism k' with k'=min(k, vcores, m1.rlen).
	 * 
	 * NOTE: This multi-threaded mmchain operation has additional memory requirements of k*ncol(X)*8bytes 
	 * for partial aggregation. Current max memory: 256KB; otherwise redirectly to sequential execution.
	 * 
	 * @param mX
	 * @param mV
	 * @param mW
	 * @param ret
	 * @param ct
	 * @param k
	 * @throws DMLRuntimeException
	 * @throws DMLUnsupportedOperationException
	 */
	public static void matrixMultChain(MatrixBlock mX, MatrixBlock mV, MatrixBlock mW, MatrixBlock ret, ChainType ct, int k) 
		throws DMLRuntimeException, DMLUnsupportedOperationException
	{		
		//check inputs / outputs (after that mV and mW guaranteed to be dense)
		if( mX.isEmptyBlock(false) || mV.isEmptyBlock(false) || (mW !=null && mW.isEmptyBlock(false)) )
			return;

		//check too high additional memory requirements (fallback to sequential)
		if( mV.rlen * 8 * k > 256*1024 ) { //256KB
			matrixMultChain(mX, mV, mW, ret, ct);
			return;
		}
		
		//Timing time = new Timing(true);
				
		//pre-processing
		ret.sparse = false;
		ret.allocateDenseBlock();
		
		//core matrix mult chain computation
		//(currently: always parallelization over number of rows)
		try {
			ExecutorService pool = Executors.newFixedThreadPool( k );
			ArrayList<MatrixMultChainTask> tasks = new ArrayList<MatrixMultChainTask>();
			int blklen = (int)(Math.ceil((double)mX.rlen/k));
			blklen += (blklen%24 != 0)?24-blklen%24:0;
			for( int i=0; i<k & i*blklen<mX.rlen; i++ )
				tasks.add(new MatrixMultChainTask(mX, mV, mW, ret, ct, i*blklen, Math.min((i+1)*blklen, mX.rlen)));
			pool.invokeAll(tasks);	
			pool.shutdown();
			//aggregate partial results
			for( MatrixMultChainTask task : tasks )
				vectAdd(task.getResult().denseBlock, ret.denseBlock, 0, 0, mX.clen);
		}
		catch(Exception ex) {
			throw new DMLRuntimeException(ex);
		}
		
		//post-processing
		ret.recomputeNonZeros();
		ret.examSparsity();
		
		//System.out.println("MMChain "+ct.toString()+" k="+k+" ("+mX.isInSparseFormat()+","+mX.getNumRows()+","+mX.getNumColumns()+","+mX.getNonZeros()+")x" +
		//		              "("+mV.isInSparseFormat()+","+mV.getNumRows()+","+mV.getNumColumns()+","+mV.getNonZeros()+") in "+time.stop());
	}

	/**
	 * 
	 * @param m1
	 * @param ret
	 * @param leftTranspose
	 * @throws DMLRuntimeException 
	 * @throws DMLUnsupportedOperationException 
	 */
	public static void matrixMultTransposeSelf( MatrixBlock m1, MatrixBlock ret, boolean leftTranspose )
		throws DMLUnsupportedOperationException, DMLRuntimeException
	{
		//Timing time = new Timing(true);
		
		if( m1.sparse )
			matrixMultTransposeSelfSparse(m1, ret, leftTranspose);
		else 
			matrixMultTransposeSelfDense(m1, ret, leftTranspose);

		//System.out.println("TSMM ("+m1.isInSparseFormat()+","+m1.getNumRows()+","+m1.getNumColumns()+","+m1.getNonZeros()+","+leftTranspose+") in "+time.stop());
	}
	
	/**
	 * 
	 * @param m1
	 * @param m2
	 * @param ret1
	 * @param ret2
	 * @throws DMLUnsupportedOperationException
	 * @throws DMLRuntimeException
	 */
	public static void matrixMultPermute( MatrixBlock pm1, MatrixBlock m2, MatrixBlock ret1, MatrixBlock ret2 )
		throws DMLUnsupportedOperationException, DMLRuntimeException
	{
		//Timing time = new Timing(true);
		
		if( m2.sparse )
			matrixMultPermuteSparse(pm1, m2, ret1, ret2);
		else if( ret1.sparse )
			matrixMultPermuteDenseSparse(pm1, m2, ret1, ret2);
		else 
			matrixMultPermuteDense(pm1, m2, ret1, ret2);
		
		//System.out.println("PMM ("+m1.isInSparseFormat()+","+m1.getNumRows()+","+m1.getNumColumns()+","+m1.getNonZeros()+")x" +
		//		              "("+m2.isInSparseFormat()+","+m2.getNumRows()+","+m2.getNumColumns()+","+m2.getNonZeros()+") in "+time.stop());
	}	
	
	
	//////////////////////////////////////////
	// optimized matrix mult implementation //
	//////////////////////////////////////////
	
	/**
	 * 
	 * @param m1
	 * @param m2
	 * @param ret
	 * @throws DMLRuntimeException 
	 */
	private static void matrixMultDenseDense(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret, int rl, int ru) 
		throws DMLRuntimeException
	{			
		double[] a = m1.denseBlock;
		double[] b = m2.denseBlock;
		double[] c = ret.denseBlock;
		final int m = m1.rlen;
		final int n = m2.clen;
		final int cd = m1.clen;

		if( LOW_LEVEL_OPTIMIZATION )
		{
			if( m==1 && n==1 ) 		   //DOT PRODUCT
			{
				c[0] = dotProduct(a, b, cd);
			}
			else if( n>1 && cd == 1 )  //OUTER PRODUCT
			{
				for( int i=rl, cix=rl*n; i < ru; i++, cix+=n) {
					if( a[i] == 1 )
						System.arraycopy(b, 0, c, cix, n);
				    else if( a[i] != 0 )
						vectMultiplyWrite(a[i], b, c, 0, cix, n);
					else
						Arrays.fill(c, cix, cix+n, 0);
				}
			}
			else if( n==1 && cd == 1 ) //VECTOR-SCALAR
			{
				vectMultiplyWrite(b[0], a, c, rl, rl, ru-rl);
			}
			else if( n==1 )            //MATRIX-VECTOR
			{
				for( int i=rl, aix=rl*cd; i < ru; i++, aix+=cd) 
					c[ i ] = dotProduct(a, b, aix, 0, cd);	
			}
			else //MATRIX-MATRIX
			{	
				//1) Unrolled inner loop (for better instruction-level parallelism)
				//2) Blocked execution (for less cache trashing in parallel exec) 	
				//3) Asymmetric block sizes (for less misses in inner loop, yet blocks in L1/L2)
				
				final int blocksizeI = 32; //64//256KB c block (typical L2 size per core), 32KB a block 
				final int blocksizeK = 24; //64//256KB b block (typical L2 size per core), used while read 512B of a / read/write 4KB of c 
				final int blocksizeJ = 1024; //512//4KB (typical main-memory page size), for scan 

				//temporary arrays (nnz a, b index)
				double[] ta = new double[ blocksizeK ];
				int[]  tbi  = new int[ blocksizeK ];
				
				//blocked execution
				for( int bi = rl; bi < ru; bi+=blocksizeI )
					for( int bk = 0, bimin = Math.min(ru, bi+blocksizeI); bk < cd; bk+=blocksizeK ) 
						for( int bj = 0, bkmin = Math.min(cd, bk+blocksizeK); bj < n; bj+=blocksizeJ ) 
						{
							int bklen = bkmin-bk;
							int bjlen = Math.min(n, bj+blocksizeJ)-bj;
							
							//core sub block matrix multiplication
				    		for( int i = bi; i < bimin; i++) 
				    		{
				    			int aixi = i * cd + bk; //start index on a
				    			int cixj = i * n + bj; //scan index on c
				    			
				    			//determine nnz of a (for sparsity-aware skipping of rows)
				    			int knnz = copyNonZeroElements(a, aixi, bk, bj, n, ta, tbi, bklen);
				    			//if( knnz > 0 ) //for skipping empty rows
				    			
			    				//rest not aligned to blocks of 4 rows
				    			final int bn = knnz % 4;
				    			switch( bn ){
					    			case 1: vectMultiplyAdd(ta[0], b, c, tbi[0], cixj, bjlen); break;
					    	    	case 2: vectMultiplyAdd2(ta[0],ta[1], b, c, tbi[0], tbi[1], cixj, bjlen); break;
					    			case 3: vectMultiplyAdd3(ta[0],ta[1],ta[2], b, c, tbi[0], tbi[1],tbi[2], cixj, bjlen); break;
				    			}
				    			
				    			//compute blocks of 4 rows (core inner loop)
				    			for( int k = bn; k<knnz; k+=4 ){
				    				vectMultiplyAdd4( ta[k], ta[k+1], ta[k+2], ta[k+3], b, c, 
				    						          tbi[k], tbi[k+1], tbi[k+2], tbi[k+3], cixj, bjlen );
				    			}
				    		}
						}
			}
		}
		else
		{
			//init empty result
			Arrays.fill(c, 0, c.length, 0);
			
			double val;
			for( int i = rl, aix=rl*cd, cix=rl*n; i < ru; i++, cix+=n) 
				for( int k = 0, bix=0; k < cd; k++, aix++, bix+=n)
				{			
					val = a[ aix ];
					if( val != 0 )
						for( int j = 0; j < n; j++) 
							c[ cix+j ] += val * b[ bix+j ];
				}	
		}
		
	}
	
	/**
	 * 
	 * @param m1
	 * @param m2
	 * @param ret
	 * @throws DMLRuntimeException 
	 */
	private static void matrixMultDenseSparse(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret, int rl, int ru) 
		throws DMLRuntimeException 
	{	
		double[] a = m1.denseBlock;
		double[] c = ret.denseBlock;
		int n = m2.clen;
		int cd = m1.clen;
	
		// MATRIX-MATRIX (VV, MV not applicable here because V always dense)
		if( LOW_LEVEL_OPTIMIZATION )
		{
			final int blocksizeI = 32; //256KB c block (typical L2 size per core), 32KB a block 
			final int blocksizeK = 32; 
			//note: in contrast to dense-dense, no blocking over j (would require maintaining blocksizeK indexes, counter-productive on skew)
			
			SparseRow[] b = m2.sparseRows;
			
			//blocked execution
			for( int bi = rl; bi < ru; bi+=blocksizeI )
				for( int bk = 0, bimin = Math.min(ru, bi+blocksizeI); bk < cd; bk+=blocksizeK ) 
				{
					int bklen = Math.min(cd, bk+blocksizeK)-bk;
					
					//core sub block matrix multiplication
		    		for( int i = bi; i < bimin; i++) 
		    		{
		    			int aixi = i * cd + bk; //start index on a
		    			int cixj = i * n + 0; //scan index on c
		    			
		    			for( int k = 0; k < bklen; k++ )
						{
							double val = a[aixi+k];
							SparseRow brow = b[ bk+k ];
							if( val != 0 && brow != null && !brow.isEmpty() ) {
								int blen = brow.size();
								int[] bix = brow.getIndexContainer();
								double[] bvals = brow.getValueContainer();								
								vectMultiplyAdd(val, bvals, c, bix, cixj, blen);
							}
						}
		    		}
				}	
		}
		else
		{
			for( int i=rl, aix=rl*cd, cix=rl*n; i < ru; i++, cix+=n ) 
				for(int k = 0; k < cd; k++, aix++ ) 
				{
					double val = a[aix];
					if( val!=0 )
					{
						SparseRow brow = m2.sparseRows[ k ];
						if( brow != null && !brow.isEmpty() ) 
						{
							int blen = brow.size();
							int[] bix = brow.getIndexContainer();
							double[] bvals = brow.getValueContainer();	
							for(int j = 0; j < blen; j++)
								c[cix+bix[j]] += val * bvals[j];								
						}
					}
				}		
		}
	}
	
	/**
	 * 
	 * @param m1
	 * @param m2
	 * @param ret
	 * @throws DMLRuntimeException 
	 */
	private static void matrixMultSparseDense(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret, int rl, int ru) 
		throws DMLRuntimeException
	{	
		double[] b = m2.denseBlock;
		double[] c = ret.denseBlock;
		final int m = m1.rlen;
		final int n = m2.clen;

		if( LOW_LEVEL_OPTIMIZATION )
		{
		
			if( m==1 && n==1 ) //DOT PRODUCT
			{
				SparseRow arow = m1.sparseRows[0];
				if( arow != null && !arow.isEmpty() )
				{
					int alen = arow.size();
					int[] aix = arow.getIndexContainer();
					double[] avals = arow.getValueContainer();
					
					c[0] = dotProduct(avals, b, aix, 0, alen);
				}
			}
			else if( n==1 ) //MATRIX-VECTOR
			{
				for( int i=rl; i<Math.min(ru, m1.sparseRows.length); i++ )
				{
					SparseRow arow = m1.sparseRows[i];
					if( arow != null && !arow.isEmpty() ) 
					{
						int alen = arow.size();
						int[] aix = arow.getIndexContainer();
						double[] avals = arow.getValueContainer();					
					
						c[i] = dotProduct(avals, b, aix, 0, alen);							
					}
				}
			}
			else //MATRIX-MATRIX
			{
				for( int i=rl, cix=rl*n; i<Math.min(ru, m1.sparseRows.length); i++, cix+=n )
				{
					SparseRow arow = m1.sparseRows[i];
					if( arow != null && !arow.isEmpty() ) 
					{
						int alen = arow.size();
						int[] aix = arow.getIndexContainer();
						double[] avals = arow.getValueContainer();					
						
						if( alen==1 && avals[0]==1 ) //ROW SELECTION 
						{
							//plain memcopy for permutation matrices
							System.arraycopy(b, aix[0]*n, c, cix, n);
						}
						else //GENERAL CASE
						{
							//rest not aligned to blocks of 4 rows
			    			final int bn = alen % 4;
			    			switch( bn ){
				    			case 1: vectMultiplyAdd(avals[0], b, c, aix[0]*n, cix, n); break;
				    	    	case 2: vectMultiplyAdd2(avals[0],avals[1], b, c, aix[0]*n, aix[1]*n, cix, n); break;
				    			case 3: vectMultiplyAdd3(avals[0],avals[1],avals[2], b, c, aix[0]*n, aix[1]*n, aix[2]*n, cix, n); break;
			    			}
			    			
			    			//compute blocks of 4 rows (core inner loop)
			    			for( int k = bn; k<alen; k+=4 ) {
			    				vectMultiplyAdd4( avals[k], avals[k+1], avals[k+2], avals[k+3], b, c, 
			    						          aix[k]*n, aix[k+1]*n, aix[k+2]*n, aix[k+3]*n, cix, n );
			    			}
						}
					}
				}					
			}
		}
		else
		{
			//init empty result
			Arrays.fill(c, 0, c.length, 0);
			
			for( int i=rl, cix=rl*n; i<Math.min(ru, m1.sparseRows.length); i++, cix+=n )
			{
				SparseRow arow = m1.sparseRows[i];
				if( arow != null && !arow.isEmpty() ) 
				{
					int alen = arow.size();
					int[] aix = arow.getIndexContainer();
					double[] avals = arow.getValueContainer();					
					
					for(int k = 0; k < alen; k++) 
					{
						double val = avals[k];
						for(int j = 0, bix=aix[k]*n; j < n; j++)
							c[cix+j] += val * b[bix+j];								
					}						
				}
			}
		}
	}
	
	/**
	 * 
	 * @param m1
	 * @param m2
	 * @param ret
	 * @throws DMLRuntimeException 
	 */
	private static void matrixMultSparseSparse(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret, int rl, int ru) 
		throws DMLRuntimeException
	{	
		double[] c = ret.denseBlock;
		int n = m2.clen;
		
		// MATRIX-MATRIX (VV, MV not applicable here because V always dense)
		if(LOW_LEVEL_OPTIMIZATION)
		{
			for( int i=rl, cix=rl*n; i<Math.min(ru, m1.sparseRows.length); i++, cix+=n )
			{
				SparseRow arow = m1.sparseRows[i];
				if( arow != null && !arow.isEmpty() ) 
				{
					int alen = arow.size();
					int[] aix = arow.getIndexContainer();
					double[] avals = arow.getValueContainer();					
					
					for(int k = 0; k < alen; k++) 
					{
						double val = avals[k];
						SparseRow brow = m2.sparseRows[ aix[k] ];
						if( brow != null && !brow.isEmpty() ) 
						{
							int blen = brow.size();
							int[] bix = brow.getIndexContainer();
							double[] bvals = brow.getValueContainer();	
							
							vectMultiplyAdd(val, bvals, c, bix, cix, blen);
						}
					}						
				}
			}
		}
		else
		{
			for( int i=rl, cix=rl*n; i<Math.min(ru, m1.sparseRows.length); i++, cix+=n )
			{
				SparseRow arow = m1.sparseRows[i];
				if( arow != null && !arow.isEmpty() ) 
				{
					int alen = arow.size();
					int[] aix = arow.getIndexContainer();
					double[] avals = arow.getValueContainer();					
					
					for(int k = 0; k < alen; k++) 
					{
						double val = avals[k];
						SparseRow brow = m2.sparseRows[ aix[k] ];
						if( brow != null && !brow.isEmpty() ) 
						{
							int blen = brow.size();
							int[] bix = brow.getIndexContainer();
							double[] bvals = brow.getValueContainer();	
							for(int j = 0; j < blen; j++)
								c[cix+bix[j]] += val * bvals[j];								
						}
					}						
				}
			}
		}
	}

	/**
	 * This implementation applies to any combination of dense/sparse if at least one
	 * input is ultrasparse (sparse and very few nnz). In that case, most importantly,
	 * we want to create a sparse output and only iterate over the few nnz as the major
	 * dimension. Low-level optimization have less importance in that case and having
	 * this generic implementation helps to reduce the implementations from (2+1)^2
	 * to 2^2+1.
	 * 
	 * @param m1
	 * @param m2
	 * @param ret
	 * @throws DMLRuntimeException
	 */
	private static void matrixMultUltraSparse(MatrixBlock m1, MatrixBlock m2, MatrixBlock ret, int rl, int ru) 
		throws DMLRuntimeException 
	{
		boolean leftUS = m1.isUltraSparse();
		final int m  = m1.rlen;
		final int cd = m1.clen;
		final int n  = m2.clen;
		
		if( leftUS ) //left is ultra-sparse (IKJ)
		{
			boolean rightSparse = m2.sparse;
			
			for( int i=rl; i<ru; i++ )
			{
				SparseRow arow = m1.sparseRows[ i ];
				if( arow != null && !arow.isEmpty() ) 
				{
					int alen = arow.size();
					int[] aixs = arow.getIndexContainer();
					double[] avals = arow.getValueContainer();	
					
					if( alen==1 && avals[0]==1 ) //ROW SELECTION (no aggregation)
					{
						int aix = aixs[0];
						if( rightSparse ) { //sparse right matrix (full row copy)
							if( m2.sparseRows!=null && m2.sparseRows[aix]!=null ) {
								ret.rlen=m;
								ret.allocateSparseRowsBlock(false); //allocation on demand
								ret.sparseRows[i] = new SparseRow(m2.sparseRows[aix]); 
								ret.nonZeros += ret.sparseRows[i].size();
							}
						}
						else { //dense right matrix (append all values)
							for( int j=0; j<n; j++ )
								ret.appendValue(i, j, m2.quickGetValue(aix, j));
						}
					}
					else //GENERAL CASE
					{
						for( int k=0; k<alen; k++ )
						{
							double aval = avals[k];
							int aix = aixs[k];
							for( int j=0; j<n; j++ )
							{
								double cval = ret.quickGetValue(i, j);
								double cvald = aval*m2.quickGetValue(aix, j);
								if( cvald != 0 )
									ret.quickSetValue(i, j, cval+cvald);
							}
						}
					}
				}
			}
		}
		else //right is ultra-sparse (KJI)
		{
			for(int k = 0; k < cd; k++ ) 
			{			
				SparseRow brow = m2.sparseRows[ k ];
				if( brow != null && !brow.isEmpty() ) 
				{
					int blen = brow.size();
					int[] bixs = brow.getIndexContainer();
					double[] bvals = brow.getValueContainer();								
					for( int j=0; j<blen; j++ )
					{
						double bval = bvals[j];
						int bix = bixs[j];
						for( int i=rl; i<ru; i++ )
						{
							double cvald = bval*m1.quickGetValue(i, k);
							if( cvald != 0 ){
								double cval = ret.quickGetValue(i, bix);
								ret.quickSetValue(i, bix, cval+cvald);
							}
						}
					}
				}
			}
		}
		//no need to recompute nonzeros because maintained internally
	}

	/**
	 * 
	 * @param mX
	 * @param mV
	 * @param mW
	 * @param ret
	 * @param ct
	 * @throws DMLRuntimeException
	 * @throws DMLUnsupportedOperationException
	 */
	private static void matrixMultChainDense(MatrixBlock mX, MatrixBlock mV, MatrixBlock mW, MatrixBlock ret, ChainType ct, int rl, int ru) 
	{
		double[] a = mX.denseBlock;
		double[] b = mV.denseBlock;
		double[] w = (mW!=null) ? mW.denseBlock : null;
		double[] c = ret.denseBlock;
		final int cd = mX.clen; //features in X
		boolean weights = (ct == ChainType.XtwXv);

		//temporary array for cache blocking
		//(blocksize chosen to fit b+v in L2 (256KB) for default 1k blocks)
		final int blocksize = 24; // constraint: factor of 4
		double[] tmp = new double[blocksize];
			
		//blockwise mmchain computation
		final int bn = ru - ru % blocksize; //rl blocksize aligned
		for( int bi=rl; bi < bn; bi+=blocksize ) 
		{
			//compute 1st matrix-vector for row block
			for( int j=0, aix=bi*cd; j < blocksize; j++, aix+=cd)
				tmp[j] = dotProduct(a, b, aix, 0, cd);
			
			//multiply weights (in-place), if required
			if( weights ) 
				vectMultiply(w, tmp, bi, 0, blocksize);	
			
			//compute 2nd matrix vector for row block and aggregate
			for (int j=0, aix=bi*cd; j < blocksize; j+=4, aix+=4*cd)
				vectMultiplyAdd4(tmp[j], tmp[j+1], tmp[j+2], tmp[j+3], a, c, aix, aix+cd, aix+2*cd, aix+3*cd, 0, cd);
		}
		
		//compute rest (not aligned to blocksize)
		for( int i=bn, aix=bn*cd; i < ru; i++, aix+=cd ) {
			double val = dotProduct(a, b, aix, 0, cd);
			val *= (weights) ? w[i] : 1; 
			vectMultiplyAdd(val, a, c, aix, 0, cd);				
		}
	}
	
{'fraction_non_alphanumeric': 0.1502420208743004, 'fraction_numerical': 0.017036000605052184, 'mean_word_length': 2.6153530658281494, 'pattern_counts': {'":': 0, '<': 97, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '834133', 'n_tokens_mistral': 23924, 'n_tokens_neox': 19150, 'n_words': 5055}