splice in some more fixes

Author: shashi

src/SparsityDetection.jl

@@ -14,6 +14,7 @@ include("propagate_tags.jl")
 include("linearity.jl")
 include("jacobian.jl")
 include("hessian.jl")
+include("blas.jl")
 
 Base.@deprecate sparsity!(args...) jacobian_sparsity(args...)
 Base.@deprecate hsparsity(args...) hessian_sparsity(args...)

src/blas.jl

@@ -0,0 +1,15 @@
+# Forward BLAS calls to generic implementation
+#
+using LinearAlgebra
+import LinearAlgebra.BLAS
+
+# generic implementations
+
+@reroute LinearAlgebra.BLAS.dot(x,y) LinearAlgebra.dot(Any, Any)
+@reroute LinearAlgebra.BLAS.axpy!(x, y) LinearAlgebra.axpy!(Any,
+                                                      AbstractArray,
+                                                      AbstractArray)
+
+gengemv!(tA, α, A, x, β, y) = LinearAlgebra.generic_matvecmul!(y, tA, A, x, LinearAlgebra.MulAddMul(α, β))
+
+@reroute LinearAlgebra.BLAS.gemv!(tA, α, A, x, β, y) gengemv!(Any, Any, Any, Any, Any, Any)

src/controlflow.jl

@@ -34,10 +34,7 @@ end
 """
 function abstract_run(acc, ctx::Cassette.Context, overdub_fn, args...)
     pass_ctx = Cassette.similarcontext(ctx, pass=AbsintPass)
-    @ambrun begin
         acc(Cassette.overdub(pass_ctx, overdub_fn, args...))
-        @amb
-    end
 end
 
 """

src/hessian.jl

@@ -248,14 +248,3 @@ function hessian_sparsity(f, X, args...; raw=false)
 end
 
 
-# Forward BLAS calls to generic implementation
-#
-using LinearAlgebra
-import LinearAlgebra.BLAS
-
-# generic implementations
-
-@reroute HessianSparsityContext BLAS.dot dot(Any, Any)
-@reroute HessianSparsityContext BLAS.axpy! axpy!(Any,
-                                                 AbstractArray,
-                                                 AbstractArray)

src/jacobian.jl

@@ -128,24 +128,6 @@ function Cassette.overdub(ctx::JacobianSparsityContext,
     end
 end
 
-function Cassette.overdub(ctx::JacobianSparsityContext,
-                          f::typeof(Base.unsafe_copyto!),
-                          X::Tagged,
-                          xstart,
-                          Y::Tagged,
-                          ystart,
-                          len)
-    S = ctx.metadata
-    if metatype(Y, ctx) <: JacInput
-        val = Cassette.fallback(ctx, f, X, xstart, Y, ystart, len)
-        nometa = Cassette.NoMetaMeta()
-        X.meta.meta[xstart:xstart+len-1] .= (i->Cassette.Meta(ProvinanceSet(i), nometa)).(ystart:ystart+len-1)
-        val
-    else
-        Cassette.recurse(ctx, f, X, xstart, Y, ystart, len)
-    end
-end
-
 function jacobian_sparsity(f!, Y, X, args...;
                            sparsity=Sparsity(length(Y), length(X)),
                            raw = false)
@@ -170,29 +152,29 @@ function jacobian_sparsity(f!, Y, X, args...;
     end
 end
 
-function Cassette.overdub(ctx::SparsityContext,
+function Cassette.overdub(ctx::JacobianSparsityContext,
                           f::typeof(Base.unsafe_copyto!),
                           X::Tagged,
                           xstart,
                           Y::Tagged,
                           ystart,
                           len)
     S = ctx.metadata
-    if ismetatype(Y, ctx, JacInput) && ismetatype(X, ctx, JacOutput)
+    if metatype(Y, ctx) <: JacInput && metatype(X, ctx) <: JacOutput
         # Write directly to the output sparsity
         val = Cassette.fallback(ctx, f, X, xstart, Y, ystart, len)
         for (i, j) in zip(xstart:xstart+len-1, ystart:ystart+len-1)
             push!(S, i, j)
         end
         val
-    elseif ismetatype(Y, ctx, JacInput)
+    elseif metatype(Y, ctx) <: JacInput
         # Keep around a ProvinanceSet
         val = Cassette.fallback(ctx, f, X, xstart, Y, ystart, len)
         nometa = Cassette.NoMetaMeta()
-        rhs = (i->Cassette.Meta(pset(i), nometa)).(ystart:ystart+len-1)
+        rhs = (i->Cassette.Meta(ProvinanceSet(i), nometa)).(ystart:ystart+len-1)
         X.meta.meta[xstart:xstart+len-1] .= rhs
         val
-    elseif ismetatype(X, ctx, JacOutput)
+    elseif metatype(X, ctx) <: JacOutput
         val = Cassette.fallback(ctx, f, X, xstart, Y, ystart, len)
         for (i, j) in zip(xstart:xstart+len-1, ystart:ystart+len-1)
             y = Cassette.@overdub ctx Y[j]

src/util.jl

@@ -12,17 +12,36 @@ function metatype(x, ctx)
     end
 end
 
-macro reroute(ctx, f, g)
+# generic implementations
+
+_name(x::Symbol) = x
+_name(x::Expr) = (@assert x.head == :(::); x.args[1])
+macro reroute(f, g)
+    fname = f.args[1]
+    fargs = f.args[2:end]
+    gname = g.args[1]
+    gargs = g.args[2:end]
     quote
-        function Cassette.overdub(ctx::$ctx,
-                                  f::typeof($(esc(f))),
-                                  args...)
-            Cassette.overdub(
+        @inline function Cassette.overdub(ctx::JacobianSparsityContext,
+                                          f::typeof($(esc(fname))),
+                                          $(fargs...))
+            Cassette.recurse(
+                ctx,
+                invoke,
+                $(esc(gname)),
+                $(esc(:(Tuple{$(gargs...)}))),
+                $(map(_name, fargs)...))
+        end
+
+        @inline function Cassette.overdub(ctx::HessianSparsityContext,
+                                          f::typeof($(esc(fname))),
+                                          $(fargs...))
+            Cassette.recurse(
                 ctx,
                 invoke,
-                $(esc(g.args[1])),
-                $(esc(:(Tuple{$(g.args[2:end]...)}))),
-                args...)
+                $(esc(gname)),
+                $(esc(:(Tuple{$(gargs...)}))),
+                $(map(_name, fargs)...))
         end
     end
 end