kscience/kmath
9
0
Fork 1
Code Issues Pull Requests 6 Packages Projects Releases 16 Wiki Activity

Very experimental WASM code generation by MST in contexts of Int and Real #158

Closed
CommanderTvis wants to merge 16 commits from feature/binaryen into dev
pull from: feature/binaryen
merge into: kscience:dev
Conversation 1 Commits 16 Files Changed 17 +134 -65
2 changed files with 134 additions and 65 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit 4f22688f6f - Show all commits

2
kmath-ast/src/jsMain/kotlin/kscience/kmath/ast/Initial.kt
View File

File diff suppressed because one or more lines are too long

197
kmath-ast/src/jsMain/kotlin/kscience/kmath/ast/RealWasmBuilder.kt
View File

@ -6,77 +6,146 @@ import kscience.kmath.operations.*
private val spreader = eval("(obj, args) => obj(...args)")
public fun compileMstToWasmF64(mst: MST): Expression<Double> {
val keys = mutableListOf<String>()
internal sealed class WasmBuilder<T>(val binaryenType: binaryen.Type, val kmathAlgebra: Algebra<T>) where T : Number {
val keys: MutableList<String> = mutableListOf()
lateinit var ctx: binaryen.Module
val bin = with(binaryen.readBinary(INITIAL)) {
fun MST.visit(): binaryen.ExpressionRef = when (this) {
is MST.Symbolic -> {
var idx = keys.indexOf(value)
open fun visitSymbolic(mst: MST.Symbolic): binaryen.ExpressionRef {
try {
kmathAlgebra.symbol(mst.value)
} catch (ignored: Throwable) {
null
}?.let { return visitNumeric(MST.Numeric(it)) }
if (idx == -1) {
keys += value
idx = keys.lastIndex
}
var idx = keys.indexOf(mst.value)
local.get(idx, binaryen.f64)
}
is MST.Numeric -> f64.const(value)
is MST.Unary -> when (operation) {
SpaceOperations.MINUS_OPERATION -> f64.neg(value.visit())
SpaceOperations.PLUS_OPERATION -> value.visit()
PowerOperations.SQRT_OPERATION -> f64.sqrt(value.visit())
TrigonometricOperations.SIN_OPERATION -> call("sin", arrayOf(value.visit()), binaryen.f64)
TrigonometricOperations.COS_OPERATION -> call("cos", arrayOf(value.visit()), binaryen.f64)
TrigonometricOperations.TAN_OPERATION -> call("tan", arrayOf(value.visit()), binaryen.f64)
TrigonometricOperations.ASIN_OPERATION -> call("asin", arrayOf(value.visit()), binaryen.f64)
TrigonometricOperations.ACOS_OPERATION -> call("acos", arrayOf(value.visit()), binaryen.f64)
TrigonometricOperations.ATAN_OPERATION -> call("atan", arrayOf(value.visit()), binaryen.f64)
HyperbolicOperations.SINH_OPERATION -> call("sinh", arrayOf(value.visit()), binaryen.f64)
HyperbolicOperations.COSH_OPERATION -> call("cosh", arrayOf(value.visit()), binaryen.f64)
HyperbolicOperations.TANH_OPERATION -> call("tanh", arrayOf(value.visit()), binaryen.f64)
HyperbolicOperations.ASINH_OPERATION -> call("asinh", arrayOf(value.visit()), binaryen.f64)
HyperbolicOperations.ACOSH_OPERATION -> call("acosh", arrayOf(value.visit()), binaryen.f64)
HyperbolicOperations.ATANH_OPERATION -> call("atanh", arrayOf(value.visit()), binaryen.f64)
ExponentialOperations.EXP_OPERATION -> call("exp", arrayOf(value.visit()), binaryen.f64)
ExponentialOperations.LN_OPERATION -> call("log", arrayOf(value.visit()), binaryen.f64)
else -> throw UnsupportedOperationException()
}
is MST.Binary -> when (operation) {
SpaceOperations.PLUS_OPERATION -> f64.add(left.visit(), right.visit())
RingOperations.TIMES_OPERATION -> f64.mul(left.visit(), right.visit())
FieldOperations.DIV_OPERATION -> f64.div(left.visit(), right.visit())
PowerOperations.POW_OPERATION -> call("pow", arrayOf(left.visit(), right.visit()), binaryen.f64)
else -> throw UnsupportedOperationException()
}
if (idx == -1) {
keys += mst.value
idx = keys.lastIndex
}
val expr = mst.visit()
addFunction(
"executable",
binaryen.createType(Array(keys.size) { binaryen.f64 }),
binaryen.f64,
arrayOf(),
expr
)
binaryen.setOptimizeLevel(3)
// optimizeFunction("executable")
addFunctionExport("executable", "executable")
val res = emitBinary()
dispose()
res
return ctx.local.get(idx, binaryenType)
}
val c = WebAssembly.Module(bin)
val i = WebAssembly.Instance(c, js("{}") as Any)
abstract fun visitNumeric(mst: MST.Numeric): binaryen.ExpressionRef
return Expression { args ->
val params = keys.map { StringSymbol(it) }.map { args.getValue(it) }.toTypedArray()
spreader(i.exports.asDynamic().executable, params) as Double
open fun visitUnary(mst: MST.Unary): binaryen.ExpressionRef =
error("Unary operation ${mst.operation} not defined in $this")
open fun visitBinary(mst: MST.Binary): binaryen.ExpressionRef =
error("Binary operation ${mst.operation} not defined in $this")
open fun createModule(): binaryen.Module = binaryen.Module()
fun visit(mst: MST): binaryen.ExpressionRef = when (mst) {
is MST.Symbolic -> visitSymbolic(mst)
is MST.Numeric -> visitNumeric(mst)
is MST.Unary -> visitUnary(mst)
is MST.Binary -> visitBinary(mst)
}
fun compile(mst: MST): Expression<T> {
val keys = mutableListOf<String>()
val bin = with(createModule()) {
ctx = this
val expr = visit(mst)
addFunction(
"executable",
binaryen.createType(Array(keys.size) { binaryenType }),
binaryenType,
arrayOf(),
expr
)
binaryen.setOptimizeLevel(3)
// optimizeFunction("executable")
addFunctionExport("executable", "executable")
val res = emitBinary()
dispose()
res
}
val c = WebAssembly.Module(bin)
val i = WebAssembly.Instance(c, js("{}") as Any)
return Expression { args ->
val params = keys.map { StringSymbol(it) }.map { args.getValue(it) }.toTypedArray()
keys.clear()
spreader(i.exports.asDynamic().executable, params) as T
}
}
}
internal class RealWasmBuilder : WasmBuilder<Double>(binaryen.f64, RealField) {
override fun createModule(): binaryen.Module = binaryen.readBinary(f64StandardFunctions)
override fun visitNumeric(mst: MST.Numeric): binaryen.ExpressionRef = ctx.f64.const(mst.value)
override fun visitUnary(mst: MST.Unary): binaryen.ExpressionRef = when (mst.operation) {
SpaceOperations.MINUS_OPERATION -> ctx.f64.neg(visit(mst.value))
SpaceOperations.PLUS_OPERATION -> visit(mst.value)
PowerOperations.SQRT_OPERATION -> ctx.f64.sqrt(visit(mst.value))
TrigonometricOperations.SIN_OPERATION -> ctx.call("sin", arrayOf(visit(mst.value)), binaryen.f64)
TrigonometricOperations.COS_OPERATION -> ctx.call("cos", arrayOf(visit(mst.value)), binaryen.f64)
TrigonometricOperations.TAN_OPERATION -> ctx.call("tan", arrayOf(visit(mst.value)), binaryen.f64)
TrigonometricOperations.ASIN_OPERATION -> ctx.call("asin", arrayOf(visit(mst.value)), binaryen.f64)
TrigonometricOperations.ACOS_OPERATION -> ctx.call("acos", arrayOf(visit(mst.value)), binaryen.f64)
TrigonometricOperations.ATAN_OPERATION -> ctx.call("atan", arrayOf(visit(mst.value)), binaryen.f64)
HyperbolicOperations.SINH_OPERATION -> ctx.call("sinh", arrayOf(visit(mst.value)), binaryen.f64)
HyperbolicOperations.COSH_OPERATION -> ctx.call("cosh", arrayOf(visit(mst.value)), binaryen.f64)
HyperbolicOperations.TANH_OPERATION -> ctx.call("tanh", arrayOf(visit(mst.value)), binaryen.f64)
HyperbolicOperations.ASINH_OPERATION -> ctx.call("asinh", arrayOf(visit(mst.value)), binaryen.f64)
HyperbolicOperations.ACOSH_OPERATION -> ctx.call("acosh", arrayOf(visit(mst.value)), binaryen.f64)
HyperbolicOperations.ATANH_OPERATION -> ctx.call("atanh", arrayOf(visit(mst.value)), binaryen.f64)
ExponentialOperations.EXP_OPERATION -> ctx.call("exp", arrayOf(visit(mst.value)), binaryen.f64)
ExponentialOperations.LN_OPERATION -> ctx.call("log", arrayOf(visit(mst.value)), binaryen.f64)
else -> super.visitUnary(mst)
}
override fun visitBinary(mst: MST.Binary): binaryen.ExpressionRef = when (mst.operation) {
SpaceOperations.PLUS_OPERATION -> ctx.f64.add(visit(mst.left), visit(mst.right))
SpaceOperations.MINUS_OPERATION -> ctx.f64.sub(visit(mst.left), visit(mst.right))
RingOperations.TIMES_OPERATION -> ctx.f64.mul(visit(mst.left), visit(mst.right))
FieldOperations.DIV_OPERATION -> ctx.f64.div(visit(mst.left), visit(mst.right))
PowerOperations.POW_OPERATION -> ctx.call("pow", arrayOf(visit(mst.left), visit(mst.right)), binaryen.f64)
else -> super.visitBinary(mst)
}
}
internal class IntWasmBuilder : WasmBuilder<Int>(binaryen.i32, IntRing) {
override fun visitNumeric(mst: MST.Numeric): binaryen.ExpressionRef = ctx.i32.const(mst.value)
override fun visitUnary(mst: MST.Unary): binaryen.ExpressionRef = when (mst.operation) {
SpaceOperations.MINUS_OPERATION -> ctx.i32.sub(ctx.i32.const(0), visit(mst.value))
SpaceOperations.PLUS_OPERATION -> visit(mst.value)
else -> super.visitUnary(mst)
}
override fun visitBinary(mst: MST.Binary): binaryen.ExpressionRef = when (mst.operation) {
SpaceOperations.PLUS_OPERATION -> ctx.i32.add(visit(mst.left), visit(mst.right))
SpaceOperations.MINUS_OPERATION -> ctx.i32.sub(visit(mst.left), visit(mst.right))
RingOperations.TIMES_OPERATION -> ctx.i32.mul(visit(mst.left), visit(mst.right))
FieldOperations.DIV_OPERATION -> ctx.i32.div_s(visit(mst.left), visit(mst.right))
else -> super.visitBinary(mst)
}
}
internal class LongWasmBuilder : WasmBuilder<Long>(binaryen.i64, LongRing) {
override fun visitNumeric(mst: MST.Numeric): binaryen.ExpressionRef = ctx.i64.const(mst.value)
override fun visitUnary(mst: MST.Unary): binaryen.ExpressionRef = when (mst.operation) {
SpaceOperations.MINUS_OPERATION -> ctx.i64.sub(ctx.i64.const(0, 0), visit(mst.value))
SpaceOperations.PLUS_OPERATION -> visit(mst.value)
else -> super.visitUnary(mst)
}
override fun visitBinary(mst: MST.Binary): binaryen.ExpressionRef = when (mst.operation) {
SpaceOperations.PLUS_OPERATION -> ctx.i64.add(visit(mst.left), visit(mst.right))
SpaceOperations.MINUS_OPERATION -> ctx.i64.sub(visit(mst.left), visit(mst.right))
RingOperations.TIMES_OPERATION -> ctx.i64.mul(visit(mst.left), visit(mst.right))
FieldOperations.DIV_OPERATION -> ctx.i64.div_s(visit(mst.left), visit(mst.right))
else -> super.visitBinary(mst)
}
}