NeuroSMT solver

ksmt-neurosmt/build.gradle.kts

@@ -0,0 +1,39 @@
+plugins {
+    id("io.ksmt.ksmt-base")
+}
+
+repositories {
+    mavenCentral()
+}
+
+dependencies {
+    implementation(project(":ksmt-core"))
+    implementation(project("utils"))
+
+    implementation("com.microsoft.onnxruntime:onnxruntime:1.15.1")
+    implementation("com.github.ajalt.clikt:clikt:3.5.2")
+}
+
+tasks {
+    val standaloneModelRunnerFatJar = register<Jar>("standaloneModelRunnerFatJar") {
+        dependsOn.addAll(listOf("compileJava", "compileKotlin", "processResources"))
+
+        archiveFileName.set("standalone-model-runner.jar")
+        destinationDirectory.set(File("."))
+        duplicatesStrategy = DuplicatesStrategy.EXCLUDE
+
+        manifest {
+            attributes(mapOf("Main-Class" to "io.ksmt.solver.neurosmt.runtime.standalone.StandaloneModelRunnerKt"))
+        }
+
+        val sourcesMain = sourceSets.main.get()
+        val contents = configurations.runtimeClasspath.get()
+            .map { if (it.isDirectory) it else zipTree(it) } + sourcesMain.output
+
+        from(contents)
+    }
+
+    build {
+        dependsOn(standaloneModelRunnerFatJar)
+    }
+}

ksmt-neurosmt/src/main/kotlin/io/ksmt/solver/neurosmt/KNeuroSMTSolver.kt

@@ -0,0 +1,87 @@
+package io.ksmt.solver.neurosmt
+
+import io.ksmt.KContext
+import io.ksmt.expr.KExpr
+import io.ksmt.solver.KModel
+import io.ksmt.solver.KSolver
+import io.ksmt.solver.KSolverStatus
+import io.ksmt.solver.neurosmt.runtime.NeuroSMTModelRunner
+import io.ksmt.sort.KBoolSort
+import kotlin.time.Duration
+
+class KNeuroSMTSolver(
+    private val ctx: KContext,
+    ordinalsPath: String, embeddingPath: String, convPath: String, decoderPath: String,
+    private val threshold: Double = 0.5
+) : KSolver<KNeuroSMTSolverConfiguration> {
+
+    private val modelRunner = NeuroSMTModelRunner(ctx, ordinalsPath, embeddingPath, convPath, decoderPath)
+    private val asserts = mutableListOf<MutableList<KExpr<KBoolSort>>>(mutableListOf())
+
+    override fun configure(configurator: KNeuroSMTSolverConfiguration.() -> Unit) {
+        TODO("Not yet implemented")
+    }
+
+    override fun assert(expr: KExpr<KBoolSort>) {
+        asserts.last().add(expr)
+    }
+
+    override fun assertAndTrack(expr: KExpr<KBoolSort>) {
+        assert(expr)
+    }
+
+    override fun push() {
+        asserts.add(mutableListOf())
+    }
+
+    override fun pop(n: UInt) {
+        repeat(n.toInt()) {
+            asserts.removeLast()
+        }
+    }
+
+    override fun check(timeout: Duration): KSolverStatus {
+        val prob = with(ctx) {
+            modelRunner.run(mkAnd(asserts.flatten()))
+        }
+
+        return if (prob > threshold) {
+            KSolverStatus.SAT
+        } else {
+            KSolverStatus.UNSAT
+        }
+    }
+
+    override fun checkWithAssumptions(assumptions: List<KExpr<KBoolSort>>, timeout: Duration): KSolverStatus {
+        val prob = with(ctx) {
+            modelRunner.run(mkAnd(asserts.flatten() + assumptions))
+        }
+
+        return if (prob > threshold) {
+            KSolverStatus.SAT
+        } else {
+            KSolverStatus.UNSAT
+        }
+    }
+
+    override fun model(): KModel {
+        TODO("Not yet implemented")
+    }
+
+    override fun unsatCore(): List<KExpr<KBoolSort>> {
+        TODO("Not yet implemented")
+    }
+
+    override fun reasonOfUnknown(): String {
+        TODO("Not yet implemented")
+    }
+
+    override fun interrupt() {
+        TODO("Not yet implemented")
+    }
+
+    override fun close() {
+        modelRunner.close()
+        asserts.clear()
+    }
+}

ksmt-neurosmt/src/main/kotlin/io/ksmt/solver/neurosmt/KNeuroSMTSolverConfiguration.kt

