[SYSTEMDS-2951] Multi-GPU Support for End-to-End ML Pipelines

src/main/java/org/apache/sysds/runtime/controlprogram/ParForProgramBlock.java

@@ -755,6 +755,8 @@ private void executeLocalParFor( ExecutionContext ec, IntObject from, IntObject
 
 		//restrict recompilation to thread local memory
 		setMemoryBudget();
+
+		_numThreads = Math.min(_numThreads, ec.getNumGPUContexts());
 
 		final LocalTaskQueue<Task> queue = new LocalTaskQueue<>();
 		final Thread[] threads         = new Thread[_numThreads];

src/main/java/org/apache/sysds/runtime/controlprogram/caching/CacheableData.java

@@ -22,8 +22,8 @@
 import java.io.File;
 import java.io.IOException;
 import java.lang.ref.SoftReference;
-import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicLong;
 
 import org.apache.commons.lang3.mutable.MutableBoolean;
@@ -214,7 +214,7 @@ public enum CacheStatus {
 	//for lazily evaluated RDDs, and (2) as abstraction for environments that do not necessarily have spark libraries available
 	private RDDObject _rddHandle = null; //RDD handle
 	private BroadcastObject<T> _bcHandle = null; //Broadcast handle
-	protected HashMap<GPUContext, GPUObject> _gpuObjects = null; //Per GPUContext object allocated on GPU
+	protected ConcurrentHashMap<GPUContext, GPUObject> _gpuObjects = null; //Per GPUContext object allocated on GPU
 
 	private LineageItem _lineage = null;
 
@@ -229,7 +229,7 @@ protected CacheableData(DataType dt, ValueType vt) {
 		_uniqueID = _seq.getNextID();
 		_cacheStatus = CacheStatus.EMPTY;
 		_numReadThreads = 0;
-		_gpuObjects = DMLScript.USE_ACCELERATOR ? new HashMap<>() : null;
+		_gpuObjects = DMLScript.USE_ACCELERATOR ? new ConcurrentHashMap<>() : null;
 	}
 
 	/**
@@ -472,7 +472,7 @@ public synchronized GPUObject getGPUObject(GPUContext gCtx) {
 
 	public synchronized void setGPUObject(GPUContext gCtx, GPUObject gObj) {
 		if( _gpuObjects == null )
-			_gpuObjects = new HashMap<>();
+			_gpuObjects = new ConcurrentHashMap<>();
 		GPUObject old = _gpuObjects.put(gCtx, gObj);
 		if (old != null)
 				throw new DMLRuntimeException("GPU : Inconsistent internal state - this CacheableData already has a GPUObject assigned to the current GPUContext (" + gCtx + ")");