Helpers to compile to WebAssembly (#114)

Author: tshort

.github/workflows/ci.yml

@@ -19,7 +19,7 @@ jobs:
       matrix:
         version:
           - '1.8'
-          - '^1.9.0-rc2'
+          - '^1.9.0-rc3'
         os:
           - ubuntu-latest
           - macOS-latest

src/StaticCompiler.jl

@@ -8,10 +8,11 @@ using Libdl: Libdl, dlsym, dlopen
 using Base: RefValue
 using Serialization: serialize, deserialize
 using Clang_jll: clang
+using LLD_jll: lld
 using StaticTools
 using StaticTools: @symbolcall, @c_str, println
 
-export compile, load_function, compile_shlib, compile_executable
+export compile, load_function, compile_shlib, compile_executable, compile_wasm
 export native_code_llvm, native_code_typed, native_llvm_module, native_code_native
 export @device_override, @print_and_throw
 
@@ -154,7 +155,8 @@ function generate_obj(f, tt, external = true, path::String = tempname(), name =
                         mixtape = NoContext(),
                         strip_llvm = false,
                         strip_asm  = true,
-                        opt_level=3,
+                        opt_level = 3,
+                        remove_julia_addrspaces = false,
                         target = (),
                         kwargs...)
     mkpath(path)
@@ -180,7 +182,7 @@ function generate_obj(f, tt, external = true, path::String = tempname(), name =
 
     # Now that we've removed all the pointers from the code, we can (hopefully) safely lower all the instrinsics
     # (again, using Enzyme's pipeline)
-    post_optimize!(mod, tm)
+    post_optimize!(mod, tm; remove_julia_addrspaces)
 
     # Make sure we didn't make any glaring errors
     LLVM.verify(mod)
@@ -355,6 +357,40 @@ function compile_shlib(funcs::Array, path::String="./";
     joinpath(abspath(path), filename * "." * Libdl.dlext)
 end
 
+"""
+```julia
+compile_wasm(f::Function, types::Tuple, [path::String="./"], [name::String=repr(f)]; filename::String=name, flags=``, kwargs...)
+compile_wasm(funcs::Array, [path::String="./"]; filename="libfoo", demangle=false, flags=``, kwargs...)
+```
+As `compile_shlib`, but compiling to a WebAssembly library.
+
+The compiled function is by default given the symbol name `julia_$(name)`, i.e.,
+the function `test` in the example below is called `julia_test` in the shared library.
+The keword argument `demangle=true` will remove this prefix, but is currently only
+supported the second (multi-function-shlib) method.
+```
+"""
+function compile_wasm(f::Function, types=(); 
+        path::String="./",
+        filename=fix_name(repr(f)),
+        flags=``,
+        kwargs...
+    )
+    tt = Base.to_tuple_type(types)
+    obj_path, name = generate_obj(f, tt, true, path, filename; target = (triple = "wasm32-unknown-wasi", cpu = "", features = ""), remove_julia_addrspaces = true, kwargs...)
+    run(`$(lld()) -flavor wasm --no-entry --export-all $flags $obj_path/obj.o -o $path/$name.wasm`)
+    joinpath(abspath(path), filename * ".wasm")
+end
+function compile_wasm(funcs::Array; 
+        path::String="./",
+        filename="libfoo",
+        flags=``,
+        kwargs...
+    )
+    obj_path, name = generate_obj(funcs, true; target = (triple = "wasm32-unknown-wasi", cpu = "", features = ""), remove_julia_addrspaces = true, kwargs...)
+    run(`$(lld()) -flavor wasm --no-entry --export-all $flags $obj_path/obj.o -o $path/$filename.wasm`)
+    joinpath(abspath(path), filename * ".wasm")
+end
 
 """
 ```julia
@@ -606,7 +642,7 @@ function native_llvm_module(funcs::Array; demangle = false, kwargs...)
     return mod
 end
 
-function generate_obj(funcs::Array, external, path::String = tempname(), filenamebase::String="obj";
+function generate_obj(funcs::Array, external::Bool, path::String = tempname(), filenamebase::String="obj";
                         demangle =false,
                         strip_llvm = false,
                         strip_asm  = true,
@@ -615,7 +651,7 @@ function generate_obj(funcs::Array, external, path::String = tempname(), filenam
     f, tt = funcs[1]
     mkpath(path)
     obj_path = joinpath(path, "$filenamebase.o")
-    fakejob, kwargs = native_job(f, tt, external, kwargs...)
+    fakejob, kwargs = native_job(f, tt, external; kwargs...)
     mod = native_llvm_module(funcs; demangle = demangle, kwargs...)
     obj, _ = GPUCompiler.emit_asm(fakejob, mod; strip=strip_asm, validate=false, format=LLVM.API.LLVMObjectFile)
     open(obj_path, "w") do io

src/optimize.jl

@@ -211,8 +211,8 @@ function addJuliaLegalizationPasses!(pm, lower_intrinsics=true)
         remove_ni!(pm)
     end
 end
-
-function post_optimize!(mod, tm)
+                        
+function post_optimize!(mod, tm; remove_julia_addrspaces = false)
     # @show "pre_post", mod
     # flush(stdout)
     # flush(stderr)
@@ -224,6 +224,9 @@ function post_optimize!(mod, tm)
     LLVM.ModulePassManager() do pm
         addJuliaLegalizationPasses!(pm, true)
         addMachinePasses!(pm)
+        if remove_julia_addrspaces
+            remove_julia_addrspaces!(pm)
+        end
         run!(pm, mod)
     end
     # @show "post_mod", mod

test/testintegration.jl

@@ -386,12 +386,12 @@ end
 @testset "Cross compiling to WebAssembly" begin
 
     m2(x) = 2x
-    obj_path, name = StaticCompiler.generate_obj(m2, Tuple{Float64}, true, tempname(); target = (triple = "wasm32-unknown-unknown", cpu = "", features = ""))
-    # link with `lld` from LLD_jll
-    run(`$(lld()) -flavor wasm --no-entry --export-all $obj_path/obj.o -o $name.wasm`)
-    # On Julia v1.9, link with included linker
-    # run(`$(Base.Linking.lld()) -flavor wasm --no-entry --export-all $obj_path/obj.o -o $name.wasm`)
-    # run(`wasm2wat $name.wasm`)   # to see a text representation (wasm2wat isn't included)
+    m3(x) = 3x
+    wasm_path = compile_wasm(m2, Tuple{Float64})
+    wasm_path2 = compile_wasm([(m2, Tuple{Float64}), (m3, Tuple{Float64})])
+
+    wasm_path = compile_wasm(m2, (Float64,))
+    wasm_path2 = compile_wasm([(m2, (Float64,)), (m3, (Float64,))])
 
 end