@@ -0,0 +1,44 @@
+package io.ksmt.solver.neurosmt
+
+import io.ksmt.solver.KSolverConfiguration
+
+interface KNeuroSMTSolverConfiguration : KSolverConfiguration {
+    fun setOption(option: String, value: Boolean)
+    fun setOption(option: String, value: Int)
+    fun setOption(option: String, value: Double)
+    fun setOption(option: String, value: String)
+
+    override fun setBoolParameter(param: String, value: Boolean) {
+        setOption(param, value)
+    }
+
+    override fun setIntParameter(param: String, value: Int) {
+        setOption(param, value)
+    }
+
+    override fun setDoubleParameter(param: String, value: Double) {
+        setOption(param, value)
+    }
+
+    override fun setStringParameter(param: String, value: String) {
+        setOption(param, value)
+    }
+}
+
+class KNeuroSMTSolverConfigurationImpl(private val params: Any?) : KNeuroSMTSolverConfiguration {
+    override fun setOption(option: String, value: Boolean) {
+        TODO("Not yet implemented")
+    }
+
+    override fun setOption(option: String, value: Int) {
+        TODO("Not yet implemented")
+    }
+
+    override fun setOption(option: String, value: Double) {
+        TODO("Not yet implemented")
+    }
+
+    override fun setOption(option: String, value: String) {
+        TODO("Not yet implemented")
+    }
+}

ksmt-neurosmt/src/main/kotlin/io/ksmt/solver/neurosmt/runtime/ExprEncoder.kt

@@ -0,0 +1,126 @@
+package io.ksmt.solver.neurosmt.runtime
+
+import ai.onnxruntime.OnnxTensor
+import ai.onnxruntime.OrtEnvironment
+import io.ksmt.KContext
+import io.ksmt.expr.KApp
+import io.ksmt.expr.KConst
+import io.ksmt.expr.KExpr
+import io.ksmt.expr.KInterpretedValue
+import io.ksmt.expr.transformer.KNonRecursiveTransformer
+import io.ksmt.sort.*
+import java.nio.FloatBuffer
+import java.nio.IntBuffer
+import java.nio.LongBuffer
+import java.util.*
+
+// expression encoder
+// walks on an expression and calculates state for each node
+// based on states for its children (which are already calculated at that moment)
+class ExprEncoder(
+    override val ctx: KContext,
+    val env: OrtEnvironment,
+    val ordinalEncoder: OrdinalEncoder,
+    val embeddingLayer: ONNXModel,
+    val convLayer: ONNXModel
+) : KNonRecursiveTransformer(ctx) {
+
+    private val exprToState = IdentityHashMap<KExpr<*>, OnnxTensor>()
+
+    fun encodeExpr(expr: KExpr<*>): OnnxTensor {
+        apply(expr)
+
+        return exprToState[expr] ?: error("expression state wasn't calculated yet")
+    }
+
+    override fun <T : KSort, A : KSort> transformApp(expr: KApp<T, A>): KExpr<T> {
+        val state = when (expr) {
+            is KConst<*> -> calcSymbolicVariableState(expr)
+            is KInterpretedValue<*> -> calcValueState(expr)
+            else -> calcAppState(expr)
+        }
+
+        exprToState[expr] = state
+
+        return expr
+    }
+
+    private fun getNodeEmbedding(key: String): OnnxTensor {
+        val nodeLabel = ordinalEncoder.getOrdinal(key)
+        val labelTensor = OnnxTensor.createTensor(
+            env, IntBuffer.allocate(1).put(nodeLabel).rewind(), longArrayOf(1, 1)
+        )
+
+        return embeddingLayer.forward(mapOf("node_labels" to labelTensor))
+    }
+
+    private fun createEdgeTensor(childrenCnt: Int): OnnxTensor {
+        val edges = listOf(
+            List(childrenCnt) { it + 1L },
+            List(childrenCnt) { 0L }
+        )
+
+        val buffer = LongBuffer.allocate(childrenCnt * 2)
+        edges.forEach { row ->
+            row.forEach { node ->
+                buffer.put(node)
+            }
+        }
+        buffer.rewind()
+
+        return OnnxTensor.createTensor(env, buffer, longArrayOf(2, childrenCnt.toLong()))
+    }
+
+    private fun <T : KSort, A : KSort> calcAppState(expr: KApp<T, A>): OnnxTensor {
+        val childrenStates = expr.args.map { exprToState[it] ?: error("expression state wasn't calculated yet") }
+        val childrenCnt = childrenStates.size
+
+        val nodeEmbedding = getNodeEmbedding(expr.decl.name)
+        val embeddingSize = nodeEmbedding.info.shape.reduce { acc, l -> acc * l }
+
+        val buffer = FloatBuffer.allocate((1 + childrenCnt) * embeddingSize.toInt())
+        buffer.put(nodeEmbedding.floatBuffer)
+        childrenStates.forEach {
+            buffer.put(it.floatBuffer)
+        }
+        buffer.rewind()
+
+        val nodeFeatures = OnnxTensor.createTensor(env, buffer, longArrayOf(1L + childrenCnt, embeddingSize))
+        val edges = createEdgeTensor(childrenStates.size)
+        val result = convLayer.forward(mapOf("node_features" to nodeFeatures, "edges" to edges))
+
+        return OnnxTensor.createTensor(env, result.floatBuffer.slice(0, embeddingSize.toInt()), longArrayOf(1L, embeddingSize))
+    }
+
+    private fun <T : KSort> calcSymbolicVariableState(symbol: KConst<T>): OnnxTensor {
+        val sort = symbol.decl.sort
+
+        val key = "SYMBOLIC;" + when (sort) {
+            is KBoolSort -> "Bool"
+            is KBvSort -> "BitVec"
+            is KFpSort -> "FP"
+            is KFpRoundingModeSort -> "FP_RM"
+            is KArraySortBase<*> -> "Array"
+            is KUninterpretedSort -> sort.name
+            else -> error("unknown symbolic sort: ${sort::class.simpleName}")
+        }
+
+        return getNodeEmbedding(key)
+    }
+
+    private fun <T : KSort> calcValueState(value: KInterpretedValue<T>): OnnxTensor {
+        val sort = value.decl.sort
+
+        val key = "VALUE;" + when (sort) {
+            is KBoolSort -> "Bool"
+            is KBvSort -> "BitVec"
+            is KFpSort -> "FP"
+            is KFpRoundingModeSort -> "FP_RM"
+            is KArraySortBase<*> -> "Array"
+            is KUninterpretedSort -> sort.name
+            else -> error("unknown value sort: ${value.decl.sort::class.simpleName}")
+        }
+
+        return getNodeEmbedding(key)
+    }
+}

