Search device for all platform in tensor benchmark

Author: Atry

OpenCL/src/main/scala/com/thoughtworks/compute/OpenCL.scala

@@ -50,6 +50,29 @@ import scala.language.higherKinds
   */
 object OpenCL {
 
+  final case class PlatformId[Owner <: Singleton with OpenCL](handle: Long) extends AnyVal {
+
+    final def deviceIdsByType(deviceType: Int): Seq[DeviceId[Owner]] = {
+      val Array(numberOfDevices) = {
+        val a = Array(0)
+        checkErrorCode(clGetDeviceIDs(handle, deviceType, null, a))
+        a
+      }
+      val stack = stackPush()
+      try {
+        val deviceIdBuffer = stack.mallocPointer(numberOfDevices)
+        checkErrorCode(clGetDeviceIDs(handle, deviceType, deviceIdBuffer, null: IntBuffer))
+        for (i <- 0 until numberOfDevices) yield {
+          val deviceId = deviceIdBuffer.get(i)
+          new DeviceId[Owner](deviceId)
+        }
+      } finally {
+        stack.close()
+      }
+    }
+
+  }
+
   /** Returns a [[String]] for the C string `address`.
     *
     * @note We don't know the exact charset of the C string. Use [[memASCII]] because lwjgl treats them as ASCII.
@@ -265,25 +288,18 @@ object OpenCL {
     }
   }
 
-  trait UseFirstPlatform {
+  trait UseFirstPlatform extends OpenCL {
     @transient
-    protected lazy val platformId: Long = {
-      val stack = stackPush()
-      try {
-        val platformIdBuffer = stack.mallocPointer(1)
-        checkErrorCode(clGetPlatformIDs(platformIdBuffer, null: IntBuffer))
-        platformIdBuffer.get(0)
-      } finally {
-        stack.close()
-      }
+    protected lazy val platformId: PlatformId = {
+      platformIds.head
     }
   }
 
   trait UseAllDevices extends OpenCL {
 
     @transient
     protected lazy val deviceIds: Seq[DeviceId] = {
-      deviceIdsByType(CL_DEVICE_TYPE_ALL)
+      platformId.deviceIdsByType(CL_DEVICE_TYPE_ALL)
     }
 
   }
@@ -292,7 +308,7 @@ object OpenCL {
 
     @transient
     protected lazy val deviceIds: Seq[DeviceId] = {
-      val allDeviceIds = deviceIdsByType(CL_DEVICE_TYPE_ALL)
+      val allDeviceIds = platformId.deviceIdsByType(CL_DEVICE_TYPE_ALL)
       Seq(allDeviceIds.head)
     }
 
@@ -302,23 +318,23 @@ object OpenCL {
 
     @transient
     protected lazy val deviceIds: Seq[DeviceId] = {
-      deviceIdsByType(CL_DEVICE_TYPE_GPU)
+      platformId.deviceIdsByType(CL_DEVICE_TYPE_GPU)
     }
   }
 
   trait UseFirstGpuDevice extends OpenCL {
 
     @transient
     protected lazy val deviceIds: Seq[DeviceId] = {
-      val allDeviceIds = deviceIdsByType(CL_DEVICE_TYPE_GPU)
+      val allDeviceIds = platformId.deviceIdsByType(CL_DEVICE_TYPE_GPU)
       Seq(allDeviceIds.head)
     }
   }
   trait UseFirstCpuDevice extends OpenCL {
 
     @transient
     protected lazy val deviceIds: Seq[DeviceId] = {
-      val allDeviceIds = deviceIdsByType(CL_DEVICE_TYPE_CPU)
+      val allDeviceIds = platformId.deviceIdsByType(CL_DEVICE_TYPE_CPU)
       Seq(allDeviceIds.head)
     }
   }
@@ -327,7 +343,7 @@ object OpenCL {
 
     @transient
     protected lazy val deviceIds: Seq[DeviceId] = {
-      deviceIdsByType(CL_DEVICE_TYPE_CPU)
+      platformId.deviceIdsByType(CL_DEVICE_TYPE_CPU)
     }
   }
 
@@ -1010,20 +1026,18 @@ trait OpenCL extends MonadicCloseable[UnitContinuation] with ImplicitsSingleton
   type Event = OpenCL.Event[this.type]
   type CommandQueue = OpenCL.CommandQueue[this.type]
   type DeviceId = OpenCL.DeviceId[this.type]
+  type PlatformId = OpenCL.PlatformId[this.type]
 
-  protected final def deviceIdsByType(deviceType: Int): Seq[DeviceId] = {
-    val Array(numberOfDevices) = {
-      val a = Array(0)
-      checkErrorCode(clGetDeviceIDs(platformId, deviceType, null, a))
-      a
-    }
+  def platformIds: Seq[PlatformId] = {
     val stack = stackPush()
     try {
-      val deviceIdBuffer = stack.mallocPointer(numberOfDevices)
-      checkErrorCode(clGetDeviceIDs(platformId, deviceType, deviceIdBuffer, null: IntBuffer))
-      for (i <- 0 until deviceIdBuffer.capacity()) yield {
-        val deviceId = deviceIdBuffer.get(i)
-        new DeviceId(deviceId)
+      val numberOfPlatformsBuffer = stack.mallocInt(1)
+      checkErrorCode(clGetPlatformIDs(null, numberOfPlatformsBuffer))
+      val numberOfPlatforms = numberOfPlatformsBuffer.get(0)
+      val platformIdBuffer = stack.mallocPointer(numberOfPlatforms)
+      checkErrorCode(clGetPlatformIDs(platformIdBuffer, null: IntBuffer))
+      (0 until numberOfPlatforms).map { i =>
+        new PlatformId(platformIdBuffer.get(i))
       }
     } finally {
       stack.close()
@@ -1174,20 +1188,20 @@ trait OpenCL extends MonadicCloseable[UnitContinuation] with ImplicitsSingleton
 
   import OpenCL._
 
-  protected val platformId: Long
+  protected val platformId: PlatformId
   protected val deviceIds: Seq[DeviceId]
 
   @transient
   protected lazy val platformCapabilities: CLCapabilities = {
-    CL.createPlatformCapabilities(platformId)
+    CL.createPlatformCapabilities(platformId.handle)
   }
 
   protected def createCommandQueue(deviceId: DeviceId, properties: Map[Int, Long]): CommandQueue = new CommandQueue(
     if (deviceCapabilities(deviceId).OpenCL20) {
       val cl20Properties = (properties.view.flatMap { case (key, value) => Seq(key, value) } ++ Seq(0L)).toArray
       val a = Array(0)
       val commandQueue =
-        clCreateCommandQueueWithProperties(platformId, deviceId.handle, cl20Properties, a)
+        clCreateCommandQueueWithProperties(platformId.handle, deviceId.handle, cl20Properties, a)
       checkErrorCode(a(0))
       commandQueue
     } else {
@@ -1211,7 +1225,7 @@ trait OpenCL extends MonadicCloseable[UnitContinuation] with ImplicitsSingleton
     val stack = stackPush()
     try {
       val errorCodeBuffer = stack.ints(CL_SUCCESS)
-      val contextProperties = stack.pointers(CL_CONTEXT_PLATFORM, platformId, 0)
+      val contextProperties = stack.pointers(CL_CONTEXT_PLATFORM, platformId.handle, 0)
       val deviceIdBuffer = stack.pointers(deviceIds.view.map(_.handle): _*)
       val context =
         clCreateContext(contextProperties,

benchmarks/src/jmh/scala/com/thoughtworks/compute/benchmarks.scala

@@ -1,5 +1,6 @@
 package com.thoughtworks.compute
 
+import com.thoughtworks.compute.OpenCL.Exceptions.DeviceNotFound
 import com.thoughtworks.compute.benchmarks.RandomNormalState
 import com.thoughtworks.feature.Factory
 import com.thoughtworks.future._
@@ -25,23 +26,49 @@ object benchmarks {
 
   trait TensorState {
     @Param(Array("CPU", "GPU"))
-    protected var deviceType: String = _
+    protected var tensorDeviceType: String = _
 
     trait BenchmarkTensors
         extends StrictLogging
         with Tensors.UnsafeMathOptimizations
         with Tensors.SuppressWarnings
         with OpenCL.LogContextNotification
         with OpenCL.GlobalExecutionContext
-        with OpenCL.UseFirstPlatform
         with OpenCL.CommandQueuePool
         with OpenCL.DontReleaseEventTooEarly
         with Tensors.WangHashingRandomNumberGenerator {
-
       @transient
-      protected lazy val deviceIds: Seq[DeviceId] = {
-        deviceIdsByType(classOf[CL10].getField(s"CL_DEVICE_TYPE_$deviceType").get(null).asInstanceOf[Int])
+      protected lazy val (platformId: PlatformId, deviceIds: Seq[DeviceId]) = {
+        val deviceType = classOf[CL10].getField(s"CL_DEVICE_TYPE_$tensorDeviceType").get(null).asInstanceOf[Int]
+
+        object MatchDeviceType {
+          def unapply(platformId: PlatformId): Option[Seq[DeviceId]] = {
+            (try {
+              platformId.deviceIdsByType(deviceType)
+            } catch {
+              case e: DeviceNotFound =>
+                return None
+            }) match {
+              case devices if devices.nonEmpty =>
+                Some(devices)
+              case _ =>
+                None
+            }
+
+          }
+        }
+
+        platformIds.collectFirst {
+          case platformId @ MatchDeviceType(deviceIds) =>
+            (platformId, deviceIds)
+        } match {
+          case None =>
+            throw new DeviceNotFound(s"$tensorDeviceType device is not found")
+          case Some(pair) =>
+            pair
+        }
       }
+
     }
   }