[SYSTEMDS-3556] Counter based random number generator

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

@@ -21,11 +21,14 @@
 package org.apache.sysds.runtime.matrix.data;
 
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
 import java.util.Random;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import java.util.stream.LongStream;
+import java.util.stream.Stream;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -37,12 +40,17 @@
 import org.apache.sysds.runtime.data.DenseBlock;
 import org.apache.sysds.runtime.data.SparseBlock;
 import org.apache.sysds.runtime.util.CommonThreadPool;
+import org.apache.sysds.runtime.util.CounterBasedPRNGenerator;
+import org.apache.sysds.runtime.util.IPRNGenerator;
 import org.apache.sysds.runtime.util.NormalPRNGenerator;
 import org.apache.sysds.runtime.util.PRNGenerator;
+import org.apache.sysds.runtime.util.PhiloxNormalCBPRNGenerator;
+import org.apache.sysds.runtime.util.PhiloxUniformCBPRNGenerator;
 import org.apache.sysds.runtime.util.PoissonPRNGenerator;
 import org.apache.sysds.runtime.util.UniformPRNGenerator;
 import org.apache.sysds.runtime.util.UtilFunctions;
 
+
 public class LibMatrixDatagen 
 {
 	private static final Log LOG = LogFactory.getLog(LibMatrixDatagen.class.getName());
@@ -140,9 +148,11 @@ public static RandomMatrixGenerator createRandomMatrixGenerator(String pdfStr, i
 		RandomMatrixGenerator rgen = null;
 		switch (pdf) {
 			case UNIFORM:
+			case CB_UNIFORM:
 				rgen = new RandomMatrixGenerator(pdf, r, c, blen, sp, min, max);
 				break;
 			case NORMAL:
+			case CB_NORMAL:
 				rgen = new RandomMatrixGenerator(pdf, r, c, blen, sp);
 				break;
 			case POISSON:
@@ -465,13 +475,17 @@ private static void genRandomNumbers(boolean invokedFromCP, int rl, int ru, int
 		int ncb = (int) Math.ceil((double)cols/blen);
 		int counter = 0;
 
+		long[] ctr = {0, 0, 0, 0}; // Counter based RNG counter
+
 		// Setup Pseudo Random Number Generator for cell values based on 'pdf'.
-		PRNGenerator valuePRNG = rgen._valuePRNG;
+		IPRNGenerator valuePRNG = rgen._valuePRNG;
 		if (valuePRNG == null) {
 			switch (rgen._pdf) {
 				case UNIFORM: valuePRNG = new UniformPRNGenerator(); break;
 				case NORMAL:  valuePRNG = new NormalPRNGenerator(); break;
 				case POISSON: valuePRNG = new PoissonPRNGenerator(); break;
+				case CB_UNIFORM: valuePRNG = new PhiloxUniformCBPRNGenerator(); break;
+				case CB_NORMAL: valuePRNG = new PhiloxNormalCBPRNGenerator(); break;
 				default:
 					throw new DMLRuntimeException("Unsupported distribution function for Rand: " + rgen._pdf);
 			}
@@ -505,7 +519,13 @@ private static void genRandomNumbers(boolean invokedFromCP, int rl, int ru, int
 				// Also note that we cannot use the same seed here, because for ultra-sparse generation
 				// the number of calls to the valuePRNG and nnzPRNG are the same, thus creating correlated
 				// outcomes (bias toward the end of the value range)
-				nnzPRNG.setSeed((long)(valuePRNG.nextDouble()*Long.MAX_VALUE));
+
+				if (valuePRNG instanceof CounterBasedPRNGenerator) {
+					nnzPRNG.setSeed((long)(((CounterBasedPRNGenerator) valuePRNG).getDoubles(ctr, 1)[0]*Long.MAX_VALUE ));
+				} else {
+					nnzPRNG.setSeed((long)(((PRNGenerator) valuePRNG).nextDouble()*Long.MAX_VALUE));
+				}
+
 				boolean localSparse = sparsity < 1 && MatrixBlock.evalSparseFormatInMemory(
 					blockrows, blockcols, (long)(sparsity*blockrows*blockcols));
 				if ( localSparse) {
@@ -515,17 +535,22 @@ private static void genRandomNumbers(boolean invokedFromCP, int rl, int ru, int
 						out.sparse = true; //otherwise ignored
 						c = out.sparseBlock;
 					}
-					genSparse(c, clen, blockrows, blockcols, rowoffset, coloffset,
-						sparsity, estnnzRow, min, range, valuePRNG, nnzPRNG);
+					if (valuePRNG instanceof PRNGenerator) {
+						genSparse(c, clen, blockrows, blockcols, rowoffset, coloffset,
+								sparsity, estnnzRow, min, range, (PRNGenerator) valuePRNG, nnzPRNG);
+					}
 				}
 				else {
 					if (sparsity == 1.0) {
 						genFullyDense(out.getDenseBlock(), blockrows, blockcols,
-							rowoffset, coloffset, min, range, valuePRNG);
+							rowoffset, coloffset, min, range, valuePRNG, ctr);
 					}
 					else {
-						genDense(out, clen, blockrows, blockcols, rowoffset, coloffset,
-							sparsity, estnnzRow, min, range, valuePRNG, nnzPRNG);
+						if (valuePRNG instanceof PRNGenerator) {
+							genDense(out, clen, blockrows, blockcols, rowoffset, coloffset,
+									sparsity, estnnzRow, min, range, (PRNGenerator) valuePRNG, nnzPRNG);
+						}
+
 					}
 				} // sparse or dense 
 			} // cbj
@@ -590,13 +615,22 @@ private static void genDense(MatrixBlock out, int clen, int blockrows, int block
 	}
 
 	private static void genFullyDense(DenseBlock c, int blockrows, int blockcols, int rowoffset, int coloffset, 
-		double min, double range, PRNGenerator valuePRNG)
+		double min, double range, IPRNGenerator valuePRNG, long[] ctr)
 	{
-		for(int i = rowoffset; i < rowoffset+blockrows; i++) {
+		Iterator<Double> rngStream;
+		if (valuePRNG instanceof PRNGenerator) {
+			rngStream = Stream.generate(() -> min + (range * ((PRNGenerator) valuePRNG).nextDouble())).iterator();
+		} else if (valuePRNG instanceof CounterBasedPRNGenerator) {
+			rngStream = Arrays.stream(((CounterBasedPRNGenerator)valuePRNG).getDoubles(ctr, blockrows * blockcols)).map(i -> min + (range * i)).iterator();
+		} else {
+			throw new DMLRuntimeException("Unsupported PRNGenerator for fully dense generation");
+		}
+		for (int i = rowoffset; i < rowoffset + blockrows; i++) {
 			double[] cvals = c.values(i);
 			int cix = c.pos(i, coloffset);
-			for(int j = 0; j < blockcols; j++)
-				cvals[cix+j] = min + (range * valuePRNG.nextDouble());
+			for (int j = 0; j < blockcols; j++) {
+				cvals[cix + j] = rngStream.next();
+			}
 		}
 	}
 

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

@@ -20,8 +20,10 @@
 package org.apache.sysds.runtime.matrix.data;
 
 import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.util.IPRNGenerator;
 import org.apache.sysds.runtime.util.NormalPRNGenerator;
-import org.apache.sysds.runtime.util.PRNGenerator;
+import org.apache.sysds.runtime.util.PhiloxNormalCBPRNGenerator;
+import org.apache.sysds.runtime.util.PhiloxUniformCBPRNGenerator;
 import org.apache.sysds.runtime.util.PoissonPRNGenerator;
 import org.apache.sysds.runtime.util.UniformPRNGenerator;
 
@@ -31,14 +33,14 @@ public class RandomMatrixGenerator {
 	 * Types of Probability density functions
 	 */
 	public enum PDF {
-		NORMAL, UNIFORM, POISSON
+		NORMAL, UNIFORM, POISSON, CB_UNIFORM, CB_NORMAL
 	}
 
 	PDF _pdf;
 	int _rows, _cols, _blocksize;
 	double _sparsity, _mean; 
 	double _min, _max; 
-	PRNGenerator _valuePRNG;
+	IPRNGenerator _valuePRNG;
 
 	public RandomMatrixGenerator() {
 		_pdf = PDF.UNIFORM;
@@ -166,6 +168,12 @@ protected void setupValuePRNG() {
 				throw new DMLRuntimeException("Invalid parameter (" + _mean + ") for Poisson distribution.");
 			_valuePRNG = new PoissonPRNGenerator(_mean);
 			break;
+		case CB_UNIFORM:
+			_valuePRNG = new PhiloxUniformCBPRNGenerator();
+			break;
+		case CB_NORMAL:
+			_valuePRNG = new PhiloxNormalCBPRNGenerator();
+			break;
 		default:
 			throw new DMLRuntimeException("Unsupported probability density function");
 		}

src/main/java/org/apache/sysds/runtime/util/CounterBasedPRNGenerator.java

@@ -0,0 +1,28 @@
+/*
+ * 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.util;
+
+public abstract class CounterBasedPRNGenerator implements IPRNGenerator {
+
+    public abstract void setSeed(long sd);
+
+    public abstract double[] getDoubles(long[] ctr, int size);
+}

src/main/java/org/apache/sysds/runtime/util/IPRNGenerator.java

@@ -0,0 +1,25 @@
+/*
+ * 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.util;
+
+public interface IPRNGenerator {
+    public void setSeed(long seed);
+}

src/main/java/org/apache/sysds/runtime/util/PRNGenerator.java

@@ -20,7 +20,7 @@
 
 package org.apache.sysds.runtime.util;
 
-public abstract class PRNGenerator {
+public abstract class PRNGenerator implements IPRNGenerator {
 
 	public abstract void setSeed(long sd);
 

src/main/java/org/apache/sysds/runtime/util/PhiloxNormalCBPRNGenerator.java

@@ -0,0 +1,65 @@
+/*
+ * 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.util;
+
+public class PhiloxNormalCBPRNGenerator extends CounterBasedPRNGenerator {
+    private long[] seed;
+    private PhiloxUniformCBPRNGenerator uniformGen;
+
+    public void setSeed(long sd) {
+        this.seed = new long[2];
+        this.seed[0] = sd;
+        this.seed[1] = sd;
+        uniformGen = new PhiloxUniformCBPRNGenerator();
+        uniformGen.setSeed(this.seed[0]);
+    }
+
+    /**
+     * Generate a sequence of random doubles using the Philox4x64 counter-based PRNG.
+     *
+     * @param ctr  The start counter to use for the PRNG
+     * @param size The number of doubles to generate
+     * @return An array of random doubles distributed normally with mean 0 and variance 1
+     */
+    public double[] getDoubles(long[] ctr, int size) {
+        // Ensure the key is correct size
+        if (this.seed.length != 2) {
+            throw new IllegalArgumentException("Key must be 128 bits");
+        }
+        // Ensure the counter is correct size
+        if (ctr.length != 4) {
+            throw new IllegalArgumentException("Counter must be 256 bits");
+        }
+
+        double[] uniform = uniformGen.getDoubles(ctr, size + size % 2);
+        double[] normal = new double[size];
+        for (int i = 0; i < size; i+=2) {
+            double v1 = Math.sqrt(-2*Math.log(uniform[i]));
+            double v2 = 2*Math.PI*uniform[i + 1];
+            normal[i] = v1 * Math.cos(v2);
+            if (i + 1 < size) {
+                normal[i + 1] = v1 * Math.sin(v2);
+            }
+        }
+
+        return normal;
+    }
+}