src/main/java/org/apache/sysds/runtime/instructions/gpu/context/GPUContextPool.java - systemds - Git at Google

 /*
  * 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.instructions.gpu.context;

 import static jcuda.driver.JCudaDriver.cuDeviceGetCount;
 import static jcuda.driver.JCudaDriver.cuInit;
 import static jcuda.runtime.JCuda.cudaGetDeviceProperties;

 import java.util.ArrayList;
 import java.util.LinkedList;
 import java.util.List;

 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.sysds.runtime.DMLRuntimeException;
 import org.apache.sysds.utils.GPUStatistics;

 import jcuda.driver.JCudaDriver;
 import jcuda.jcublas.JCublas2;
 import jcuda.jcudnn.JCudnn;
 import jcuda.jcusparse.JCusparse;
 import jcuda.runtime.JCuda;
 import jcuda.runtime.cudaDeviceProp;

 public class GPUContextPool {

 	protected static final Log LOG = LogFactory.getLog(GPUContextPool.class.getName());

 	/**
 	 * GPUs to use, can specify -1 to use all, comma separated list of GPU numbers, a specific GPU or a range
 	 */
 	public static String AVAILABLE_GPUS;


 	private static long INITIAL_GPU_MEMORY_BUDGET = -1;

 	/**
 	 * Whether cuda has been initialized
 	 */
 	static boolean initialized = false;

 	/**
 	 * The total number of cuda devices on this machine
 	 */
 	static int deviceCount = -1;

 	/**
 	 * Stores the cached deviceProperties
 	 */
 	static cudaDeviceProp[] deviceProperties;

 	/**
 	 * Set of free GPUContexts
 	 */
 	static List<GPUContext> pool = new LinkedList<>();

 	/**
 	 * Whether the pool of GPUs is reserved or not
 	 */
 	static boolean reserved = false;

 	/**
 	 * Static initialization of the number of devices
 	 * Also sets behaviour for J{Cuda, Cudnn, Cublas, Cusparse} in case of error
 	 * Initializes the CUDA driver
 	 * All these need be done once, and not per GPU
 	 */
 	public synchronized static void initializeGPU() {
 		if (initialized)
 			return;
 		initialized = true;
 		GPUContext.LOG.info("Initializing CUDA");
 		long start = System.nanoTime();
 		JCuda.setExceptionsEnabled(true);
 		JCudnn.setExceptionsEnabled(true);
 		JCublas2.setExceptionsEnabled(true);
 		JCusparse.setExceptionsEnabled(true);
 		JCudaDriver.setExceptionsEnabled(true);
 		cuInit(0); // Initialize the driver

 		int deviceCountArray[] = { 0 };
 		cuDeviceGetCount(deviceCountArray);        // Obtain the number of devices
 		deviceCount = deviceCountArray[0];
 		deviceProperties = new cudaDeviceProp[deviceCount];

 		try {
 			ArrayList<Integer> listOfGPUs = parseListString(AVAILABLE_GPUS, deviceCount);

 			// Initialize the list of devices & the pool of GPUContexts
 			for (int i : listOfGPUs) {
 				cudaDeviceProp properties = new cudaDeviceProp();
 				cudaGetDeviceProperties(properties, i);
 				deviceProperties[i] = properties;
 				GPUContext gCtx = new GPUContext(i);
 				pool.add(gCtx);
 			}

 		} catch (IllegalArgumentException e) {
 			LOG.warn("Invalid setting for setting systemds.gpu.availableGPUs, defaulting to use ALL GPUs");

 			// Initialize the list of devices & the pool of GPUContexts
 			for (int i = 0; i < deviceCount; i++) {
 				cudaDeviceProp properties = new cudaDeviceProp();
 				cudaGetDeviceProperties(properties, i);
 				deviceProperties[i] = properties;
 				GPUContext gCtx = new GPUContext(i);
 				pool.add(gCtx);
 			}
 		}


 		// Initialize the initial memory budget
 		// If there are heterogeneous GPUs on the machine (different memory sizes)
 		// initially available memory is set to the GPU with the lowest memory
 		// This is because at runtime, we wouldn't know which GPU a certain
 		// operation gets scheduled on
 		long minAvailableMemory = Long.MAX_VALUE;
 		for (GPUContext gCtx : pool) {
 			gCtx.initializeThread();
 			minAvailableMemory = Math.min(minAvailableMemory, gCtx.getAvailableMemory());
 		}
 		INITIAL_GPU_MEMORY_BUDGET = minAvailableMemory;


 		GPUContext.LOG.info("Total number of GPUs on the machine: " + deviceCount);
 		GPUContext.LOG.info("GPUs being used: " + AVAILABLE_GPUS);
 		GPUContext.LOG.info("Initial GPU memory: " + initialGPUMemBudget());

 		//int[] device = {-1};
 		//cudaGetDevice(device);
 		//cudaDeviceProp prop = getGPUProperties(device[0]);
 		//int maxBlocks = prop.maxGridSize[0];
 		//int maxThreadsPerBlock = prop.maxThreadsPerBlock;
 		//long sharedMemPerBlock = prop.sharedMemPerBlock;
 		//LOG.debug("Active CUDA device number : " + device[0]);
 		//LOG.debug("Max Blocks/Threads/SharedMem on active device: " + maxBlocks + "/" + maxThreadsPerBlock + "/" + sharedMemPerBlock);
 		GPUStatistics.cudaInitTime = System.nanoTime() - start;
 	}

 	/**
 	 * Parses a string into a list. The string can be of these forms:
 	 * 1. "-1" : all integers from range 0 to max - [0,1,2,3....max]
 	 * 2. "2,3,0" : comma separated list of integers - [0,2,3]
 	 * 3. "4" : a specific integer - [4]
 	 * 4. "0-4" : a range of integers - [0,1,2,3,4]
 	 * In ranges and comma separated lists, all values must be positive. Anything else is invalid.
 	 * @param str input string
 	 * @param max maximum range of integers
 	 * @return the list of integers in the parsed string
 	 */
 	public static ArrayList<Integer> parseListString(String str, int max) {
 		ArrayList<Integer> result = new ArrayList<>();
 		str = str.trim();
 		if (str.equalsIgnoreCase("-1")) {  // all
 			for (int i=0; i<max; i++){
 				result.add(i);
 			}
 		} else if (str.contains("-")){  // range
 			String[] numbersStr = str.split("-");
 			if (numbersStr.length != 2) {
 				throw new IllegalArgumentException("Invalid string to parse to a list of numbers : " + str);
 			}
 			String beginStr = numbersStr[0];
 			String endStr = numbersStr[1];
 			int begin = Integer.parseInt(beginStr);
 			int end = Integer.parseInt(endStr);

 			for (int i=begin; i<=end; i++){
 				result.add(i);
 			}
 		} else if (str.contains(",")) { // comma separated list
 			String[] numbers = str.split(",");
 			for (int i = 0; i < numbers.length; i++) {
 				int n = Integer.parseInt(numbers[i].trim());
 				result.add(n);
 			}
 		} else {  // single number
 			int number = Integer.parseInt(str);
 			result.add(number);
 		}
 		// Check if all numbers between 0 and max
 		for (int n : result){
 			if (n < 0 || n >= max) {
 				throw new IllegalArgumentException("Invalid string (" + str + ") parsed to a list of numbers (" + result + ") which exceeds the maximum range : ");
 			}
 		}
 		return result;
 	}

 	/**
 	 * Reserves and gets an initialized list of GPUContexts
 	 *
 	 * @return null if no GPUContexts in pool, otherwise a valid list of GPUContext
 	 */
 	public static synchronized List<GPUContext> reserveAllGPUContexts() {
 		if (reserved)
 			throw new DMLRuntimeException("Trying to re-reserve GPUs");
 		initializeGPU();
 		reserved = true;
 		LOG.trace("GPU : Reserved all GPUs");
 		return pool;
 	}

 	/**
 	 * Get the number of free GPUContexts
 	 *
 	 * @return number of free GPUContexts
 	 */
 	public static synchronized int getAvailableCount() {
 		return pool.size();
 	}

 	/**
 	 * Gets the device properties
 	 *
 	 * @param device the device number (on a machine with more than 1 GPU)
 	 * @return the device properties
 	 */
 	static cudaDeviceProp getGPUProperties(int device) {
 		// do once - initialization of GPU
 		if (!initialized)
 			initializeGPU();
 		return deviceProperties[device];
 	}

 	/**
 	 * Number of available devices on this machine
 	 *
 	 * @return number of available GPUs on this machine
 	 */
 	public static int getDeviceCount() {
 		if (!initialized)
 			initializeGPU();
 		return deviceCount;
 	}

 	/**
 	 * Unreserves all GPUContexts
 	 */
 	public static synchronized void freeAllGPUContexts() {
 		if (!reserved)
 			throw new DMLRuntimeException("Trying to free unreserved GPUs");
 		reserved = false;
 		LOG.trace("GPU : Unreserved all GPUs");

 	}

 	/**
 	 * Gets the initial GPU memory budget. This is the minimum of the
 	 * available memories across all the GPUs on the machine(s)
 	 * @return minimum available memory
 	 * @throws RuntimeException if error initializing the GPUs
 	 */
 	public static synchronized long initialGPUMemBudget() throws RuntimeException {
 		try {
 			if (!initialized)
 				initializeGPU();
 			return INITIAL_GPU_MEMORY_BUDGET;
 		} catch (DMLRuntimeException e){
 			throw new RuntimeException(e);
 		}
 	}
 }
	/*
	* 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.instructions.gpu.context;

	import static jcuda.driver.JCudaDriver.cuDeviceGetCount;
	import static jcuda.driver.JCudaDriver.cuInit;
	import static jcuda.runtime.JCuda.cudaGetDeviceProperties;

	import java.util.ArrayList;
	import java.util.LinkedList;
	import java.util.List;

	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.sysds.runtime.DMLRuntimeException;
	import org.apache.sysds.utils.GPUStatistics;

	import jcuda.driver.JCudaDriver;
	import jcuda.jcublas.JCublas2;
	import jcuda.jcudnn.JCudnn;
	import jcuda.jcusparse.JCusparse;
	import jcuda.runtime.JCuda;
	import jcuda.runtime.cudaDeviceProp;

	public class GPUContextPool {

	protected static final Log LOG = LogFactory.getLog(GPUContextPool.class.getName());

	/**
	* GPUs to use, can specify -1 to use all, comma separated list of GPU numbers, a specific GPU or a range
	*/
	public static String AVAILABLE_GPUS;


	private static long INITIAL_GPU_MEMORY_BUDGET = -1;

	/**
	* Whether cuda has been initialized
	*/
	static boolean initialized = false;

	/**
	* The total number of cuda devices on this machine
	*/
	static int deviceCount = -1;

	/**
	* Stores the cached deviceProperties
	*/
	static cudaDeviceProp[] deviceProperties;

	/**
	* Set of free GPUContexts
	*/
	static List<GPUContext> pool = new LinkedList<>();

	/**
	* Whether the pool of GPUs is reserved or not
	*/
	static boolean reserved = false;

	/**
	* Static initialization of the number of devices
	* Also sets behaviour for J{Cuda, Cudnn, Cublas, Cusparse} in case of error
	* Initializes the CUDA driver
	* All these need be done once, and not per GPU
	*/
	public synchronized static void initializeGPU() {
	if (initialized)
	return;
	initialized = true;
	GPUContext.LOG.info("Initializing CUDA");
	long start = System.nanoTime();
	JCuda.setExceptionsEnabled(true);
	JCudnn.setExceptionsEnabled(true);
	JCublas2.setExceptionsEnabled(true);
	JCusparse.setExceptionsEnabled(true);
	JCudaDriver.setExceptionsEnabled(true);
	cuInit(0); // Initialize the driver

	int deviceCountArray[] = { 0 };
	cuDeviceGetCount(deviceCountArray); // Obtain the number of devices
	deviceCount = deviceCountArray[0];
	deviceProperties = new cudaDeviceProp[deviceCount];

	try {
	ArrayList<Integer> listOfGPUs = parseListString(AVAILABLE_GPUS, deviceCount);

	// Initialize the list of devices & the pool of GPUContexts
	for (int i : listOfGPUs) {
	cudaDeviceProp properties = new cudaDeviceProp();
	cudaGetDeviceProperties(properties, i);
	deviceProperties[i] = properties;
	GPUContext gCtx = new GPUContext(i);
	pool.add(gCtx);
	}

	} catch (IllegalArgumentException e) {
	LOG.warn("Invalid setting for setting systemds.gpu.availableGPUs, defaulting to use ALL GPUs");

	// Initialize the list of devices & the pool of GPUContexts
	for (int i = 0; i < deviceCount; i++) {
	cudaDeviceProp properties = new cudaDeviceProp();
	cudaGetDeviceProperties(properties, i);
	deviceProperties[i] = properties;
	GPUContext gCtx = new GPUContext(i);
	pool.add(gCtx);
	}
	}


	// Initialize the initial memory budget
	// If there are heterogeneous GPUs on the machine (different memory sizes)
	// initially available memory is set to the GPU with the lowest memory
	// This is because at runtime, we wouldn't know which GPU a certain
	// operation gets scheduled on
	long minAvailableMemory = Long.MAX_VALUE;
	for (GPUContext gCtx : pool) {
	gCtx.initializeThread();
	minAvailableMemory = Math.min(minAvailableMemory, gCtx.getAvailableMemory());
	}
	INITIAL_GPU_MEMORY_BUDGET = minAvailableMemory;


	GPUContext.LOG.info("Total number of GPUs on the machine: " + deviceCount);
	GPUContext.LOG.info("GPUs being used: " + AVAILABLE_GPUS);
	GPUContext.LOG.info("Initial GPU memory: " + initialGPUMemBudget());

	//int[] device = {-1};
	//cudaGetDevice(device);
	//cudaDeviceProp prop = getGPUProperties(device[0]);
	//int maxBlocks = prop.maxGridSize[0];
	//int maxThreadsPerBlock = prop.maxThreadsPerBlock;
	//long sharedMemPerBlock = prop.sharedMemPerBlock;
	//LOG.debug("Active CUDA device number : " + device[0]);
	//LOG.debug("Max Blocks/Threads/SharedMem on active device: " + maxBlocks + "/" + maxThreadsPerBlock + "/" + sharedMemPerBlock);
	GPUStatistics.cudaInitTime = System.nanoTime() - start;
	}

	/**
	* Parses a string into a list. The string can be of these forms:
	* 1. "-1" : all integers from range 0 to max - [0,1,2,3....max]
	* 2. "2,3,0" : comma separated list of integers - [0,2,3]
	* 3. "4" : a specific integer - [4]
	* 4. "0-4" : a range of integers - [0,1,2,3,4]
	* In ranges and comma separated lists, all values must be positive. Anything else is invalid.
	* @param str input string
	* @param max maximum range of integers
	* @return the list of integers in the parsed string
	*/
	public static ArrayList<Integer> parseListString(String str, int max) {
	ArrayList<Integer> result = new ArrayList<>();
	str = str.trim();
	if (str.equalsIgnoreCase("-1")) { // all
	for (int i=0; i<max; i++){
	result.add(i);
	}
	} else if (str.contains("-")){ // range
	String[] numbersStr = str.split("-");
	if (numbersStr.length != 2) {
	throw new IllegalArgumentException("Invalid string to parse to a list of numbers : " + str);
	}
	String beginStr = numbersStr[0];
	String endStr = numbersStr[1];
	int begin = Integer.parseInt(beginStr);
	int end = Integer.parseInt(endStr);

	for (int i=begin; i<=end; i++){
	result.add(i);
	}
	} else if (str.contains(",")) { // comma separated list
	String[] numbers = str.split(",");
	for (int i = 0; i < numbers.length; i++) {
	int n = Integer.parseInt(numbers[i].trim());
	result.add(n);
	}
	} else { // single number
	int number = Integer.parseInt(str);
	result.add(number);
	}
	// Check if all numbers between 0 and max
	for (int n : result){
	if (n < 0 \|\| n >= max) {
	throw new IllegalArgumentException("Invalid string (" + str + ") parsed to a list of numbers (" + result + ") which exceeds the maximum range : ");
	}
	}
	return result;
	}

	/**
	* Reserves and gets an initialized list of GPUContexts
	*
	* @return null if no GPUContexts in pool, otherwise a valid list of GPUContext
	*/
	public static synchronized List<GPUContext> reserveAllGPUContexts() {
	if (reserved)
	throw new DMLRuntimeException("Trying to re-reserve GPUs");
	initializeGPU();
	reserved = true;
	LOG.trace("GPU : Reserved all GPUs");
	return pool;
	}

	/**
	* Get the number of free GPUContexts
	*
	* @return number of free GPUContexts
	*/
	public static synchronized int getAvailableCount() {
	return pool.size();
	}

	/**
	* Gets the device properties
	*
	* @param device the device number (on a machine with more than 1 GPU)
	* @return the device properties
	*/
	static cudaDeviceProp getGPUProperties(int device) {
	// do once - initialization of GPU
	if (!initialized)
	initializeGPU();
	return deviceProperties[device];
	}

	/**
	* Number of available devices on this machine
	*
	* @return number of available GPUs on this machine
	*/
	public static int getDeviceCount() {
	if (!initialized)
	initializeGPU();
	return deviceCount;
	}

	/**
	* Unreserves all GPUContexts
	*/
	public static synchronized void freeAllGPUContexts() {
	if (!reserved)
	throw new DMLRuntimeException("Trying to free unreserved GPUs");
	reserved = false;
	LOG.trace("GPU : Unreserved all GPUs");

	}

	/**
	* Gets the initial GPU memory budget. This is the minimum of the
	* available memories across all the GPUs on the machine(s)
	* @return minimum available memory
	* @throws RuntimeException if error initializing the GPUs
	*/
	public static synchronized long initialGPUMemBudget() throws RuntimeException {
	try {
	if (!initialized)
	initializeGPU();
	return INITIAL_GPU_MEMORY_BUDGET;
	} catch (DMLRuntimeException e){
	throw new RuntimeException(e);
	}
	}
	}