Scal

Project.toml

@@ -18,7 +18,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 [compat]
 CEnum = "0.4"
 EnzymeCore = "0.3"
-Enzyme_jll = "0.0.62"
+Enzyme_jll = "0.0.63"
 GPUCompiler = "0.19"
 LLVM = "5"
 ObjectFile = "0.3"

src/api.jl

@@ -216,7 +216,9 @@ EnzymeGradientUtilsLookup(gutils, val, B) = ccall((:EnzymeGradientUtilsLookup, l
 EnzymeGradientUtilsInvertPointer(gutils, val, B) = ccall((:EnzymeGradientUtilsInvertPointer, libEnzyme), LLVMValueRef, (EnzymeGradientUtilsRef, LLVMValueRef, LLVM.API.LLVMBuilderRef), gutils, val, B)
 EnzymeGradientUtilsDiffe(gutils, val, B) = ccall((:EnzymeGradientUtilsDiffe, libEnzyme), LLVMValueRef, (EnzymeGradientUtilsRef, LLVMValueRef, LLVM.API.LLVMBuilderRef), gutils, val, B)
 EnzymeGradientUtilsAddToDiffe(gutils, val, diffe, B, T) = ccall((:EnzymeGradientUtilsAddToDiffe, libEnzyme), Cvoid, (EnzymeGradientUtilsRef, LLVMValueRef, LLVMValueRef, LLVM.API.LLVMBuilderRef, LLVMTypeRef), gutils, val, diffe, B, T)
-EnzymeGradientUtilsAddToInvertedPointerDiffeTT(gutils, orig, vd, size, origptr, prediff, B, align, premask) = ccall((:EnzymeGradientUtilsAddToInvertedPointerDiffeTT, libEnzyme), Cvoid, (EnzymeGradientUtilsRef, LLVMValueRef, CTypeTreeRef, Cuint, LLVMValueRef, LLVMValueRef, LLVM.API.LLVMBuilderRef, Cuint, LLVMValueRef), gutils, orig, vd, size, origptr, prediff, B, align, premask)
+function EnzymeGradientUtilsAddToInvertedPointerDiffeTT(gutils, orig, origVal, vd, size, origptr, prediff, B, align, premask) 
+    ccall((:EnzymeGradientUtilsAddToInvertedPointerDiffeTT, libEnzyme), Cvoid, (EnzymeGradientUtilsRef, LLVMValueRef, LLVMValueRef, CTypeTreeRef, Cuint, LLVMValueRef, LLVMValueRef, LLVM.API.LLVMBuilderRef, Cuint, LLVMValueRef), gutils, orig, origVal, vd, size, origptr, prediff, B, align, premask)
+end
 
 EnzymeGradientUtilsSetDiffe(gutils, val, diffe, B) = ccall((:EnzymeGradientUtilsSetDiffe, libEnzyme), Cvoid, (EnzymeGradientUtilsRef, LLVMValueRef, LLVMValueRef, LLVM.API.LLVMBuilderRef), gutils, val, diffe, B)
 EnzymeGradientUtilsIsConstantValue(gutils, val) = ccall((:EnzymeGradientUtilsIsConstantValue, libEnzyme), UInt8, (EnzymeGradientUtilsRef, LLVMValueRef), gutils, val)

src/compiler.jl

@@ -4548,7 +4548,7 @@ function enzyme_custom_common_rev(forward::Bool, B::LLVM.API.LLVMBuilderRef, ori
                 size = sizeof(Ty)
                 align = 0
                 premask = C_NULL
-                API.EnzymeGradientUtilsAddToInvertedPointerDiffeTT(gutils, C_NULL, TT, size, v,           ext, B, align, premask)
+                API.EnzymeGradientUtilsAddToInvertedPointerDiffeTT(gutils, orig, C_NULL, TT, size, v,           ext, B, align, premask)
             else
                 @assert value_type(ext) == shadowVType
                 API.EnzymeGradientUtilsAddToDiffe(gutils, v, ext, B, Typ)

test/blas/scal.jl

@@ -0,0 +1,65 @@
+using Enzyme
+using LinearAlgebra
+using Test
+
+@testset "BLAS scal" begin
+
+    x1 = [1.0, 2.0, 3.0]
+    x2 = [2.0, 3.0, 4.0]
+    x3 = [2.0, 4.0, 6.0]
+    x4 = [2.0, 4.0, 6.0]
+    x5 = [2.0, 4.0, 6.0]
+
+    dx1 = [0.0, 0.0, 0.0]
+    dx2 = [1.0, 0.0, 0.0]
+    dx3 = [1.0, 1.0, 1.0]
+    dx4 = [1.0, 1.0, 2.0]
+    dx5 = [1.0, -1.0, 2.0]
+
+    ret1 = autodiff(Reverse, BLAS.scal!, Const, Active(1.0), Duplicated(x1, dx1))[1][1]
+    ret2 = autodiff(Reverse, BLAS.scal!, Const, Active(1.0), Duplicated(x2, dx2))[1][1]
+    ret3 = autodiff(Reverse, BLAS.scal!, Const, Active(1.0), Duplicated(x3, dx3))[1][1]
+    ret4 = autodiff(Reverse, BLAS.scal!, Const, Active(1.0), Duplicated(x4, dx4))[1][1]
+    ret5 = autodiff(Reverse, BLAS.scal!, Const, Active(1.0), Duplicated(x5, dx5))[1][1]
+    @test ret1 == 0.0
+    @test ret2 == 2.0
+    @test ret3 == 12.0
+    @test ret4 == 18.0
+    @test ret5 == 10.0
+    @test dx1 == [0.0, 0.0, 0.0]
+    @test dx2 == [1.0, 0.0, 0.0]
+    @test dx3 == [1.0, 1.0, 1.0]
+    @test dx4 == [1.0, 1.0, 2.0]
+    @test dx5 == [1.0, -1.0, 2.0]
+
+
+
+    x1 = [1.0, 2.0, 3.0]
+    x2 = [2.0, 3.0, 4.0]
+    x3 = [2.0, 4.0, 6.0]
+    x4 = [2.0, 4.0, 6.0]
+    x5 = [2.0, 4.0, 6.0]
+
+    dx1 = [0.0, 0.0, 0.0]
+    dx2 = [1.0, 0.0, 0.0]
+    dx3 = [1.0, 1.0, 1.0]
+    dx4 = [1.0, 1.0, 2.0]
+    dx5 = [1.0, -1.0, 2.0]
+
+    ret1 = autodiff(Reverse, BLAS.scal!, Const, Active(2.0), Duplicated(x1, dx1))[1][1]
+    ret2 = autodiff(Reverse, BLAS.scal!, Const, Active(2.0), Duplicated(x2, dx2))[1][1]
+    ret3 = autodiff(Reverse, BLAS.scal!, Const, Active(2.0), Duplicated(x3, dx3))[1][1]
+    ret4 = autodiff(Reverse, BLAS.scal!, Const, Active(2.0), Duplicated(x4, dx4))[1][1]
+    ret5 = autodiff(Reverse, BLAS.scal!, Const, Active(2.0), Duplicated(x5, dx5))[1][1]
+    @test ret1 == 0.0
+    @test ret2 == 2.0
+    @test ret3 == 12.0
+    @test ret4 == 18.0
+    @test ret5 == 10.0
+    @test dx1 == [0.0, 0.0, 0.0]
+    @test dx2 == [2.0, 0.0, 0.0]
+    @test dx3 == [2.0, 2.0, 2.0]
+    @test dx4 == [2.0, 2.0, 4.0]
+    @test dx5 == [2.0, -2.0, 4.0]
+
+end

test/runtests.jl

@@ -61,6 +61,8 @@ function test_matrix_to_number(f, x; rtol=1e-9, atol=1e-9, fdm=central_fdm(5, 1)
     @test isapprox(dx_fwd, dx_fd; rtol=rtol, atol=atol, kwargs...)
 end
 
+include("blas/scal.jl")
+
 include("abi.jl")
 include("typetree.jl")