loosen fwd runtime activity blas restrictions (#2398)

* loosen fwd runtime activity blas restrictions

* fix'

Author: wsmoses

Project.toml

@@ -39,7 +39,7 @@ BFloat16s = "0.2, 0.3, 0.4, 0.5"
 CEnum = "0.4, 0.5"
 ChainRulesCore = "1"
 EnzymeCore = "0.8.8"
-Enzyme_jll = "0.0.179"
+Enzyme_jll = "0.0.180"
 GPUArraysCore = "0.1.6, 0.2"
 GPUCompiler = "1.3"
 LLVM = "6.1, 7, 8, 9"

src/errors.jl

@@ -310,7 +310,15 @@ function julia_error(
                     println(io)
                 end
             end
-            emit_error(B, nothing, msg2, EnzymeNoDerivativeError)
+    	    if data2 != C_NULL
+        		data2 = LLVM.Value(data2)
+                if value_type(data2) != LLVM.IntType(1)
+                    data2 = nothing
+                end
+            else
+                data2 = nothing
+            end
+            emit_error(B, nothing, msg2, EnzymeNoDerivativeError, data2)
             return C_NULL
         end
         throw(NoDerivativeException(msg, ir, bt))

src/jlrt.jl

@@ -199,7 +199,48 @@ function emit_jl_throw!(B::LLVM.IRBuilder, @nospecialize(val::LLVM.Value))::LLVM
     T_prjlvalue = LLVM.PointerType(T_jlvalue, 12)
     FT = LLVM.FunctionType(T_void, [T_prjlvalue])
     fn, _ = get_function!(mod, "jl_throw", FT)
-    call!(B, FT, fn, LLVM.Value[val])
+    cb = call!(B, FT, fn, LLVM.Value[val])
+    LLVM.API.LLVMAddCallSiteAttribute(
+		cb,
+		reinterpret(LLVM.API.LLVMAttributeIndex, LLVM.API.LLVMAttributeFunctionIndex),
+		EnumAttribute("noreturn"),
+	    )
+    return cb
+end
+
+function emit_conditional_throw!(B::LLVM.IRBuilder, @nospecialize(val::LLVM.Value), @nospecialize(cond::LLVM.Value))::LLVM.Value
+    curent_bb = position(B)
+    fn = LLVM.parent(curent_bb)
+    mod = LLVM.parent(fn)
+    T_void = LLVM.VoidType()
+    T_jlvalue = LLVM.StructType(LLVMType[])
+    T_prjlvalue = LLVM.PointerType(T_jlvalue, 10)
+    FT = LLVM.FunctionType(T_void, [T_prjlvalue, LLVM.IntType(1)])
+
+    name = "jl_conditional_throw"
+    if haskey(functions(mod), name)
+        fn = functions(mod)[name]
+    else
+        fn = LLVM.Function(mod, "jl_conditional_throw", FT)
+    	linkage!(fn, LLVM.API.LLVMInternalLinkage)
+        err, rcond = LLVM.parameters(fn)
+	 builder = LLVM.IRBuilder()
+         entry = BasicBlock(fn, "entry")
+         errb = BasicBlock(fn, "err")
+         exitb = BasicBlock(fn, "errb")
+         position!(builder, entry)
+	 br!(builder, rcond, errb, exitb)
+         position!(builder, errb)
+	 err = addrspacecast!(builder, err, LLVM.PointerType(T_jlvalue, 12))
+	 thrown = emit_jl_throw!(builder, err)
+	 unreachable!(builder)
+	 position!(builder, exitb)
+	 ret!(builder)
+
+        push!(LLVM.function_attributes(fn), LLVM.EnumAttribute("alwaysinline", 0))
+    end
+
+    call!(B, FT, fn, LLVM.Value[val, cond])
 end
 
 function emit_box_int32!(B::LLVM.IRBuilder, @nospecialize(val::LLVM.Value))::LLVM.Value
@@ -1004,7 +1045,7 @@ function allocate_sret!(gutils::API.EnzymeGradientUtilsRef, @nospecialize(N::LLV
     allocate_sret!(B, N)
 end
 
-function emit_error(B::LLVM.IRBuilder, @nospecialize(orig::Union{Nothing, LLVM.Instruction}), string::Union{String, LLVM.Value}, @nospecialize(errty::Type) = EnzymeRuntimeException)
+function emit_error(B::LLVM.IRBuilder, @nospecialize(orig::Union{Nothing, LLVM.Instruction}), string::Union{String, LLVM.Value}, @nospecialize(errty::Type) = EnzymeRuntimeException, @nospecialize(cond::Union{Nothing, LLVM.Value}) = nothing)
     curent_bb = position(B)
     fn = LLVM.parent(curent_bb)
     mod = LLVM.parent(fn)
@@ -1061,24 +1102,30 @@ function emit_error(B::LLVM.IRBuilder, @nospecialize(orig::Union{Nothing, LLVM.I
         err = emit_allocobj!(B, errty)
         err2 = bitcast!(B, err, LLVM.PointerType(LLVM.PointerType(LLVM.Int8Type()), 10))
         store!(B, string, err2)
-        emit_jl_throw!(
-            B,
-            addrspacecast!(B, err, LLVM.PointerType(LLVM.StructType(LLVMType[]), 12)),
-        )
+	if cond !== nothing
+        	emit_conditional_throw!(B, err, cond)
+	else
+		emit_jl_throw!(
+		    B,
+		    addrspacecast!(B, err, LLVM.PointerType(LLVM.StructType(LLVMType[]), 12)),
+		)
+	end
     end
 
     # 2. Call error function and insert unreachable
-    LLVM.API.LLVMAddCallSiteAttribute(
-        ct,
-        reinterpret(LLVM.API.LLVMAttributeIndex, LLVM.API.LLVMAttributeFunctionIndex),
-        EnumAttribute("noreturn"),
-    )
-    if EnzymeMutabilityException != errty
-        LLVM.API.LLVMAddCallSiteAttribute(
-            ct,
-            reinterpret(LLVM.API.LLVMAttributeIndex, LLVM.API.LLVMAttributeFunctionIndex),
-            StringAttribute("enzyme_error"),
-        )
+    if cond === nothing
+	    LLVM.API.LLVMAddCallSiteAttribute(
+		ct,
+		reinterpret(LLVM.API.LLVMAttributeIndex, LLVM.API.LLVMAttributeFunctionIndex),
+		EnumAttribute("noreturn"),
+	    )
+	    if EnzymeMutabilityException != errty
+		LLVM.API.LLVMAddCallSiteAttribute(
+		    ct,
+		    reinterpret(LLVM.API.LLVMAttributeIndex, LLVM.API.LLVMAttributeFunctionIndex),
+		    StringAttribute("enzyme_error"),
+		)
+	    end
     end
     return ct
 end

test/internal_rules.jl

@@ -282,6 +282,43 @@ end
     end
 end
 
+function two_blas(a, b)
+  a = copy(a)
+  @inline LinearAlgebra.LAPACK.potrf!('L', a)
+  @inline LinearAlgebra.LAPACK.potrf!('L', b)
+  return a[1,1] + b[1,1]
+end
+
+@testset "Forward Mode runtime activity" begin
+
+  a = [2.7 3.5; 7.4 9.2]
+  da = [7.2 5.3; 4.7 2.9]
+
+  b = [3.1 5.6; 13 19]
+  db = [1.3 6.5; .13 .19]
+  
+  res = Enzyme.autodiff(Forward, two_blas, Duplicated(a, da), Duplicated(b, db))[1]
+  @test res ≈ 2.5600654222812564
+
+  a = [2.7 3.5; 7.4 9.2]
+  da = [7.2 5.3; 4.7 2.9]
+
+  b = [3.1 5.6; 13 19]
+  db = [1.3 6.5; .13 .19]
+
+  res = Enzyme.autodiff(set_runtime_activity(Forward), two_blas, Duplicated(a, da), Duplicated(b, db))[1]
+  @test res ≈ 2.5600654222812564
+
+  a = [2.7 3.5; 7.4 9.2]
+  da = [7.2 5.3; 4.7 2.9]
+
+  b = [3.1 5.6; 13 19]
+  db = [1.3 6.5; .13 .19]
+
+  @test_throws Enzyme.Compiler.EnzymeNoDerivativeError Enzyme.autodiff(set_runtime_activity(Forward), two_blas, Duplicated(a, da), Duplicated(b, b))
+
+end
+
 @testset "Cholesky" begin
     function symmetric_definite(n :: Int=10)
         α = one(Float64)