diff --git a/buildSrc/src/main/kotlin/Dependencies.kt b/buildSrc/src/main/kotlin/Dependencies.kt index 064f023960..29b8696e5d 100644 --- a/buildSrc/src/main/kotlin/Dependencies.kt +++ b/buildSrc/src/main/kotlin/Dependencies.kt @@ -241,6 +241,12 @@ object Libs { version = Versions.onnxruntime ) + val onnxruntime_gpu = dep( + group = "com.microsoft.onnxruntime", + name = "onnxruntime_gpu", + version = Versions.onnxruntime + ) + val ktor_server_core = dep( group = "io.ktor", name = "ktor-server-core", diff --git a/settings.gradle.kts b/settings.gradle.kts index 10dbc5dabb..7f737460c6 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -21,6 +21,7 @@ findProject(":usvm-python:usvm-python-runner")?.name = "usvm-python-runner" include("usvm-python:usvm-python-commons") findProject(":usvm-python:usvm-python-commons")?.name = "usvm-python-commons" include("usvm-ml-gameserver") +include("usvm-runner") pluginManagement { resolutionStrategy { diff --git a/usvm-core/build.gradle.kts b/usvm-core/build.gradle.kts index e142271ec3..f43db19eb1 100644 --- a/usvm-core/build.gradle.kts +++ b/usvm-core/build.gradle.kts @@ -13,6 +13,7 @@ dependencies { testImplementation(Libs.ksmt_yices) implementation(Libs.onnxruntime) + implementation(Libs.onnxruntime_gpu) } publishing { diff --git a/usvm-core/src/main/kotlin/org/usvm/utils/OnnxModelImpl.kt b/usvm-core/src/main/kotlin/org/usvm/utils/OnnxModelImpl.kt index dd589b13b6..20699124df 100644 --- a/usvm-core/src/main/kotlin/org/usvm/utils/OnnxModelImpl.kt +++ b/usvm-core/src/main/kotlin/org/usvm/utils/OnnxModelImpl.kt @@ -4,18 +4,31 @@ import ai.onnxruntime.OnnxTensor import ai.onnxruntime.OrtEnvironment import ai.onnxruntime.OrtSession import org.usvm.StateId +import org.usvm.logger import org.usvm.statistics.BasicBlock import org.usvm.statistics.BlockGraph import org.usvm.util.OnnxModel import java.nio.FloatBuffer import java.nio.LongBuffer -import java.nio.file.Path + +enum class Mode { + CPU, + GPU +} class OnnxModelImpl( - pathToONNX: Path + pathToONNX: String, + mode: Mode ): OnnxModel> { private val env: OrtEnvironment = OrtEnvironment.getEnvironment() - private val session: OrtSession = env.createSession(pathToONNX.toString()) + private val sessionOptions: OrtSession.SessionOptions = OrtSession.SessionOptions().apply { + if (mode == Mode.GPU) { + addCUDA() + } else { + logger.info("Using CPU execution provider.") + } + } + private val session: OrtSession = env.createSession(pathToONNX, sessionOptions) override fun predictState(game: Game): UInt { val stateIds = mutableMapOf() diff --git a/usvm-ml-gameserver/build.gradle.kts b/usvm-ml-gameserver/build.gradle.kts index 23f0fcaf10..a294230383 100644 --- a/usvm-ml-gameserver/build.gradle.kts +++ b/usvm-ml-gameserver/build.gradle.kts @@ -32,6 +32,7 @@ tasks.test { dependencies { implementation(project(":usvm-jvm")) implementation(project(":usvm-core")) + implementation(project(":usvm-runner")) implementation(Libs.jacodb_api_jvm) implementation(Libs.jacodb_core) @@ -44,6 +45,5 @@ dependencies { implementation(Libs.kotlinx_cli) implementation(Libs.slf4j_simple) - implementation("io.github.rchowell:dotlin:1.0.2") testImplementation(kotlin("test")) } diff --git a/usvm-ml-gameserver/src/main/kotlin/org.usvm/JavaMethodRunner.kt b/usvm-ml-gameserver/src/main/kotlin/org.usvm/JavaMethodRunner.kt deleted file mode 100644 index ef6e0437ab..0000000000 --- a/usvm-ml-gameserver/src/main/kotlin/org.usvm/JavaMethodRunner.kt +++ /dev/null @@ -1,57 +0,0 @@ -package org.usvm - -import org.usvm.gameserver.GameMap -import org.usvm.gameserver.Searcher.BFS -import org.usvm.gameserver.Searcher.DFS -import org.usvm.machine.JcMachine -import org.usvm.machine.state.JcState -import org.usvm.util.JacoDBContainer -import org.usvm.util.Predictor -import org.usvm.util.getJcMethodByName -import org.usvm.util.loadClasspathFromJar -import org.usvm.util.mapsKey -import java.io.File -import kotlin.time.Duration - -class JavaMethodRunner(gameMap: GameMap, oracle: Predictor<*>? = null) { - val jacodbCpKey: String - get() = mapsKey - - val classpath: List = - loadClasspathFromJar(gameMap.assemblyFullName) - - val cp by lazy { - JacoDBContainer(jacodbCpKey, classpath).cp - } - - val stepLimit = gameMap.stepsToPlay + gameMap.stepsToStart - - val defaultSearcher = when (gameMap.defaultSearcher) { - BFS -> PathSelectionStrategy.BFS - DFS -> PathSelectionStrategy.DFS - } - - val options: UMachineOptions = UMachineOptions( - pathSelectionStrategies = listOf(defaultSearcher, PathSelectionStrategy.AI), - pathSelectorCombinationStrategy = PathSelectorCombinationStrategy.SEQUENTIAL, - coverageZone = CoverageZone.METHOD, - exceptionsPropagation = true, - solverTimeout = Duration.INFINITE, - timeout = Duration.INFINITE, - typeOperationsTimeout = Duration.INFINITE, - useSolverForForks = false, - - stepLimit = stepLimit.toULong(), - stepsToStart = gameMap.stepsToStart, - oracle = oracle - ) - - fun cover(className: String, methodName: String): Pair, Float> { - val jcMethod = cp.getJcMethodByName(className, methodName) - - val (states, percentageCoverage) = JcMachine(cp, options, interpreterObserver = null).use { machine -> - machine.analyze(jcMethod.method, targets = emptyList()) - } - return states to percentageCoverage - } -} diff --git a/usvm-ml-gameserver/src/main/kotlin/org.usvm/OracleImpl.kt b/usvm-ml-gameserver/src/main/kotlin/org.usvm/OracleImpl.kt index 1550fe6101..f40cf7840e 100644 --- a/usvm-ml-gameserver/src/main/kotlin/org.usvm/OracleImpl.kt +++ b/usvm-ml-gameserver/src/main/kotlin/org.usvm/OracleImpl.kt @@ -4,7 +4,7 @@ import org.usvm.gameserver.GameState import org.usvm.statistics.BasicBlock import org.usvm.util.Oracle import org.usvm.utils.Game -import org.usvm.util.createGameState +import org.usvm.gameserver.serialization.createGameState class OracleImpl( val predict: (GameState) -> UInt diff --git a/usvm-ml-gameserver/src/main/kotlin/org.usvm/gameserver/Explorer.kt b/usvm-ml-gameserver/src/main/kotlin/org.usvm/gameserver/Explorer.kt index 1e64efd797..079d347f8c 100644 --- a/usvm-ml-gameserver/src/main/kotlin/org.usvm/gameserver/Explorer.kt +++ b/usvm-ml-gameserver/src/main/kotlin/org.usvm/gameserver/Explorer.kt @@ -1,9 +1,10 @@ -@file:Suppress("UNUSED_PARAMETER", "UNUSED_VARIABLE") - package org.usvm.gameserver -import org.usvm.JavaMethodRunner +import org.usvm.runner.JavaMethodRunner import org.usvm.OracleImpl +import org.usvm.PathSelectionStrategy +import org.usvm.UMachineOptions +import org.usvm.runner.defaultOptions import org.usvm.statistics.BasicBlock import kotlin.math.floor @@ -22,10 +23,10 @@ fun randomExplorer( step.gameStep.stateId } - val (className, methodName) = extractClassAndMethod(gameMap.nameOfObjectToCover) - val runner = JavaMethodRunner(gameMap, OracleImpl(predict)) + val options = cloneDefaultOptions(gameMap, predict) + val runner = JavaMethodRunner(options, gameMap.assemblyFullName) - val (results, percentageCoverage) = runner.cover(className, methodName) + val (results, percentageCoverage) = runner.cover(gameMap.nameOfObjectToCover) val errors = results.count { it.isExceptional } val tests = results.size - errors @@ -36,10 +37,16 @@ fun randomExplorer( ) } -private fun extractClassAndMethod(fullName: String): Pair { - val parts = fullName.split('.') - val className = parts.dropLast(1).joinToString(".") - val methodName = parts.last() - - return Pair(className, methodName) +private fun cloneDefaultOptions(gameMap: GameMap, predict: (GameState) -> UInt): UMachineOptions { + val defaultSearcher = when (gameMap.defaultSearcher) { + Searcher.BFS -> PathSelectionStrategy.BFS + Searcher.DFS -> PathSelectionStrategy.DFS + } + val stepLimit = (gameMap.stepsToStart + gameMap.stepsToStart).toULong() + return defaultOptions.copy( + pathSelectionStrategies = listOf(defaultSearcher, PathSelectionStrategy.AI), + stepLimit = stepLimit, + stepsToStart = gameMap.stepsToStart, + oracle = OracleImpl(predict) + ) } \ No newline at end of file diff --git a/usvm-ml-gameserver/src/main/kotlin/org.usvm/util/Conversion.kt b/usvm-ml-gameserver/src/main/kotlin/org.usvm/gameserver/serialization/Conversion.kt similarity index 98% rename from usvm-ml-gameserver/src/main/kotlin/org.usvm/util/Conversion.kt rename to usvm-ml-gameserver/src/main/kotlin/org.usvm/gameserver/serialization/Conversion.kt index a87bfc67a4..21a937d618 100644 --- a/usvm-ml-gameserver/src/main/kotlin/org.usvm/util/Conversion.kt +++ b/usvm-ml-gameserver/src/main/kotlin/org.usvm/gameserver/serialization/Conversion.kt @@ -1,4 +1,4 @@ -package org.usvm.util +package org.usvm.gameserver.serialization import org.usvm.gameserver.GameEdgeLabel import org.usvm.gameserver.GameMapEdge diff --git a/usvm-runner/build.gradle.kts b/usvm-runner/build.gradle.kts new file mode 100644 index 0000000000..5cbbc53df0 --- /dev/null +++ b/usvm-runner/build.gradle.kts @@ -0,0 +1,32 @@ +import org.jetbrains.kotlin.gradle.tasks.KotlinCompile + +plugins { + id("usvm.kotlin-conventions") + kotlin +} + +repositories { + mavenCentral() +} + +tasks.test { + useJUnitPlatform() +} + +tasks.withType { + kotlinOptions.allWarningsAsErrors = false +} + +dependencies { + implementation(project(":usvm-jvm")) + implementation(project(":usvm-core")) + implementation(Libs.jacodb_api_jvm) + implementation(Libs.jacodb_core) + implementation(Libs.kotlinx_cli) + + implementation(Libs.slf4j_simple) + implementation(Libs.logback) + + implementation("io.github.rchowell:dotlin:1.0.2") + testImplementation(kotlin("test")) +} diff --git a/usvm-runner/src/main/kotlin/org/usvm/runner/JavaMethodRunner.kt b/usvm-runner/src/main/kotlin/org/usvm/runner/JavaMethodRunner.kt new file mode 100644 index 0000000000..79ace440ce --- /dev/null +++ b/usvm-runner/src/main/kotlin/org/usvm/runner/JavaMethodRunner.kt @@ -0,0 +1,57 @@ +package org.usvm.runner + +import org.usvm.CoverageZone +import org.usvm.PathSelectionStrategy +import org.usvm.PathSelectorCombinationStrategy +import org.usvm.UMachineOptions +import org.usvm.machine.JcMachine +import org.usvm.machine.state.JcState +import org.usvm.runner.util.JacoDBContainer +import org.usvm.runner.util.getJcMethodByName +import org.usvm.runner.util.loadClasspathFromJar +import org.usvm.runner.util.mapsKey +import kotlin.time.Duration + +class JavaMethodRunner(val options: UMachineOptions, classpath: String) { + val jacodbCpKey: String + get() = mapsKey + + val cp by lazy { + val classpath = loadClasspathFromJar(classpath) + JacoDBContainer(jacodbCpKey, classpath).cp + } + + fun cover(methodFullName: String): Pair, Float> { + val (className, methodName) = extractClassAndMethod(methodFullName) + return cover(className, methodName) + } + + fun cover(className: String, methodName: String): Pair, Float> { + val jcMethod = cp.getJcMethodByName(className, methodName) + + val (states, percentageCoverage) = JcMachine(cp, options, interpreterObserver = null).use { machine -> + machine.analyze(jcMethod.method, targets = emptyList()) + } + return states to percentageCoverage + } + + private fun extractClassAndMethod(fullName: String): Pair { + val parts = fullName.split('.') + val className = parts.dropLast(1).joinToString(".") + val methodName = parts.last() + + return Pair(className, methodName) + } +} + +val defaultOptions = UMachineOptions( + pathSelectionStrategies = listOf(PathSelectionStrategy.AI), + pathSelectorCombinationStrategy = PathSelectorCombinationStrategy.SEQUENTIAL, + coverageZone = CoverageZone.METHOD, + exceptionsPropagation = true, + solverTimeout = Duration.INFINITE, + timeout = Duration.INFINITE, + typeOperationsTimeout = Duration.INFINITE, + useSolverForForks = false, + stepLimit = 2000U +) \ No newline at end of file diff --git a/usvm-runner/src/main/kotlin/org/usvm/runner/Main.kt b/usvm-runner/src/main/kotlin/org/usvm/runner/Main.kt new file mode 100644 index 0000000000..a0121236fa --- /dev/null +++ b/usvm-runner/src/main/kotlin/org/usvm/runner/Main.kt @@ -0,0 +1,63 @@ +package org.usvm.runner + +import ch.qos.logback.classic.Level +import kotlinx.cli.ArgParser +import kotlinx.cli.ArgType +import kotlinx.cli.default +import kotlinx.cli.required +import org.slf4j.LoggerFactory +import org.usvm.statistics.BasicBlock +import org.usvm.utils.Mode +import org.usvm.utils.OnnxModelImpl + + +fun main(args: Array) { + val parser = ArgParser("usvm.Runner") + val model by parser.option( + ArgType.String, + fullName = "model", + shortName = "m", + description = "Path to ONNX model." + ).required() + + val path by parser.option( + ArgType.String, + fullName = "classpath", + shortName = "cp", + description = "Path to Java classpath." + ).required() + + val mode by parser.option( + ArgType.Choice(), + fullName = "mode", + description = "Mode to run inference in. Could be either CPU or GPU, the default is former." + ).default(Mode.CPU) + + val methodFullName by parser.option( + ArgType.String, + fullName = "method", + description = "Full name of method to cover including package and class name." + ).required() + + parser.parse(args) + + + val runner = JavaMethodRunner( + defaultOptions.copy(oracle = OnnxModelImpl(model, mode)), + classpath = path + ) + + // setLogLevel("ROOT", Level.DEBUG) + + + runner.cover(methodFullName) +} + +// evil logger hack +@Suppress("unused") +private fun setLogLevel(loggerName: String, level: Level) { + val loggerContext = LoggerFactory.getILoggerFactory() as ch.qos.logback.classic.LoggerContext + val logger = loggerContext.getLogger(loggerName) + logger.level = level +} + diff --git a/usvm-ml-gameserver/src/main/kotlin/org.usvm/util/DotGraph.kt b/usvm-runner/src/main/kotlin/org/usvm/runner/util/DotGraph.kt similarity index 99% rename from usvm-ml-gameserver/src/main/kotlin/org.usvm/util/DotGraph.kt rename to usvm-runner/src/main/kotlin/org/usvm/runner/util/DotGraph.kt index fa9a5e3974..9f6a4b49aa 100644 --- a/usvm-ml-gameserver/src/main/kotlin/org.usvm/util/DotGraph.kt +++ b/usvm-runner/src/main/kotlin/org/usvm/runner/util/DotGraph.kt @@ -1,3 +1,5 @@ +package org.usvm.runner.util + import io.github.rchowell.dotlin.DotNodeShape import io.github.rchowell.dotlin.digraph import org.usvm.statistics.BasicBlock diff --git a/usvm-ml-gameserver/src/main/kotlin/org.usvm/util/JacoDBContainer.kt b/usvm-runner/src/main/kotlin/org/usvm/runner/util/JacoDBContainer.kt similarity index 97% rename from usvm-ml-gameserver/src/main/kotlin/org.usvm/util/JacoDBContainer.kt rename to usvm-runner/src/main/kotlin/org/usvm/runner/util/JacoDBContainer.kt index 2572174503..1baadecd10 100644 --- a/usvm-ml-gameserver/src/main/kotlin/org.usvm/util/JacoDBContainer.kt +++ b/usvm-runner/src/main/kotlin/org/usvm/runner/util/JacoDBContainer.kt @@ -1,4 +1,4 @@ -package org.usvm.util +package org.usvm.runner.util import kotlinx.coroutines.runBlocking import org.jacodb.api.jvm.JcClasspath diff --git a/usvm-ml-gameserver/src/main/kotlin/org.usvm/util/Util.kt b/usvm-runner/src/main/kotlin/org/usvm/runner/util/Util.kt similarity index 95% rename from usvm-ml-gameserver/src/main/kotlin/org.usvm/util/Util.kt rename to usvm-runner/src/main/kotlin/org/usvm/runner/util/Util.kt index a4c955eb20..bc9816102a 100644 --- a/usvm-ml-gameserver/src/main/kotlin/org.usvm/util/Util.kt +++ b/usvm-runner/src/main/kotlin/org/usvm/runner/util/Util.kt @@ -1,4 +1,4 @@ -package org.usvm.util +package org.usvm.runner.util import org.jacodb.api.jvm.JcClasspath import org.jacodb.api.jvm.JcTypedMethod diff --git a/usvm-runner/src/main/resources/logback.xml b/usvm-runner/src/main/resources/logback.xml new file mode 100644 index 0000000000..7d47eaa2ba --- /dev/null +++ b/usvm-runner/src/main/resources/logback.xml @@ -0,0 +1,11 @@ + + + + %d{YYYY-MM-dd HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n + + + + + + + \ No newline at end of file