Cassette-type functionality with Mixtape (#69)

Author: tshort

.github/workflows/ci-integration-nightly.yml

@@ -18,8 +18,8 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.7'
           - '1.8'
+          - '^1.9.0-rc2'
         os:
           - ubuntu-latest
           - macOS-latest

.github/workflows/ci-integration.yml

@@ -1,6 +1,6 @@
 Copyright (c) 2019-2022 Mason Protter, William Moses, Valentin Churavy, 
-                        Brenhin Keller, Julian Samaroo, Tom Short, and 
-                        other contributors
+                        McCoy R. Becker, Brenhin Keller, Julian Samaroo,
+                        Tom Short, and other contributors
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

.github/workflows/ci-julia-nightly.yml

@@ -1,18 +1,30 @@
 name = "StaticCompiler"
 uuid = "81625895-6c0f-48fc-b932-11a18313743c"
 authors = ["Tom Short and contributors"]
-version = "0.4.9"
+version = "0.4.10"
 
 [deps]
 Clang_jll = "0ee61d77-7f21-5576-8119-9fcc46b10100"
+CodeInfoTools = "bc773b8a-8374-437a-b9f2-0e9785855863"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 StaticTools = "86c06d3c-3f03-46de-9781-57580aa96d0a"
 
 [compat]
-GPUCompiler = "0.16, 0.17"
+CodeInfoTools = "0.3"
+GPUCompiler = "0.17"
 LLVM = "4.8, 5"
 StaticTools = "0.8"
-julia = "1.7"
+MacroTools = "0.5"
+julia = "1.8, 1.9"
+
+[extras]
+Formatting = "59287772-0a20-5a39-b81b-1366585eb4c0"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["Test", "Formatting"]

.github/workflows/ci.yml

@@ -65,6 +65,14 @@ Hello, world!
 ```
 This latter approach comes with substantially more limitations, as you cannot rely on `libjulia` (see, e.g., [StaticTools.jl](https://github.com/brenhinkeller/StaticTools.jl) for some ways to work around these limitations).
 
+### Mixtape
+
+This feature allows one to change functionality when statically compiling. This uses code and API from [Mixtape](https://github.com/JuliaCompilerPlugins/Mixtape.jl) to transform lowered code much like [Cassette](https://github.com/JuliaLabs/Cassette.jl).
+
+To use the Mixtape feature, define a `CompilationContext` struct and pass this to any of the compilation functions with the `mixtape` keyword. Define `transform` and `allow` functions for this `CompilationContext` to define the transformation to be done.
+
+See [here](https://github.com/tshort/StaticCompiler.jl/blob/master/test/testintegration.jl#L329) for an example.
+
 ## Approach
 
 This package uses the [GPUCompiler package](https://github.com/JuliaGPU/GPUCompiler.jl) to generate code.

LICENSE

@@ -1,5 +1,5 @@
 module StaticCompiler
-
+using InteractiveUtils
 using GPUCompiler: GPUCompiler
 using LLVM
 using LLVM.Interop
@@ -15,6 +15,8 @@ export compile, load_function, compile_shlib, compile_executable
 export native_code_llvm, native_code_typed, native_llvm_module, native_code_native
 export @device_override, @print_and_throw
 
+include("mixtape.jl")
+include("interpreter.jl")
 include("target.jl")
 include("pointer_patching.jl")
 include("code_loading.jl")
@@ -87,19 +89,21 @@ with `load_function(path)`. `LazyStaticcompiledfunction`s contain the requisite
 `StaticCompiledFunction` and may be called with arguments of type `types` as if it were a function with a
 single method (the method determined by `types`).
 """
-function compile(f, _tt, path::String = tempname();  name = GPUCompiler.safe_name(repr(f)), filename="obj",
+function compile(f, _tt, path::String = tempname(); 
+                 mixtape = NoContext(), 
+                 name = GPUCompiler.safe_name(repr(f)), 
+                 filename = "obj",
                  strip_llvm = false,
                  strip_asm  = true,
                  opt_level=3,
                  kwargs...)
     tt = Base.to_tuple_type(_tt)
     isconcretetype(tt) || error("input type signature $_tt is not concrete")
 
-    rt = last(only(native_code_typed(f, tt)))
+    rt = last(only(native_code_typed(f, tt, mixtape = mixtape)))
     isconcretetype(rt) || error("$f on $_tt did not infer to a concrete type. Got $rt")
-
     f_wrap!(out::Ref, args::Ref{<:Tuple}) = (out[] = f(args[]...); nothing)
-    _, _, table = generate_obj(f_wrap!, Tuple{RefValue{rt}, RefValue{tt}}, false, path, name; opt_level, strip_llvm, strip_asm, filename, kwargs...)
+    _, _, table = generate_obj(f_wrap!, Tuple{RefValue{rt}, RefValue{tt}}, false, path, name; mixtape = mixtape, opt_level, strip_llvm, strip_asm, filename, kwargs...)
 
     lf = LazyStaticCompiledFunction{rt, tt}(Symbol(f), path, name, filename, table)
     cjl_path = joinpath(path, "$filename.cjl")
@@ -136,18 +140,26 @@ shell> tree \$path
 ```
 """
 function generate_obj(f, tt, external = true, path::String = tempname(), name = GPUCompiler.safe_name(repr(f)), filenamebase::String="obj";
+                        mixtape = NoContext(),
                         strip_llvm = false,
                         strip_asm  = true,
                         opt_level=3,
                         kwargs...)
     mkpath(path)
     obj_path = joinpath(path, "$filenamebase.o")
     tm = GPUCompiler.llvm_machine(external ? ExternalNativeCompilerTarget() : NativeCompilerTarget())
-    job, kwargs = native_job(f, tt, external; name, kwargs...)
     #Get LLVM to generated a module of code for us. We don't want GPUCompiler's optimization passes.
+    job = CompilerJob(NativeCompilerTarget(), 
+                      FunctionSpec(f, tt, false, name), 
+                      StaticCompilerParams(; 
+                                            opt = true, 
+                                            mixtape = mixtape,
+                                            optlevel = Base.JLOptions().opt_level))
+
     mod, meta = GPUCompiler.JuliaContext() do context
         GPUCompiler.codegen(:llvm, job; strip=strip_llvm, only_entry=false, validate=false, optimize=false, ctx=context)
-    end
+    end  
+
     # Use Enzyme's annotation and optimization pipeline
     annotate!(mod)
     optimize!(mod, tm)
@@ -247,7 +259,7 @@ function compile_executable(f::Function, types=(), path::String="./", name=GPUCo
     isexecutableargtype = tt == Tuple{} || tt == Tuple{Int, Ptr{Ptr{UInt8}}}
     isexecutableargtype || @warn "input type signature $types should be either `()` or `(Int, Ptr{Ptr{UInt8}})` for standard executables"
 
-    rt = last(only(native_code_typed(f, tt)))
+    rt = last(only(native_code_typed(f, tt; kwargs...)))
     isconcretetype(rt) || error("`$f$types` did not infer to a concrete type. Got `$rt`")
     nativetype = isprimitivetype(rt) || isa(rt, Ptr)
     nativetype || @warn "Return type `$rt` of `$f$types` does not appear to be a native type. Consider returning only a single value of a native machine type (i.e., a single float, int/uint, bool, or pointer). \n\nIgnoring this warning may result in Undefined Behavior!"
@@ -302,7 +314,7 @@ function compile_shlib(f::Function, types=(), path::String="./", name=GPUCompile
     tt = Base.to_tuple_type(types)
     isconcretetype(tt) || error("input type signature `$types` is not concrete")
 
-    rt = last(only(native_code_typed(f, tt)))
+    rt = last(only(native_code_typed(f, tt; kwargs...)))
     isconcretetype(rt) || error("`$f$types` did not infer to a concrete type. Got `$rt`")
     nativetype = isprimitivetype(rt) || isa(rt, Ptr)
     nativetype || @warn "Return type `$rt` of `$f$types` does not appear to be a native type. Consider returning only a single value of a native machine type (i.e., a single float, int/uint, bool, or pointer). \n\nIgnoring this warning may result in Undefined Behavior!"

Project.toml

@@ -49,7 +49,6 @@ function instantiate(p::LazyStaticCompiledFunction{rt, tt}) where {rt, tt}
     prefix = LLVM.get_prefix(lljit)
     dg = LLVM.CreateDynamicLibrarySearchGeneratorForProcess(prefix)
     LLVM.add!(jd, dg)
-    
     LLVM.add!(lljit, jd, ofile)
     fptr = pointer(LLVM.lookup(lljit, "julia_" * p.name))