ksmt-neurosmt/src/main/kotlin/io/ksmt/solver/neurosmt/runtime/NeuroSMTModelRunner.kt

@@ -0,0 +1,36 @@
+package io.ksmt.solver.neurosmt.runtime
+
+import ai.onnxruntime.OrtEnvironment
+import io.ksmt.KContext
+import io.ksmt.expr.KExpr
+import kotlin.math.exp
+
+// wrapper for NeuroSMT model
+// runs the whole model pipeline
+class NeuroSMTModelRunner(
+    val ctx: KContext,
+    ordinalsPath: String, embeddingPath: String, convPath: String, decoderPath: String
+) {
+    val env = OrtEnvironment.getEnvironment()
+
+    val ordinalEncoder = OrdinalEncoder(ordinalsPath)
+    val embeddingLayer = ONNXModel(env, embeddingPath)
+    val convLayer = ONNXModel(env, convPath)
+
+    val decoder = ONNXModel(env, decoderPath)
+
+    fun run(expr: KExpr<*>): Float {
+        val encoder = ExprEncoder(ctx, env, ordinalEncoder, embeddingLayer, convLayer)
+        val exprFeatures = encoder.encodeExpr(expr)
+        val result = decoder.forward(mapOf("expr_features" to exprFeatures))
+        val logit = result.floatBuffer[0]
+
+        return 1f / (1f + exp(-logit)) // sigmoid calculation
+    }
+
+    fun close() {
+        embeddingLayer.close()
+        convLayer.close()
+        decoder.close()
+    }
+}

ksmt-neurosmt/src/main/kotlin/io/ksmt/solver/neurosmt/runtime/ONNXModel.kt

@@ -0,0 +1,18 @@
+package io.ksmt.solver.neurosmt.runtime
+
+import ai.onnxruntime.OnnxTensor
+import ai.onnxruntime.OrtEnvironment
+
+// wrapper for any exported ONNX model
+class ONNXModel(env: OrtEnvironment, modelPath: String) : AutoCloseable {
+    val session = env.createSession(modelPath)
+
+    fun forward(input: Map<String, OnnxTensor>): OnnxTensor {
+        val result = session.run(input)
+        return result.get(0) as OnnxTensor
+    }
+
+    override fun close() {
+        session.close()
+    }
+}

ksmt-neurosmt/src/main/kotlin/io/ksmt/solver/neurosmt/runtime/OrdinalEncoder.kt

@@ -0,0 +1,23 @@
+package io.ksmt.solver.neurosmt.runtime
+
+import java.nio.file.Files
+import java.nio.file.Path
+import kotlin.streams.asSequence
+
+const val UNKNOWN_VALUE = 1999
+
+// wrapper for single-feature sklearn OrdinalEncoder (for each string we should provide its ordinal)
+// used to convert strings to integers
+class OrdinalEncoder(ordinalsPath: String, private val unknownValue: Int = UNKNOWN_VALUE) {
+    private val lookup = HashMap<String, Int>()
+
+    init {
+        Files.lines(Path.of(ordinalsPath)).asSequence().forEachIndexed { index, s ->
+            lookup[s] = index
+        }
+    }
+
+    fun getOrdinal(s: String): Int {
+        return lookup[s] ?: unknownValue
+    }
+}